// // Copyright (c) Microsoft. All rights reserved. // Licensed under the MIT license. See LICENSE file in the project root for full license information. // using System; using System.Data.Common; using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.SqlServer.Management.SqlParser.Parser; using Microsoft.SqlTools.ServiceLayer.Connection; using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; using Microsoft.SqlTools.ServiceLayer.SqlContext; namespace Microsoft.SqlTools.ServiceLayer.QueryExecution { /// /// Internal representation of an active query /// public class Query : IDisposable { #region Properties /// /// The batches underneath this query /// internal Batch[] Batches { get; set; } /// /// The summaries of the batches underneath this query /// public BatchSummary[] BatchSummaries { get { if (!HasExecuted) { throw new InvalidOperationException("Query has not been executed."); } return Batches.Select((batch, index) => new BatchSummary { Id = index, HasError = batch.HasError, Messages = batch.ResultMessages.ToArray(), ResultSetSummaries = batch.ResultSummaries }).ToArray(); } } /// /// Cancellation token source, used for cancelling async db actions /// private readonly CancellationTokenSource cancellationSource; /// /// The connection info associated with the file editor owner URI, used to create a new /// connection upon execution of the query /// private ConnectionInfo EditorConnection { get; set; } /// /// Whether or not the query has completed executed, regardless of success or failure /// /// /// Don't touch the setter unless you're doing unit tests! /// public bool HasExecuted { get { return Batches.All(b => b.HasExecuted); } internal set { foreach (var batch in Batches) { batch.HasExecuted = value; } } } /// /// The text of the query to execute /// public string QueryText { get; set; } #endregion /// /// Constructor for a query /// /// The text of the query to execute /// The information of the connection to use to execute the query /// Settings for how to execute the query, from the user public Query(string queryText, ConnectionInfo connection, QueryExecutionSettings settings) { // Sanity check for input if (String.IsNullOrEmpty(queryText)) { throw new ArgumentNullException(nameof(queryText), "Query text cannot be null"); } if (connection == null) { throw new ArgumentNullException(nameof(connection), "Connection cannot be null"); } if (settings == null) { throw new ArgumentNullException(nameof(settings), "Settings cannot be null"); } // Initialize the internal state QueryText = queryText; EditorConnection = connection; cancellationSource = new CancellationTokenSource(); // Process the query into batches ParseResult parseResult = Parser.Parse(queryText, new ParseOptions { BatchSeparator = settings.BatchSeparator }); Batches = parseResult.Script.Batches.Select(b => new Batch(b.Sql)).ToArray(); } /// /// Executes this query asynchronously and collects all result sets /// public async Task Execute() { // Open up a connection for querying the database string connectionString = ConnectionService.BuildConnectionString(EditorConnection.ConnectionDetails); using (DbConnection conn = EditorConnection.Factory.CreateSqlConnection(connectionString)) { await conn.OpenAsync(); // We need these to execute synchronously, otherwise the user will be very unhappy foreach (Batch b in Batches) { await b.Execute(conn, cancellationSource.Token); } } } /// /// Retrieves a subset of the result sets /// /// The index for selecting the batch item /// The index for selecting the result set /// The starting row of the results /// How many rows to retrieve /// A subset of results public ResultSetSubset GetSubset(int batchIndex, int resultSetIndex, int startRow, int rowCount) { // Sanity check that the results are available if (!HasExecuted) { throw new InvalidOperationException("The query has not completed, yet."); } // Sanity check to make sure that the batch is within bounds if (batchIndex < 0 || batchIndex >= Batches.Length) { throw new ArgumentOutOfRangeException(nameof(batchIndex), "Result set index cannot be less than 0" + "or greater than the number of result sets"); } return Batches[batchIndex].GetSubset(resultSetIndex, startRow, rowCount); } /// /// Cancels the query by issuing the cancellation token /// public void Cancel() { // Make sure that the query hasn't completed execution if (HasExecuted) { throw new InvalidOperationException("The query has already completed, it cannot be cancelled."); } // Issue the cancellation token for the query cancellationSource.Cancel(); } #region IDisposable Implementation private bool disposed; public void Dispose() { Dispose(true); GC.SuppressFinalize(this); } protected virtual void Dispose(bool disposing) { if (disposed) { return; } if (disposing) { cancellationSource.Dispose(); } disposed = true; } ~Query() { Dispose(false); } #endregion } }