From 6099746922f8637d0c05dc0d153a3fe3e3784f10 Mon Sep 17 00:00:00 2001
From: Cheena Malhotra <13396919+cheenamalhotra@users.noreply.github.com>
Date: Wed, 2 Aug 2023 14:04:30 -0700
Subject: [PATCH] Improve Query editor Read performance + cancel timely for
 large data (#2161)

---
 .../QueryExecution/Batch.cs                   |  1 +
 .../DataStorage/StorageDataReader.cs          | 22 +++++++++++++++++++
 .../QueryExecution/ResultSet.cs               |  3 ++-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
index 28126227..c96b6ad8 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs
@@ -429,6 +429,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         {
                             do
                             {
+                                cancellationToken.ThrowIfCancellationRequested();
                                 columnSchemas.Add(reader.GetColumnSchema().ToArray());
                             } while (reader.NextResult());
                         }
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/StorageDataReader.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/StorageDataReader.cs
index 49209154..dc8de79e 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/StorageDataReader.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/DataStorage/StorageDataReader.cs
@@ -88,11 +88,33 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.DataStorage
         /// </summary>
         /// <param name="cancellationToken">The cancellation token to use for cancelling a query</param>
         /// <returns></returns>
+        [Obsolete("Deprecated due to performance issues, please use Read() instead.")]
         public Task<bool> ReadAsync(CancellationToken cancellationToken)
         {
             return DbDataReader.ReadAsync(cancellationToken);
         }
 
+        /// <summary>
+        /// Pass-through to DbDataReader.Read()
+        /// 
+        /// ************** IMPORTANT ****************
+        /// M.D.SqlClient's ReadAsync() implementation is not as 
+        /// performant as Read() and doesn't respect Cancellation Token 
+        /// due to long existing design issues like below:
+        /// 
+        /// https://github.com/dotnet/SqlClient/issues/593
+        /// https://github.com/dotnet/SqlClient/issues/44
+        /// 
+        /// Until these issues are resolved, prefer using Sync APIs.
+        /// *****************************************
+        /// 
+        /// </summary>
+        /// <returns></returns>
+        public bool Read()
+        {
+            return DbDataReader.Read();
+        }
+
         /// <summary>
         /// Retrieves a value
         /// </summary>
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs
index 1487114d..19406c85 100644
--- a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs
+++ b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultSet.cs
@@ -396,8 +396,9 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     //
                     availableTask = SendCurrentResults();
 
-                    while (await dataReader.ReadAsync(cancellationToken))
+                    while (dataReader.Read())
                     {
+                        cancellationToken.ThrowIfCancellationRequested();
                         fileOffsets.Add(totalBytesWritten);
                         totalBytesWritten += fileWriter.WriteRow(dataReader);
                     }