using System; using System.Collections.Generic; using System.Data; using System.Data.Common; using System.Linq; using System.Threading; using System.Threading.Tasks; using Microsoft.SqlTools.ServiceLayer.Connection; using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts; namespace Microsoft.SqlTools.ServiceLayer.QueryExecution { public class Query //: IDisposable { #region Properties public string QueryText { get; set; } public ConnectionInfo EditorConnection { get; set; } private readonly CancellationTokenSource cancellationSource; public List ResultSets { get; set; } public ResultSetSummary[] ResultSummary { get { return ResultSets.Select((set, index) => new ResultSetSummary { ColumnInfo = set.Columns, Id = index, RowCount = set.Rows.Count }).ToArray(); } } public bool HasExecuted { get; set; } #endregion public Query(string queryText, ConnectionInfo connection) { // Sanity check for input if (queryText == null) { throw new ArgumentNullException(nameof(queryText), "Query text cannot be null"); } if (connection == null) { throw new ArgumentNullException(nameof(connection), "Connection cannot be null"); } // Initialize the internal state QueryText = queryText; EditorConnection = connection; HasExecuted = false; ResultSets = new List(); cancellationSource = new CancellationTokenSource(); } public async Task Execute() { // Sanity check to make sure we haven't already run this query if (HasExecuted) { throw new InvalidOperationException("Query has already executed."); } // Create a connection from the connection details string connectionString = ConnectionService.BuildConnectionString(EditorConnection.ConnectionDetails); using (DbConnection conn = EditorConnection.Factory.CreateSqlConnection(connectionString)) { await conn.OpenAsync(cancellationSource.Token); // Create a command that we'll use for executing the query using (DbCommand command = conn.CreateCommand()) { command.CommandText = QueryText; command.CommandType = CommandType.Text; // Execute the command to get back a reader using (DbDataReader reader = await command.ExecuteReaderAsync(cancellationSource.Token)) { do { // TODO: This doesn't properly handle scenarios where the query is SELECT but does not have rows if (!reader.HasRows) { continue; } // Read until we hit the end of the result set ResultSet resultSet = new ResultSet(); while (await reader.ReadAsync(cancellationSource.Token)) { resultSet.AddRow(reader); } // Read off the column schema information if (reader.CanGetColumnSchema()) { resultSet.Columns = reader.GetColumnSchema().ToArray(); } // Add the result set to the results of the query ResultSets.Add(resultSet); } while (await reader.NextResultAsync(cancellationSource.Token)); } } } // Mark that we have executed HasExecuted = true; } public ResultSetSubset GetSubset(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 we have valid numbers if (resultSetIndex < 0 || resultSetIndex >= ResultSets.Count) { throw new ArgumentOutOfRangeException(nameof(resultSetIndex), "Result set index cannot be less than 0" + "or greater than the number of result sets"); } ResultSet targetResultSet = ResultSets[resultSetIndex]; if (startRow < 0 || startRow >= targetResultSet.Rows.Count) { throw new ArgumentOutOfRangeException(nameof(startRow), "Start row cannot be less than 0 " + "or greater than the number of rows in the resultset"); } if (rowCount <= 0) { throw new ArgumentOutOfRangeException(nameof(rowCount), "Row count must be a positive integer"); } // Retrieve the subset of the results as per the request object[][] rows = targetResultSet.Rows.Skip(startRow).Take(rowCount).ToArray(); return new ResultSetSubset { Rows = rows, RowCount = rows.Length }; } } }