ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-14 01:25:40 -05:00
Code Issues Packages Projects Releases Wiki Activity

Changed RefreshAuthToken in ConnectionService and KustoClient to follow TryParse pattern. (#1245)

Browse Source
This commit is contained in:
Justin M
2021-09-17 12:25:11 -07:00
committed by GitHub
parent 6f5cac4cb5
commit 60c847ff60
2 changed files with 27 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs
View File

@@ -213,9 +213,13 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
return completeParams;
}
internal string RefreshAuthToken(string ownerUri)
internal bool TryRefreshAuthToken(string ownerUri, out string token)
{
TryFindConnection(ownerUri, out ConnectionInfo connection);
token = string.Empty;
if (!TryFindConnection(ownerUri, out ConnectionInfo connection))
{
return false;
}
var requestMessage = new RequestSecurityTokenParams
{
@@ -227,8 +231,8 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
var response = _serviceHost.SendRequest(SecurityTokenRequest.Type, requestMessage, true).Result;
connection.UpdateAuthToken(response.Token);
return response.Token;
token = response.Token;
return true;
}
private void TryCloseConnectionTemporaryConnection(ConnectParams connectionParams, ConnectionInfo connectionInfo)

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

@@ -54,9 +54,13 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
DatabaseName = databaseName?.ToString() ?? "";
}
private void RefreshAuthToken()
private bool TryRefreshAuthToken()
{
string accountToken = ConnectionService.Instance.RefreshAuthToken(_ownerUri);
if (!ConnectionService.Instance.TryRefreshAuthToken(_ownerUri, out string accountToken))
{
return false;
}
_kustoQueryProvider.Dispose();
_kustoAdminProvider.Dispose();
@@ -69,6 +73,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
};
Initialize(connectionDetails);
return true;
}
private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(DataSourceConnectionDetails connectionDetails)
@@ -163,7 +168,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
exception.InnerException is KustoRequestException innerException
&& innerException.FailureCode == 401) // Unauthorized
{
RefreshAuthToken();
if (!TryRefreshAuthToken())
{
throw;
}
retryCount--;
return ExecuteQuery(query, cancellationToken, databaseName, retryCount);
}
@@ -184,7 +192,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
}
catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
{
RefreshAuthToken();
if (!TryRefreshAuthToken())
{
throw;
}
retryCount--;
await ExecuteControlCommandAsync(command, throwOnError, retryCount);
}
@@ -232,7 +243,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
}
catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
{
RefreshAuthToken();
if (!TryRefreshAuthToken())
{
throw;
}
retryCount--;
ExecuteControlCommand(command, retryCount);
}