ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-29 09:35:38 -05:00
Code Issues Packages Projects Releases Wiki Activity

3490 Kusto Connection Refresh Fix (#1085)

Browse Source 
* 3490 Injected OwnerUri into KustClient to store for token refreshing. Removed UpdateAzureToken from IDataSource, DataSourceBase, and KustoDataSource. Removed logic for retrying queries related to Unauthorized datasource in Batch and Query. Changed ScriptingService, ScriptingScriptOperation, and ScriptAsScriptingOperation to take DataSource in the constructor instead of datasourcefactory. Changed ScriptingService to inject ConnectionService through InitializeService function.

* 3490 Removed Catch block for DataSourceUnauthorizedException in ExecuteControlCommandAsync

* 3490 Removed OwnerUri from KustoClient and used azureAccountToken to refresh token in ConnectionService.

* 3490 Reverted unneeded changes.

* 3490 Split ExecuteQuery in KustoClient to execute first query then remaining queries after

* 3490 Passed OwnerUri down into KustoClient to refresh token.

* 3490 Removed DataSourceUnauthorizedException. Refactored ExecuteQuery to catch aggregate exception. Added RefreshAzureToken logic to ExecuteControlCommand

* 3490 Added logic to update ReliableDataSourceConnection azure token within ConnectionInfo.

* 3490 Add retry logic to ExecuteQuery and ExecuteControlCommand in KustoClient
This commit is contained in:
Justin M
2020-10-06 17:20:13 -07:00
committed by GitHub
parent 2b022d2e48
commit f2eb590d97
25 changed files with 120 additions and 212 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceBase.cs
View File

@@ -89,8 +89,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
public abstract string GenerateExecuteFunctionScript(string functionName);
public abstract void UpdateAzureToken(string azureToken);
/// <inheritdoc/>
public DataSourceType DataSourceType { get; protected set; }

4
src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceFactory.cs
View File

@@ -13,7 +13,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
[Export(typeof(IDataSourceFactory))]
public class DataSourceFactory : IDataSourceFactory
{
public IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken)
public IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken, string ownerUri)
{
ValidationUtils.IsArgumentNotNullOrWhiteSpace(connectionString, nameof(connectionString));
ValidationUtils.IsArgumentNotNullOrWhiteSpace(azureAccountToken, nameof(azureAccountToken));
@@ -22,7 +22,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
{
case DataSourceType.Kusto:
{
var kustoClient = new KustoClient(connectionString, azureAccountToken);
var kustoClient = new KustoClient(connectionString, azureAccountToken, ownerUri);
return new KustoDataSource(kustoClient);
}

11
src/Microsoft.Kusto.ServiceLayer/DataSource/Exceptions/DataSourceUnauthorizedException.cs
View File

@@ -1,11 +0,0 @@
using System;
namespace Microsoft.Kusto.ServiceLayer.DataSource.Exceptions
{
public class DataSourceUnauthorizedException : Exception
{
public DataSourceUnauthorizedException(Exception ex) : base (ex.Message, ex)
{
}
}
}

6
src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSource.cs
View File

@@ -111,11 +111,5 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
/// <param name="functionName"></param>
/// <returns></returns>
string GenerateExecuteFunctionScript(string functionName);
/// <summary>
/// Updates Azure Token
/// </summary>
/// <param name="azureToken"></param>
void UpdateAzureToken(string azureToken);
}
}

2
src/Microsoft.Kusto.ServiceLayer/DataSource/IDataSourceFactory.cs
View File

@@ -2,6 +2,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
{
public interface IDataSourceFactory
{
IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken);
IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken, string ownerUri);
}
}

7
src/Microsoft.Kusto.ServiceLayer/DataSource/IKustoClient.cs
View File

@@ -17,9 +17,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
string DatabaseName { get; }
void UpdateAzureToken(string azureAccountToken);
IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null);
IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null, int retryCount = 1);
/// <summary>
/// Executes a query or command against a kusto cluster and returns a sequence of result row instances.
@@ -37,7 +35,8 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
/// Executes a Kusto control command.
/// </summary>
/// <param name="command">The command.</param>
void ExecuteControlCommand(string command);
/// <param name="retryCount"></param>
void ExecuteControlCommand(string command, int retryCount = 1);
void UpdateDatabase(string databaseName);

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

@@ -14,14 +14,16 @@ using Kusto.Data.Net.Client;
using Kusto.Language;
using Kusto.Language.Editor;
using Microsoft.Data.SqlClient;
using Microsoft.Kusto.ServiceLayer.Connection;
using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
using Microsoft.Kusto.ServiceLayer.DataSource.Exceptions;
using Microsoft.Kusto.ServiceLayer.Utility;
namespace Microsoft.Kusto.ServiceLayer.DataSource
{
public class KustoClient : IKustoClient
{
private readonly string _ownerUri;
[DebuggerBrowsable(DebuggerBrowsableState.Never)]
private ICslAdminProvider _kustoAdminProvider;
@@ -36,8 +38,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
public string ClusterName { get; }
public string DatabaseName { get; private set; }
public KustoClient(string connectionString, string azureAccountToken)
public KustoClient(string connectionString, string azureAccountToken, string ownerUri)
{
_ownerUri = ownerUri;
ClusterName = GetClusterName(connectionString);
var databaseName = new SqlConnectionStringBuilder(connectionString).InitialCatalog;
Initialize(ClusterName, databaseName, azureAccountToken);
@@ -75,8 +78,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
_kustoAdminProvider = KustoClientFactory.CreateCslAdminProvider(stringBuilder);
}
public void UpdateAzureToken(string azureAccountToken)
private void RefreshAzureToken()
{
string azureAccountToken = ConnectionService.Instance.RefreshAzureToken(_ownerUri);
_kustoQueryProvider.Dispose();
_kustoAdminProvider.Dispose();
Initialize(ClusterName, DatabaseName, azureAccountToken);
@@ -162,7 +166,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
}
}
public IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null)
public IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null, int retryCount = 1)
{
ValidationUtils.IsArgumentNotNullOrWhiteSpace(query, nameof(query));
@@ -175,29 +179,33 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
var script = CodeScript.From(query, GlobalState.Default);
IDataReader[] origReaders = new IDataReader[script.Blocks.Count];
Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
try
{
var minimalQuery = codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
try
Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
{
var minimalQuery =
codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
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);
});
return new KustoResultsReader(origReaders);
}
catch (AggregateException exception)
when (retryCount > 0 &&
exception.InnerException is KustoRequestException innerException
&& innerException.FailureCode == 401) // Unauthorized
{
RefreshAzureToken();
retryCount--;
return ExecuteQuery(query, cancellationToken, databaseName, retryCount);
}
}
/// <summary>
/// Executes a query or command against a kusto cluster and returns a sequence of result row instances.
/// </summary>
@@ -211,10 +219,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
var tableReader = results[WellKnownDataSet.PrimaryResult].Single().TableData.CreateDataReader();
return new ObjectReader<T>(tableReader);
}
catch (DataSourceUnauthorizedException)
{
throw;
}
catch (Exception) when (!throwOnError)
{
return null;
@@ -243,12 +247,22 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
/// Executes a Kusto control command.
/// </summary>
/// <param name="command">The command.</param>
public void ExecuteControlCommand(string command)
/// <param name="retryCount"></param>
public void ExecuteControlCommand(string command, int retryCount = 1)
{
ValidationUtils.IsArgumentNotNullOrWhiteSpace(command, nameof(command));
using (var adminOutput = _kustoAdminProvider.ExecuteControlCommand(command, null))
try
{
using (var adminOutput = _kustoAdminProvider.ExecuteControlCommand(command, null))
{
}
}
catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
{
RefreshAzureToken();
retryCount--;
ExecuteControlCommand(command, retryCount);
}
}

4
src/Microsoft.Kusto.ServiceLayer/DataSource/KustoDataSource.cs
View File

@@ -796,10 +796,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
return string.IsNullOrWhiteSpace(objectName) ? databaseName : $"{databaseName}.{objectName}";
}
public override void UpdateAzureToken(string azureToken)
{
_kustoClient.UpdateAzureToken(azureToken);
}
#endregion
}
}

16
src/Microsoft.Kusto.ServiceLayer/DataSource/ReliableDataSourceConnection.cs
View File

@@ -42,8 +42,9 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
private readonly Guid _azureSessionId = Guid.NewGuid();
private readonly string _connectionString;
private readonly string _azureAccountToken;
private string _azureAccountToken;
private readonly IDataSourceFactory _dataSourceFactory;
private readonly string _ownerUri;
/// <summary>
/// Initializes a new instance of the ReliableKustoClient class with a given connection string
@@ -55,13 +56,15 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
/// <param name="commandRetryPolicy">The retry policy defining whether to retry a request if a command fails to be executed.</param>
/// <param name="azureAccountToken"></param>
/// <param name="dataSourceFactory"></param>
/// <param name="ownerUri"></param>
public ReliableDataSourceConnection(string connectionString, RetryPolicy connectionRetryPolicy,
RetryPolicy commandRetryPolicy, string azureAccountToken, IDataSourceFactory dataSourceFactory)
RetryPolicy commandRetryPolicy, string azureAccountToken, IDataSourceFactory dataSourceFactory, string ownerUri)
{
_connectionString = connectionString;
_azureAccountToken = azureAccountToken;
_dataSourceFactory = dataSourceFactory;
_dataSource = dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccountToken);
_ownerUri = ownerUri;
_dataSource = dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccountToken, ownerUri);
_connectionRetryPolicy = connectionRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
_commandRetryPolicy = commandRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
@@ -190,7 +193,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
{
_connectionRetryPolicy.ExecuteAction(() =>
{
_dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, _connectionString, _azureAccountToken);
_dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, _connectionString, _azureAccountToken, _ownerUri);
});
}
}
@@ -247,6 +250,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
{
get { return _dataSource.DatabaseName; }
}
public void UpdateAzureToken(string token)
{
_azureAccountToken = token;
}
}
}