/// <summary> /// Executes this query asynchronously and collects all result sets /// </summary> private async Task ExecuteInternal() { ReliableSqlConnection sqlConn = null; try { // Mark that we've internally executed hasExecuteBeenCalled = true; // Don't actually execute if there aren't any batches to execute if (Batches.Length == 0) { if (BatchMessageSent != null) { await BatchMessageSent(new ResultMessage(SR.QueryServiceCompletedSuccessfully, false, null)); } if (QueryCompleted != null) { await QueryCompleted(this); } return; } // Locate and setup the connection DbConnection queryConnection = await ConnectionService.Instance.GetOrOpenConnection(editorConnection.OwnerUri, ConnectionType.Query); sqlConn = queryConnection as ReliableSqlConnection; if (sqlConn != null) { // Subscribe to database informational messages sqlConn.GetUnderlyingConnection().FireInfoMessageEventOnUserErrors = true; sqlConn.GetUnderlyingConnection().InfoMessage += OnInfoMessage; } // Execute beforeBatches synchronously, before the user defined batches foreach (Batch b in BeforeBatches) { await b.Execute(queryConnection, cancellationSource.Token); } // We need these to execute synchronously, otherwise the user will be very unhappy foreach (Batch b in Batches) { // Add completion callbacks b.BatchStart += BatchStarted; b.BatchCompletion += BatchCompleted; b.BatchMessageSent += BatchMessageSent; b.ResultSetCompletion += ResultSetCompleted; await b.Execute(queryConnection, cancellationSource.Token); } // Execute afterBatches synchronously, after the user defined batches foreach (Batch b in AfterBatches) { await b.Execute(queryConnection, cancellationSource.Token); } // Call the query execution callback if (QueryCompleted != null) { await QueryCompleted(this); } } catch (Exception e) { // Call the query failure callback if (QueryFailed != null) { await QueryFailed(this, e); } } finally { // Remove the message handler from the connection if (sqlConn != null) { // Subscribe to database informational messages sqlConn.GetUnderlyingConnection().InfoMessage -= OnInfoMessage; } // If any message notified us we had changed databases, then we must let the connection service know if (newDatabaseName != null) { ConnectionService.Instance.ChangeConnectionDatabaseContext(editorConnection.OwnerUri, newDatabaseName); } foreach (Batch b in Batches) { if (b.HasError) { ConnectionService.EnsureConnectionIsOpen(sqlConn); break; } } } }
private async Task ExecuteOnce(DbConnection conn, CancellationToken cancellationToken) { // Make sure we haven't cancelled yet cancellationToken.ThrowIfCancellationRequested(); // Create a command that we'll use for executing the query using (DbCommand dbCommand = CreateCommand(conn)) { // Make sure that we cancel the command if the cancellation token is cancelled cancellationToken.Register(() => dbCommand?.Cancel()); // Setup the command for executing the batch dbCommand.CommandText = BatchText; dbCommand.CommandType = CommandType.Text; dbCommand.CommandTimeout = 0; List <DbColumn[]> columnSchemas = null; if (getFullColumnSchema) { // Executing the same query twice with different command behavior causes the second // execution to return no rows if there's a trailing comment with no newline after, // so add a newline to the end of the query. See https://github.com/Microsoft/sqlopsstudio/issues/1424 dbCommand.CommandText += Environment.NewLine; ConnectionService.EnsureConnectionIsOpen(conn); // Fetch schema info separately, since CommandBehavior.KeyInfo will include primary // key columns in the result set, even if they weren't part of the select statement. // Extra key columns get added to the end, so just correlate via Column Ordinal. columnSchemas = new List <DbColumn[]>(); using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(CommandBehavior.KeyInfo | CommandBehavior.SchemaOnly, cancellationToken)) { if (reader != null && reader.CanGetColumnSchema()) { do { columnSchemas.Add(reader.GetColumnSchema().ToArray()); } while (await reader.NextResultAsync(cancellationToken)); } } } ConnectionService.EnsureConnectionIsOpen(conn); // Execute the command to get back a reader using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(cancellationToken)) { do { // Verify that the cancellation token hasn't been canceled cancellationToken.ThrowIfCancellationRequested(); // Skip this result set if there aren't any rows (i.e. UPDATE/DELETE/etc queries) if (!reader.HasRows && reader.FieldCount == 0) { continue; } // This resultset has results (i.e. SELECT/etc queries) ResultSet resultSet = new ResultSet(resultSets.Count, Id, outputFileFactory); resultSet.ResultCompletion += ResultSetCompletion; // Add the result set to the results of the query lock (resultSets) { resultSets.Add(resultSet); } // Read until we hit the end of the result set await resultSet.ReadResultToEnd(reader, cancellationToken); } while (await reader.NextResultAsync(cancellationToken)); // If there were no messages, for whatever reason (NO COUNT set, messages // were emitted, records returned), output a "successful" message if (!messagesSent) { await SendMessage(SR.QueryServiceCompletedSuccessfully, false); } } if (columnSchemas != null) { ExtendResultMetadata(columnSchemas, resultSets); } } }