mirror of
https://github.com/ckaczor/sqltoolsservice.git
synced 2026-02-16 10:58:30 -05:00
Query Execution: Adding better cancellation support (#498)
* Adding better cancellation support * Changing implementation as per @pensivebrian's comment
This commit is contained in:
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user