//
// 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
}
}