ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-02-16 10:58:30 -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); clientRequestProperties.SetOption(ClientRequestProperties.OptionNoTruncation, true);
cancellationToken.Register(() => CancelQuery(clientRequestProperties.ClientRequestId)); cancellationToken.Register(() => CancelQuery(clientRequestProperties.ClientRequestId));
var kustoCodeService = new KustoCodeService(query); var script = CodeScript.From(query, GlobalState.Default);
var minimalQuery = kustoCodeService.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments); IDataReader[] origReaders = new IDataReader[script.Blocks.Count];
try Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
{ {
IDataReader origReader = _kustoQueryProvider.ExecuteQuery( var minimalQuery = codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
KustoQueryUtils.IsClusterLevelQuery(minimalQuery) ? "" : databaseName,
minimalQuery,
clientRequestProperties);
return new KustoResultsReader(origReader); try
} {
catch (KustoRequestException exception) when (exception.FailureCode == 401) // Unauthorized IDataReader origReader = _kustoQueryProvider.ExecuteQuery(
{ KustoQueryUtils.IsClusterLevelQuery(minimalQuery) ? "" : databaseName,
throw new DataSourceUnauthorizedException(exception); minimalQuery,
} clientRequestProperties);
origReaders[index] = origReader;
}
catch (KustoRequestException exception) when (exception.FailureCode == 401) // Unauthorized
{
throw new DataSourceUnauthorizedException(exception);
}
});
return new KustoResultsReader(origReaders);
} }
/// <summary> /// <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. /// 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. /// ADS just needs query results. When returning query results we need to trim off the last 3 tables.
/// </summary> /// </summary>
public KustoResultsReader(IDataReader reader) public KustoResultsReader(IDataReader[] readers)
: base() : base()
{ {
// Read out all tables // Read out all tables
List<DataTable> results = new List<DataTable>(); 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 foreach (var reader in readers)
if(results.Count > 3) results.RemoveRange(results.Count - 3, 3); {
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 // Create a DataReader for the trimmed set
_resultDataSet = new DataSet(); _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()); SetDataReader(_resultDataSet.CreateDataReader());