ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-14 17:23:27 -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

32
src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptAsScriptingOperation.cs
View File

@@ -24,20 +24,15 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
{
private readonly IScripter _scripter;
private static readonly Dictionary<string, SqlServerVersion> scriptCompatibilityMap = LoadScriptCompatibilityMap();
private string _serverName;
private string _databaseName;
public ScriptAsScriptingOperation(ScriptingParams parameters, string azureAccountToken, IScripter scripter, IDataSourceFactory dataSourceFactory) : base(parameters, dataSourceFactory)
public ScriptAsScriptingOperation(ScriptingParams parameters, IScripter scripter, IDataSource datasource) :
base(parameters, datasource)
{
DataSource = _dataSourceFactory.Create(DataSourceType.Kusto, this.Parameters.ConnectionString,
azureAccountToken);
_scripter = scripter;
}
internal IDataSource DataSource { get; set; }
private string serverName;
private string databaseName;
private bool disconnectAtDispose = false;
public override void Execute()
{
try
@@ -49,7 +44,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
this.CancellationToken.ThrowIfCancellationRequested();
string resultScript = string.Empty;
UrnCollection urns = CreateUrns(DataSource);
UrnCollection urns = CreateUrns(_dataSource);
ScriptingOptions options = new ScriptingOptions();
SetScriptBehavior(options);
ScriptAsOptions scriptAsOptions = new ScriptAsOptions(this.Parameters.ScriptOptions);
@@ -65,12 +60,12 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
switch (this.Parameters.Operation)
{
case ScriptingOperationType.Select:
resultScript = GenerateScriptSelect(DataSource, urns);
resultScript = GenerateScriptSelect(_dataSource, urns);
break;
case ScriptingOperationType.Alter:
case ScriptingOperationType.Execute:
resultScript = GenerateScriptForFunction(DataSource);
resultScript = GenerateScriptForFunction(_dataSource);
break;
}
@@ -118,13 +113,6 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
});
}
}
finally
{
if (disconnectAtDispose && DataSource != null)
{
DataSource.Dispose();
}
}
}
private string GenerateScriptSelect(IDataSource dataSource, UrnCollection urns)
@@ -168,8 +156,8 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
{
IEnumerable<ScriptingObject> selectedObjects = new List<ScriptingObject>(this.Parameters.ScriptingObjects);
serverName = dataSource.ClusterName;
databaseName = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
_serverName = dataSource.ClusterName;
_databaseName = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
UrnCollection urnCollection = new UrnCollection();
foreach (var scriptingObject in selectedObjects)
{
@@ -178,7 +166,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
// TODO: get the default schema
scriptingObject.Schema = "dbo";
}
urnCollection.Add(scriptingObject.ToUrn(serverName, databaseName));
urnCollection.Add(scriptingObject.ToUrn(_serverName, _databaseName));
}
return urnCollection;
}