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

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();