3644 Kusto Azure Token Refresh (#1081)

* 3644 Created KustoClient to handle interaction with Kusto servers. Refactored KustoDataSource and moved execution logic into KustoClient.

* 3644 Added RequestSecurityTokenParams, RequestSecurityTokenResponse, and SecurityTokenRequest to Kusto. Moved intellisense functions from KustoDataSource to KustoIntellisenseHelper. Added SchemaState to readonly property on KustoDataSource. Added catch block to Batch.cs to catch KustoRequestExceptions.

* 3644 Removed unused reference from ConnectionDetails and ConnectedBindingContext. Added UpdateAzureToken function to IKustoClient, KustoClient, IDataSource, KustoDataSource, and KustoDataSource. Added dataSource.Dispose to ReliableDataSourceConnection > Close. Added RefreshAzureToken to ConnectionService to refresh azure token.

* 3644 Removed unused properties from RequestSecurityTokenParams and RequestSecurityTokenResponse

* 3644 Added default to DbColumnWrapper to UnknownTypeName when null. Moved database query logic from KustoIntellisenseHelper to KustoClient. Moved KustoIntellisenseHelper data objects out of class. Changed SchemaState to load through separate tasks.

* 3644 Changed ReRunQuery logic in Kusto Batch to flip back to false if the query fails a second time so it can be rejected

* 3644 Updated GetAutoCompleteSuggestions and GetHoverHelp with changes from main

* 3644 Added AccountId to RequestSecurityTokenParams and set value in ConnectionService. Added throw to Batch.cs to ensure exceptions that are non-Unauthorized bubble up to ADS.

* 3644 Changed KustoUnauthorizedException to take original exception as inner exception. Changed catch block to only throw KustoUnauthorizedException when FailureCode is 401

* 3644 Renamed KustoUnauthorizedException to DataSourceUnauthorizedException. Moved logic to throw exception down into KustoClient. Changed retryLogic in Batch.cs to a decrementing count

* 3644 Changed logic in Batch.cs for throwing InvalidOperationException
This commit is contained in:
Justin M
2020-09-28 13:24:23 -07:00
committed by GitHub
parent d6ff73d510
commit d061e781f4
23 changed files with 673 additions and 487 deletions

View File

@@ -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.Exceptions;
namespace Microsoft.Kusto.ServiceLayer.QueryExecution
{
@@ -67,6 +66,8 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
/// </summary>
private readonly bool getFullColumnSchema;
private int _retryCount;
#endregion
internal Batch(string batchText, SelectionData selection, int ordinalId,
@@ -87,7 +88,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
this.outputFileFactory = outputFileFactory;
specialAction = new SpecialAction();
BatchExecutionCount = executionCount > 0 ? executionCount : 1;
_retryCount = 1;
this.getFullColumnSchema = getFullColumnSchema;
}
@@ -251,7 +252,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
public async Task Execute(ReliableDataSourceConnection conn, CancellationToken cancellationToken)
{
// Sanity check to make sure we haven't already run this batch
if (HasExecuted)
if (HasExecuted && _retryCount < 0)
{
throw new InvalidOperationException("Batch has already executed.");
}
@@ -266,6 +267,12 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
{
await DoExecute(conn, cancellationToken);
}
catch (DataSourceUnauthorizedException)
{
// Rerun the query once if unauthorized
_retryCount--;
throw;
}
catch (TaskCanceledException)
{
// Cancellation isn't considered an error condition
@@ -290,7 +297,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
await BatchCompletion(this);
}
}
}
private async Task DoExecute(ReliableDataSourceConnection conn, CancellationToken cancellationToken)
@@ -308,6 +314,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
{
await ExecuteOnce(conn, cancellationToken);
}
catch (DbException dbe)
{
HasError = true;