Adding cancellation preemption to query execution (#409)

* Adding cancellation preemption to query execution

* Removing redundant cancellation check as per PR comment

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs

@@ -241,6 +241,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             try
             {
+                // Make sure we haven't cancelled yet
+                cancellationToken.ThrowIfCancellationRequested();
+
                 // Register the message listener to *this instance* of the batch
                 // Note: This is being done to associate messages with batches
                 ReliableSqlConnection sqlConn = conn as ReliableSqlConnection;
@@ -279,6 +282,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         int resultSetOrdinal = 0;
                         do
                         {
+                            // Verify that the cancellation token hasn't benn cancelled
+                            cancellationToken.ThrowIfCancellationRequested();
+
                             // Skip this result set if there aren't any rows (ie, UPDATE/DELETE/etc queries)
                             if (!reader.HasRows && reader.FieldCount == 0)
                             {
@@ -297,7 +303,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                             }
 
                             // Read until we hit the end of the result set
-                            await resultSet.ReadResultToEnd(reader, cancellationToken).ConfigureAwait(false);
+                            await resultSet.ReadResultToEnd(reader, cancellationToken);
 
                         } while (await reader.NextResultAsync(cancellationToken));
 

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs

@@ -342,6 +342,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
             try
             {
+                // Verify the request hasn't been cancelled
+                cancellationToken.ThrowIfCancellationRequested();
+
                 // Mark that result has been read
                 hasBeenRead = true;