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

dSTS Authentication (#1125)

Browse Source 
* Refactored Kusto.ServiceLayer to pass ConnectionDetails to DataSourceFactory instead of connection string. Created KustoConnectionDetails to map needed details to KustoClient.

* Removed unused ScriptingScriptOperation from KustoServiceLayer.

* Created DstsAuthenticationManager and moved logic for getting DstsToken. Updated error message for failing to create KustoConnection.

* Removed DstsAuthenticationManager.cs. Refactored DataSourceFactory to retrieve UserToken from ConnectionDetails.

* Renamed AzureAccountToken in ConnectionDetails to AccountToken. Changed mapping to KustoConnectionDetails based on the AccountToken.

* Removed Kusto.Data reference from ConnectionService and ScriptingListObjectsOperation. Moved creation of KustoConnectionStringBuilder to DataSourceFactory

* Added accountToken validation to DataSourceFactory Create.

* Renamed KustoConnectionDetails to DataSourceConnectionDetails. Renamed AzureToken to AuthToken.
This commit is contained in:
Justin M
2021-01-14 13:49:09 -08:00
committed by GitHub
parent 822ffb2908
commit f0a5e11d51
22 changed files with 174 additions and 438 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/Microsoft.Kusto.ServiceLayer/DataSource/Contracts/DataSourceConnectionDetails.cs Normal file
View File

@@ -0,0 +1,11 @@
namespace Microsoft.Kusto.ServiceLayer.DataSource.Contracts
{
public class DataSourceConnectionDetails
{
public string ServerName { get; set; }
public string DatabaseName { get; set; }
public string UserToken { get; set; }
public string ConnectionString { get; set; }
public string AuthenticationType { get; set; }
}
}

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

@@ -1,33 +1,79 @@
using System;
using System.Collections.Generic;
using System.Composition;
using Microsoft.Kusto.ServiceLayer.Utility;
using Kusto.Data;
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
using Microsoft.Kusto.ServiceLayer.DataSource.Contracts;
using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
using Microsoft.Kusto.ServiceLayer.Workspace.Contracts;
using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
using Microsoft.Kusto.ServiceLayer.Utility;
namespace Microsoft.Kusto.ServiceLayer.DataSource
{
[Export(typeof(IDataSourceFactory))]
public class DataSourceFactory : IDataSourceFactory
{
public IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken, string ownerUri)
public IDataSource Create(DataSourceType dataSourceType, ConnectionDetails connectionDetails, string ownerUri)
{
ValidationUtils.IsArgumentNotNullOrWhiteSpace(connectionString, nameof(connectionString));
ValidationUtils.IsArgumentNotNullOrWhiteSpace(azureAccountToken, nameof(azureAccountToken));
ValidationUtils.IsArgumentNotNullOrWhiteSpace(connectionDetails.AccountToken, nameof(connectionDetails.AccountToken));
switch (dataSourceType)
{
case DataSourceType.Kusto:
{
var kustoClient = new KustoClient(connectionString, azureAccountToken, ownerUri);
var kustoConnectionDetails = MapKustoConnectionDetails(connectionDetails);
var kustoClient = new KustoClient(kustoConnectionDetails, ownerUri);
return new KustoDataSource(kustoClient);
}
default:
throw new ArgumentException($"Unsupported data source type \"{dataSourceType}\"",
throw new ArgumentException($@"Unsupported data source type ""{dataSourceType}""",
nameof(dataSourceType));
}
}
private DataSourceConnectionDetails MapKustoConnectionDetails(ConnectionDetails connectionDetails)
{
return new DataSourceConnectionDetails
{
ServerName = connectionDetails.ServerName,
DatabaseName = connectionDetails.DatabaseName,
ConnectionString = connectionDetails.ConnectionString,
AuthenticationType = connectionDetails.AuthenticationType,
UserToken = connectionDetails.AccountToken
};
}
public static KustoConnectionStringBuilder CreateConnectionStringBuilder(DataSourceType dataSourceType, string serverName, string databaseName)
{
switch (dataSourceType)
{
case DataSourceType.Kusto:
{
return new KustoConnectionStringBuilder(serverName, databaseName);
}
default:
throw new ArgumentException($@"Unsupported data source type ""{dataSourceType}""",
nameof(dataSourceType));
}
}
public static KustoConnectionStringBuilder CreateConnectionStringBuilder(DataSourceType dataSourceType, string connectionString)
{
switch (dataSourceType)
{
case DataSourceType.Kusto:
{
return new KustoConnectionStringBuilder(connectionString);
}
default:
throw new ArgumentException($@"Unsupported data source type ""{dataSourceType}""",
nameof(dataSourceType));
}
}

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

@@ -1,7 +1,9 @@
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
namespace Microsoft.Kusto.ServiceLayer.DataSource
{
public interface IDataSourceFactory
{
IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken, string ownerUri);
IDataSource Create(DataSourceType dataSourceType, ConnectionDetails connectionDetails, string ownerUri);
}
}

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

@@ -15,6 +15,7 @@ using Kusto.Data.Net.Client;
using Kusto.Language;
using Kusto.Language.Editor;
using Microsoft.Kusto.ServiceLayer.Connection;
using Microsoft.Kusto.ServiceLayer.DataSource.Contracts;
using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
using Microsoft.Kusto.ServiceLayer.Utility;
@@ -38,10 +39,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
public string ClusterName { get; private set; }
public string DatabaseName { get; private set; }
public KustoClient(string connectionString, string azureAccountToken, string ownerUri)
public KustoClient(DataSourceConnectionDetails connectionDetails, string ownerUri)
{
_ownerUri = ownerUri;
Initialize(azureAccountToken, connectionString);
Initialize(connectionDetails);
SchemaState = LoadSchemaState();
}
@@ -79,29 +80,38 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
DatabaseName, ClusterName);
}
private void Initialize(string azureAccountToken, string connectionString = "")
private void Initialize(DataSourceConnectionDetails connectionDetails)
{
var stringBuilder = GetKustoConnectionStringBuilder(azureAccountToken, connectionString);
var stringBuilder = GetKustoConnectionStringBuilder(connectionDetails);
_kustoQueryProvider = KustoClientFactory.CreateCslQueryProvider(stringBuilder);
_kustoAdminProvider = KustoClientFactory.CreateCslAdminProvider(stringBuilder);
}
private void RefreshAzureToken()
private void RefreshAuthToken()
{
string azureAccountToken = ConnectionService.Instance.RefreshAzureToken(_ownerUri);
string accountToken = ConnectionService.Instance.RefreshAzureToken(_ownerUri);
_kustoQueryProvider.Dispose();
_kustoAdminProvider.Dispose();
Initialize(azureAccountToken);
var connectionDetails = new DataSourceConnectionDetails
{
ServerName = ClusterName,
DatabaseName = DatabaseName,
UserToken = accountToken,
AuthenticationType = "AzureMFA"
};
Initialize(connectionDetails);
}
private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(string userToken, string connectionString)
private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(DataSourceConnectionDetails connectionDetails)
{
ValidationUtils.IsTrue<ArgumentException>(!string.IsNullOrWhiteSpace(userToken),
$"the Kusto authentication is not specified - either set {nameof(userToken)}");
var stringBuilder = string.IsNullOrWhiteSpace(connectionString)
? new KustoConnectionStringBuilder(ClusterName, DatabaseName)
: new KustoConnectionStringBuilder(connectionString);
ValidationUtils.IsTrue<ArgumentException>(!string.IsNullOrWhiteSpace(connectionDetails.UserToken),
$"The Kusto User Token is not specified - set {nameof(connectionDetails.UserToken)}");
var stringBuilder = string.IsNullOrWhiteSpace(connectionDetails.ConnectionString)
? new KustoConnectionStringBuilder(connectionDetails.ServerName, connectionDetails.DatabaseName)
: new KustoConnectionStringBuilder(connectionDetails.ConnectionString);
ClusterName = stringBuilder.DataSource;
var databaseName = ParseDatabaseName(stringBuilder.InitialCatalog);
@@ -110,7 +120,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
ValidationUtils.IsNotNull(ClusterName, nameof(ClusterName));
return stringBuilder.WithAadUserTokenAuthentication(userToken);
return connectionDetails.AuthenticationType == "dstsAuth"
? stringBuilder.WithDstsUserTokenAuthentication(connectionDetails.UserToken)
: stringBuilder.WithAadUserTokenAuthentication(connectionDetails.UserToken);
}
private ClientRequestProperties GetClientRequestProperties(CancellationToken cancellationToken)
@@ -181,7 +193,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
exception.InnerException is KustoRequestException innerException
&& innerException.FailureCode == 401) // Unauthorized
{
RefreshAzureToken();
RefreshAuthToken();
retryCount--;
return ExecuteQuery(query, cancellationToken, databaseName, retryCount);
}
@@ -202,7 +214,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
}
catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
{
RefreshAzureToken();
RefreshAuthToken();
retryCount--;
await ExecuteControlCommandAsync(command, throwOnError, retryCount);
}
@@ -254,7 +266,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
}
catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
{
RefreshAzureToken();
RefreshAuthToken();
retryCount--;
ExecuteControlCommand(command, retryCount);
}

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

@@ -24,6 +24,7 @@ using System;
using System.Data.SqlClient;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
using Microsoft.SqlTools.Utility;
using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
using Microsoft.Kusto.ServiceLayer.DataSource;
@@ -41,8 +42,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
private RetryPolicy _commandRetryPolicy;
private readonly Guid _azureSessionId = Guid.NewGuid();
private readonly string _connectionString;
private string _azureAccountToken;
private readonly ConnectionDetails _connectionDetails;
private readonly IDataSourceFactory _dataSourceFactory;
private readonly string _ownerUri;
@@ -51,20 +51,18 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
/// and a policy defining whether to retry a request if the connection fails to be opened or a command
/// fails to be successfully executed.
/// </summary>
/// <param name="connectionString">The connection string used to open the SQL Azure database.</param>
/// <param name="connectionDetails"></param>
/// <param name="connectionRetryPolicy">The retry policy defining whether to retry a request if a connection fails to be established.</param>
/// <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, string ownerUri)
public ReliableDataSourceConnection(ConnectionDetails connectionDetails, RetryPolicy connectionRetryPolicy,
RetryPolicy commandRetryPolicy, IDataSourceFactory dataSourceFactory, string ownerUri)
{
_connectionString = connectionString;
_azureAccountToken = azureAccountToken;
_connectionDetails = connectionDetails;
_dataSourceFactory = dataSourceFactory;
_ownerUri = ownerUri;
_dataSource = dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccountToken, ownerUri);
_dataSource = dataSourceFactory.Create(DataSourceType.Kusto, connectionDetails, ownerUri);
_connectionRetryPolicy = connectionRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
_commandRetryPolicy = commandRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
@@ -193,7 +191,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
{
_connectionRetryPolicy.ExecuteAction(() =>
{
_dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, _connectionString, _azureAccountToken, _ownerUri);
_dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, _connectionDetails, _ownerUri);
});
}
}
@@ -251,9 +249,9 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
get { return _dataSource.DatabaseName; }
}
public void UpdateAzureToken(string token)
public void UpdateAuthToken(string token)
{
_azureAccountToken = token;
_connectionDetails.AccountToken = token;
}
}
}