From 164ca951dae07389be2cd4fbb2729910b20ed6f4 Mon Sep 17 00:00:00 2001
From: Monica Gupta <scorpio90m@gmail.com>
Date: Sun, 4 Oct 2020 19:12:33 -0700
Subject: [PATCH] Fix for multiple queries separated by space for Kusto (#1089)

* Fix for multiple queries separated by space

* Addressed comments

Co-authored-by: Monica Gupta <mogupt@microsoft.com>
---
 .../DataSource/KustoClient.cs                 | 33 +++++++++++--------
 .../DataSource/KustoResultsReader.cs          | 27 ++++++++-------
 2 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs
index fc9539a8..9fa4b246 100644
--- a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs
@@ -173,22 +173,29 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             clientRequestProperties.SetOption(ClientRequestProperties.OptionNoTruncation, true);
             cancellationToken.Register(() => CancelQuery(clientRequestProperties.ClientRequestId));
 
-            var kustoCodeService = new KustoCodeService(query);
-            var minimalQuery = kustoCodeService.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
+            var script = CodeScript.From(query, GlobalState.Default);
+            IDataReader[] origReaders = new IDataReader[script.Blocks.Count];
 
-            try
+            Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
             {
-                IDataReader origReader = _kustoQueryProvider.ExecuteQuery(
-                    KustoQueryUtils.IsClusterLevelQuery(minimalQuery) ? "" : databaseName,
-                    minimalQuery,
-                    clientRequestProperties);
+                var minimalQuery = codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
                 
-                return new KustoResultsReader(origReader);
-            }
-            catch (KustoRequestException exception) when (exception.FailureCode == 401) // Unauthorized
-            {
-                throw new DataSourceUnauthorizedException(exception);
-            }
+                try
+                {
+                    IDataReader origReader = _kustoQueryProvider.ExecuteQuery(
+                        KustoQueryUtils.IsClusterLevelQuery(minimalQuery) ? "" : databaseName,
+                        minimalQuery,
+                        clientRequestProperties);
+                    
+                    origReaders[index] = origReader;
+                }
+                catch (KustoRequestException exception) when (exception.FailureCode == 401) // Unauthorized
+                {
+                    throw new DataSourceUnauthorizedException(exception);
+                }
+            });
+
+            return new KustoResultsReader(origReaders);
         }
 
         /// <summary>
diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoResultsReader.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoResultsReader.cs
index baf60fae..09fff1b1 100644
--- a/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoResultsReader.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/DataSource/KustoResultsReader.cs
@@ -11,26 +11,31 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// Kusto returns atleast 4 results tables - QueryResults(sometimes more than one), QueryProperties, QueryStatus and Query Results Metadata Table. 
         /// ADS just needs query results. When returning query results we need to trim off the last 3 tables. 
         /// </summary> 
-        public KustoResultsReader(IDataReader reader)
+        public KustoResultsReader(IDataReader[] readers)
             : base()
         {
             // Read out all tables
             List<DataTable> results = new List<DataTable>();
-            while (!(reader?.IsClosed ?? true))
-            {
-                DataTable dt = new DataTable();
-                dt.Load(reader); // This calls NextResult on the reader
-                results.Add(dt);
-            }
 
-            // Trim results
-            if(results.Count > 3) results.RemoveRange(results.Count - 3, 3);
+            foreach (var reader in readers)
+            {
+                while (!(reader?.IsClosed ?? true))
+                {
+                    DataTable dt = new DataTable();
+                    dt.Load(reader); // This calls NextResult on the reader
+                    results.Add(dt);
+                }
+                
+                // Trim results
+                if(results.Count > 3) results.RemoveRange(results.Count - 3, 3);
+            }
 
             // Create a DataReader for the trimmed set
             _resultDataSet = new DataSet();
-            foreach(var result in results)
+            for(int i = 0; i < results.Count; i++)
             {
-                _resultDataSet.Tables.Add(result);
+                results[i].TableName = "Table_" + i;
+                _resultDataSet.Tables.Add(results[i]);
             }
 
             SetDataReader(_resultDataSet.CreateDataReader());