Fix for displaying multiple table (#1075)

src/Microsoft.Kusto.ServiceLayer/DataSource/KustoResultsReader.cs

@@ -1,21 +1,27 @@
-using System.Data;
+using Microsoft.Kusto.ServiceLayer.QueryExecution;
+using System.Collections.Generic;
+using System.Data;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
     internal class KustoResultsReader : DataReaderWrapper
     {
-        public KustoResultsReader(IDataReader reader) : base(reader)
+        public KustoResultsReader(IDataReader reader)
+            : base(reader)
         {
         }
 
         /// <summary>
-        /// Kusto returns 3 results tables - QueryResults, QueryProperties, QueryStatus. When returning query results	
+        /// Kusto returns atleast 4 results tables - QueryResults(sometimes more than one), QueryProperties, QueryStatus and Query Results Metadata Table. 
-        /// we want the caller to only read the first table. We override the NextResult function here to only return one table	
+        /// When returning query results we need to trim off the last 3 tables as we want the caller to only read results table. 
-        /// from the IDataReader.	
         /// </summary> 
-        public override bool NextResult()	
+
+        public void SanitizeResults(List<ResultSet> resultSets)
         {
-            return false;	
+            if (resultSets.Count > 3)
+            {
+                resultSets.RemoveRange(resultSets.Count - 3, 3);
+            }
         }
     }
 }

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

@@ -6,7 +6,6 @@ using System;
 using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
-using System.Diagnostics;
 using System.Data.SqlClient;
 using System.Linq;
 using System.Threading;
@@ -16,7 +15,7 @@ using Microsoft.Kusto.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.Kusto.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.Utility;
 using System.Globalization;
-using System.Collections.ObjectModel;
+using Microsoft.Kusto.ServiceLayer.DataSource;
 
 namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 {
@@ -369,6 +368,14 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 
                 } while (reader.NextResult());
 
+                KustoResultsReader kreader = reader as KustoResultsReader;
+                kreader.SanitizeResults(resultSets);
+
+                foreach (var resultSet in resultSets)
+                {
+                    await resultSet.SendCurrentResults();
+                }
+                
                 // If there were no messages, for whatever reason (NO COUNT set, messages 
                 // were emitted, records returned), output a "successful" message
                 if (!messagesSent)

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

@@ -338,7 +338,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             //
             Validate.IsNotNull(nameof(dbDataReader), dbDataReader);
 
-            Task availableTask = null;
             try
             {
                 // Verify the request hasn't been cancelled
@@ -357,11 +356,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                     //
                     hasStartedRead = true;
 
-                    // Invoke the SendCurrentResults() asynchronously that will send the results available notification
-                    //   and also trigger the timer to send periodic updates.
-                    //
-                    availableTask = SendCurrentResults();
-
                     while (await dataReader.ReadAsync(cancellationToken))
                     {
                         fileOffsets.Add(totalBytesWritten);
@@ -371,25 +365,9 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             }
             finally
             {
-
-                // await the completion of available notification in case it is not already done before proceeding
-                //
-                await availableTask;
-
                 // now set the flag to indicate that we are done reading. this equates to Complete flag to be marked 'True' in any future notifications.
                 //
                 hasCompletedRead = true;
-
-
-                // Make a final call to SendCurrentResults() and await its completion. If the previously scheduled task already took care of latest status send then this should be a no-op
-                //
-                await SendCurrentResults();
-
-
-                // and finally:
-                // Make a call to send ResultCompletion and await its completion. This is just for backward compatibility with older protocol
-                //
-                await (ResultCompletion?.Invoke(this) ?? Task.CompletedTask);
             }
         }
 
@@ -582,19 +560,8 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 
         #endregion
 
-        #region Private Helper Methods
+        #region Public Helper Methods
-        /// <summary>
+        public async Task SendCurrentResults()
-        /// Sends the ResultsUpdated message if the number of rows has changed since last send.
-        /// </summary>
-        /// <param name="stateInfo"></param>
-        private void SendResultAvailableOrUpdated (object stateInfo = null)
-        {
-            // Make the call to send current results and synchronously wait for it to finish
-            //
-            SendCurrentResults().Wait();
-        }
-
-        private async Task SendCurrentResults()
         {
             try
             {
@@ -659,11 +626,29 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             }
             finally
             {
+                // and finally:
+                // Make a call to send ResultCompletion and await its completion. This is just for backward compatibility with older protocol
+                //
+                await (ResultCompletion?.Invoke(this) ?? Task.CompletedTask);
+
                 // Release the sendResultsSemphore so the next invocation gets unblocked
                 //
                 sendResultsSemphore.Release();
             }
         }
+        #endregion
+
+        #region Private Helper Methods
+        /// <summary>
+        /// Sends the ResultsUpdated message if the number of rows has changed since last send.
+        /// </summary>
+        /// <param name="stateInfo"></param>
+        private void SendResultAvailableOrUpdated(object stateInfo = null)
+        {
+            // Make the call to send current results and synchronously wait for it to finish
+            //
+            SendCurrentResults().Wait();
+        }
 
         private uint ResultsIntervalMultiplier { get; set; } = 1;