Query Execution: Adding better cancellation support (#498)

* Adding better cancellation support

* Changing implementation as per @pensivebrian's comment
This commit is contained in:
Benjamin Russell
2017-10-16 11:26:15 -07:00
committed by GitHub
parent 0c7f559315
commit 5d11654a10

View File

@@ -24,7 +24,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
public class Batch : IDisposable public class Batch : IDisposable
{ {
#region Member Variables #region Member Variables
/// <summary> /// <summary>
/// For IDisposable implementation, whether or not this has been disposed /// For IDisposable implementation, whether or not this has been disposed
/// </summary> /// </summary>
@@ -247,37 +246,41 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
// Register the message listener to *this instance* of the batch // Register the message listener to *this instance* of the batch
// Note: This is being done to associate messages with batches // Note: This is being done to associate messages with batches
ReliableSqlConnection sqlConn = conn as ReliableSqlConnection; ReliableSqlConnection sqlConn = conn as ReliableSqlConnection;
DbCommand command; DbCommand dbCommand;
if (sqlConn != null) if (sqlConn != null)
{ {
// Register the message listener to *this instance* of the batch // Register the message listener to *this instance* of the batch
// Note: This is being done to associate messages with batches // Note: This is being done to associate messages with batches
sqlConn.GetUnderlyingConnection().InfoMessage += ServerMessageHandler; sqlConn.GetUnderlyingConnection().InfoMessage += ServerMessageHandler;
command = sqlConn.GetUnderlyingConnection().CreateCommand(); dbCommand = sqlConn.GetUnderlyingConnection().CreateCommand();
// Add a handler for when the command completes // Add a handler for when the command completes
SqlCommand sqlCommand = (SqlCommand)command; SqlCommand sqlCommand = (SqlCommand)dbCommand;
sqlCommand.StatementCompleted += StatementCompletedHandler; sqlCommand.StatementCompleted += StatementCompletedHandler;
} }
else else
{ {
command = conn.CreateCommand(); dbCommand = conn.CreateCommand();
} }
// Make sure we aren't using a ReliableCommad since we do not want automatic retry // Make sure we aren't using a ReliableCommad since we do not want automatic retry
Debug.Assert(!(command is ReliableSqlConnection.ReliableSqlCommand), Debug.Assert(!(dbCommand is ReliableSqlConnection.ReliableSqlCommand),
"ReliableSqlCommand command should not be used to execute queries"); "ReliableSqlCommand command should not be used to execute queries");
// Create a command that we'll use for executing the query // Create a command that we'll use for executing the query
using (command) using (dbCommand)
{ {
command.CommandText = BatchText; // Make sure that we cancel the command if the cancellation token is cancelled
command.CommandType = CommandType.Text; cancellationToken.Register(() => dbCommand?.Cancel());
command.CommandTimeout = 0;
// Setup the command for executing the batch
dbCommand.CommandText = BatchText;
dbCommand.CommandType = CommandType.Text;
dbCommand.CommandTimeout = 0;
executionStartTime = DateTime.Now; executionStartTime = DateTime.Now;
// Execute the command to get back a reader // Execute the command to get back a reader
using (DbDataReader reader = await command.ExecuteReaderAsync(cancellationToken)) using (DbDataReader reader = await dbCommand.ExecuteReaderAsync(cancellationToken))
{ {
int resultSetOrdinal = 0; int resultSetOrdinal = 0;
do do