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);
}
///
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.
///
- public KustoResultsReader(IDataReader reader)
+ public KustoResultsReader(IDataReader[] readers)
: base()
{
// Read out all tables
List results = new List();
- 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());