ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-25 17:24:17 -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;
}

9
src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingScriptOperation.cs
View File

@@ -20,18 +20,15 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
/// </summary>
public sealed class ScriptingScriptOperation : SmoScriptingOperation
{
private int scriptedObjectCount = 0;
private int totalScriptedObjectCount = 0;
private int eventSequenceNumber = 1;
private string azureAccessToken;
public ScriptingScriptOperation(ScriptingParams parameters, string azureAccessToken, IDataSourceFactory dataSourceFactory) : base(parameters, dataSourceFactory)
public ScriptingScriptOperation(ScriptingParams parameters, IDataSource dataSource) : base(parameters, dataSource)
{
this.azureAccessToken = azureAccessToken;
}
public override void Execute()
@@ -204,7 +201,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
selectedObjects.Count(),
string.Join(", ", selectedObjects)));
string server = GetServerNameFromLiveInstance(this.Parameters.ConnectionString, this.azureAccessToken);
string server = GetServerNameFromLiveInstance();
string database = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
foreach (ScriptingObject scriptingObject in selectedObjects)

48
src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingService.cs
View File

@@ -10,7 +10,6 @@ using System.Threading.Tasks;
using Microsoft.SqlTools.Hosting.Protocol;
using Microsoft.SqlTools.Hosting.Protocol.Contracts;
using Microsoft.Kusto.ServiceLayer.Connection;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Microsoft.Kusto.ServiceLayer.Scripting.Contracts;
using Microsoft.SqlTools.Utility;
using Microsoft.Kusto.ServiceLayer.Utility;
@@ -22,13 +21,11 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
/// </summary>
public sealed class ScriptingService : IDisposable
{
private const int ScriptingOperationTimeout = 60000;
private static readonly Lazy<ScriptingService> LazyInstance = new Lazy<ScriptingService>(() => new ScriptingService());
public static ScriptingService Instance => LazyInstance.Value;
private static ConnectionService connectionService;
private static ConnectionService _connectionService;
private readonly Lazy<ConcurrentDictionary<string, ScriptingOperation>> operations =
new Lazy<ConcurrentDictionary<string, ScriptingOperation>>(() => new ConcurrentDictionary<string, ScriptingOperation>());
@@ -36,26 +33,6 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
private bool disposed;
private IScripter _scripter;
private IDataSourceFactory _dataSourceFactory;
/// <summary>
/// Internal for testing purposes only
/// </summary>
internal static ConnectionService ConnectionServiceInstance
{
get
{
if (connectionService == null)
{
connectionService = ConnectionService.Instance;
}
return connectionService;
}
set
{
connectionService = value;
}
}
/// <summary>
/// The collection of active operations
@@ -66,11 +43,13 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
/// Initializes the Scripting Service instance
/// </summary>
/// <param name="serviceHost"></param>
/// <param name="context"></param>
public void InitializeService(ServiceHost serviceHost, IScripter scripter, IDataSourceFactory dataSourceFactory)
/// <param name="scripter"></param>
/// <param name="connectionService"></param>
public void InitializeService(ServiceHost serviceHost, IScripter scripter, ConnectionService connectionService)
{
_scripter = scripter;
_dataSourceFactory = dataSourceFactory;
_connectionService = connectionService;
serviceHost.SetRequestHandler(ScriptingRequest.Type, this.HandleScriptExecuteRequest);
serviceHost.SetRequestHandler(ScriptingCancelRequest.Type, this.HandleScriptCancelRequest);
serviceHost.SetRequestHandler(ScriptingListObjectsRequest.Type, this.HandleListObjectsRequest);
@@ -108,22 +87,18 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
/// </summary>
public async Task HandleScriptExecuteRequest(ScriptingParams parameters, RequestContext<ScriptingResult> requestContext)
{
SmoScriptingOperation operation = null;
try
{
// if a connection string wasn't provided as a parameter then
// use the owner uri property to lookup its associated ConnectionInfo
// and then build a connection string out of that
ConnectionInfo connInfo = null;
string accessToken = null;
if (parameters.ConnectionString == null)
{
ScriptingService.ConnectionServiceInstance.TryFindConnection(parameters.OwnerUri, out connInfo);
ConnectionInfo connInfo;
_connectionService.TryFindConnection(parameters.OwnerUri, out connInfo);
if (connInfo != null)
{
parameters.ConnectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails);
accessToken = connInfo.ConnectionDetails.AzureAccountToken;
}
else
{
@@ -131,13 +106,16 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
}
}
SmoScriptingOperation operation;
var datasource = _connectionService.GetOrOpenConnection(parameters.OwnerUri, ConnectionType.Default)
.Result.GetUnderlyingConnection();
if (!ShouldCreateScriptAsOperation(parameters))
{
operation = new ScriptingScriptOperation(parameters, accessToken, _dataSourceFactory);
operation = new ScriptingScriptOperation(parameters, datasource);
}
else
{
operation = new ScriptAsScriptingOperation(parameters, accessToken, _scripter, _dataSourceFactory);
operation = new ScriptAsScriptingOperation(parameters, _scripter, datasource);
}
operation.PlanNotification += (sender, e) => requestContext.SendEvent(ScriptingPlanNotificationEvent.Type, e).Wait();

20
src/Microsoft.Kusto.ServiceLayer/Scripting/SmoScriptingOperation.cs
View File

@@ -20,14 +20,13 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
/// </summary>
public abstract class SmoScriptingOperation : ScriptingOperation
{
protected readonly IDataSourceFactory _dataSourceFactory;
protected readonly IDataSource _dataSource;
private bool _disposed;
protected SmoScriptingOperation(ScriptingParams parameters, IDataSourceFactory dataSourceFactory)
protected SmoScriptingOperation(ScriptingParams parameters, IDataSource datasource)
{
_dataSourceFactory = dataSourceFactory;
_dataSource = datasource;
Validate.IsNotNull("parameters", parameters);
this.Parameters = parameters;
}
@@ -73,17 +72,10 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
parameters.OperationId = this.OperationId;
}
protected string GetServerNameFromLiveInstance(string connectionString, string azureAccessToken)
protected string GetServerNameFromLiveInstance()
{
string serverName = string.Empty;
using(var dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccessToken))
{
serverName = dataSource.ClusterName;
}
Logger.Write(TraceEventType.Verbose, string.Format("Resolved server name '{0}'", serverName));
return serverName;
Logger.Write(TraceEventType.Verbose, string.Format("Resolved server name '{0}'", _dataSource.ClusterName));
return _dataSource.ClusterName;
}
protected void ValidateScriptDatabaseParams()