ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-13 17:23:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

Fix for multiple queries separated by space for Kusto (#1089)

Browse Source 
* Fix for multiple queries separated by space

* Addressed comments

Co-authored-by: Monica Gupta <mogupt@microsoft.com>
This commit is contained in:
Monica Gupta
2020-10-04 19:12:33 -07:00
committed by GitHub
parent 725d4ba5db
commit 164ca951da
2 changed files with 36 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs
View File

@@ -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>

27
src/Microsoft.Kusto.ServiceLayer/DataSource/KustoResultsReader.cs
View File

@@ -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());