mirror of
https://github.com/ckaczor/sqltoolsservice.git
synced 2026-01-13 17:23:02 -05:00
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:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user