ckaczor/sqltoolsservice
1
0
Fork 0
mirror of https://github.com/ckaczor/sqltoolsservice.git synced 2026-01-26 09:35:38 -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

6
src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionInfo.cs
View File

@@ -158,13 +158,13 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
}
}
public void UpdateAzureToken(string token)
public void UpdateAuthToken(string token)
{
ConnectionDetails.AzureAccountToken = token;
ConnectionDetails.AccountToken = token;
foreach (var connection in _connectionTypeToConnectionMap.Values)
{
connection.UpdateAzureToken(token);
connection.UpdateAuthToken(token);
}
}
}

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

@@ -19,7 +19,6 @@ using Microsoft.Kusto.ServiceLayer.LanguageServices.Contracts;
using Microsoft.Kusto.ServiceLayer.Utility;
using Microsoft.SqlTools.Utility;
using System.Diagnostics;
using Kusto.Data;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
@@ -280,7 +279,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
};
var response = Instance.ServiceHost.SendRequest(SecurityTokenRequest.Type, requestMessage, true).Result;
connection.UpdateAzureToken(response.Token);
connection.UpdateAuthToken(response.Token);
return response.Token;
}
@@ -389,12 +388,13 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
if (!string.IsNullOrEmpty(connectionInfo.ConnectionDetails.ConnectionString))
{
// If the connection was set up with a connection string, use the connection string to get the details
var connectionString = new KustoConnectionStringBuilder(connection.ConnectionString);
var connectionStringBuilder = DataSourceFactory.CreateConnectionStringBuilder(DataSourceType.Kusto, connection.ConnectionString);
response.ConnectionSummary = new ConnectionSummary
{
ServerName = connectionString.DataSource,
DatabaseName = connectionString.InitialCatalog,
UserName = connectionString.UserID
ServerName = connectionStringBuilder.DataSource,
DatabaseName = connectionStringBuilder.InitialCatalog,
UserName = connectionStringBuilder.UserID
};
}
else
@@ -446,11 +446,9 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
try
{
connectionInfo.ConnectionDetails.Pooling = false;
// build the connection string from the input parameters
string connectionString = BuildConnectionString(connectionInfo.ConnectionDetails);
// create a sql connection instance
connection = connectionInfo.Factory.CreateDataSourceConnection(connectionString, connectionInfo.ConnectionDetails.AzureAccountToken, connectionInfo.OwnerUri);
// create a data source connection instance
connection = connectionInfo.Factory.CreateDataSourceConnection(connectionInfo.ConnectionDetails, connectionInfo.OwnerUri);
connectionInfo.AddConnection(connectionParams.Type, connection);
// Add a cancellation token source so that the connection OpenAsync() can be cancelled
@@ -955,28 +953,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
}
}
/// <summary>
/// Build a connection string from a connection details instance
/// </summary>
/// <param name="connectionDetails"></param>
public static string BuildConnectionString(ConnectionDetails connectionDetails)
{
return CreateConnectionStringBuilder(connectionDetails).ToString();
}
/// <summary>
/// Build a connection string builder a connection details instance
/// </summary>
/// <param name="connectionDetails"></param>
private static KustoConnectionStringBuilder CreateConnectionStringBuilder(ConnectionDetails connectionDetails)
{
var stringBuilder = string.IsNullOrWhiteSpace(connectionDetails.ConnectionString)
? new KustoConnectionStringBuilder(connectionDetails.ServerName, connectionDetails.DatabaseName)
: new KustoConnectionStringBuilder(connectionDetails.ConnectionString);
return stringBuilder.WithAadUserTokenAuthentication(connectionDetails.AzureAccountToken);
}
/// <summary>
/// Handles a request to get a connection string for the provided connection
/// </summary>
@@ -997,9 +973,9 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
info.ConnectionDetails.Password = ConnectionService.PasswordPlaceholder;
}
info.ConnectionDetails.ApplicationName = "sqlops-connection-string";
connectionString = BuildConnectionString(info.ConnectionDetails);
info.ConnectionDetails.ApplicationName = "ads-connection-string";
connectionString = DataSourceFactory.CreateConnectionStringBuilder(DataSourceType.Kusto,
info.ConnectionDetails.ServerName, info.ConnectionDetails.DatabaseName).ToString();
}
catch (Exception e)
{
@@ -1035,7 +1011,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
public ConnectionDetails ParseConnectionString(string connectionString)
{
var builder = new KustoConnectionStringBuilder(connectionString);
var builder = DataSourceFactory.CreateConnectionStringBuilder(DataSourceType.Kusto, connectionString);
return new ConnectionDetails
{
ApplicationName = builder.ApplicationNameForTracing,
@@ -1082,10 +1058,8 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
conn.Dispose();
info.RemoveConnection(key);
string connectionString = BuildConnectionString(info.ConnectionDetails);
// create a sql connection instance
ReliableDataSourceConnection connection = info.Factory.CreateDataSourceConnection(connectionString, info.ConnectionDetails.AzureAccountToken, ownerUri);
// create a kusto connection instance
ReliableDataSourceConnection connection = info.Factory.CreateDataSourceConnection(info.ConnectionDetails, ownerUri);
connection.Open();
info.AddConnection(key, connection);
}
@@ -1094,7 +1068,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
conn.ChangeDatabase(newDatabaseName);
}
}
}
// Fire a connection changed event

4
src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetails.cs
View File

@@ -496,7 +496,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts
}
}
public string AzureAccountToken
public string AccountToken
{
get
{
@@ -518,7 +518,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts
if (ServerName != other.ServerName
|| AuthenticationType != other.AuthenticationType
|| UserName != other.UserName
|| AzureAccountToken != other.AzureAccountToken)
|| AccountToken != other.AccountToken)
{
return false;
}

2
src/Microsoft.Kusto.ServiceLayer/Connection/Contracts/ConnectionDetailsExtensions.cs
View File

@@ -45,7 +45,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection.Contracts
TypeSystemVersion = details.TypeSystemVersion,
ConnectionString = details.ConnectionString,
Port = details.Port,
AzureAccountToken = details.AzureAccountToken
AccountToken = details.AccountToken
};
}
}

5
src/Microsoft.Kusto.ServiceLayer/Connection/DataSourceConnectionFactory.cs
View File

@@ -4,6 +4,7 @@
//
using System.Composition;
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Microsoft.SqlTools.ServiceLayer.Connection.ReliableConnection;
@@ -28,11 +29,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
/// <summary>
/// Creates a new SqlConnection object
/// </summary>
public ReliableDataSourceConnection CreateDataSourceConnection(string connectionString, string azureAccountToken, string ownerUri)
public ReliableDataSourceConnection CreateDataSourceConnection(ConnectionDetails connectionDetails, string ownerUri)
{
RetryPolicy connectionRetryPolicy = RetryPolicyFactory.CreateDefaultConnectionRetryPolicy();
RetryPolicy commandRetryPolicy = RetryPolicyFactory.CreateDefaultConnectionRetryPolicy();
return new ReliableDataSourceConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken, _dataSourceFactory, ownerUri);
return new ReliableDataSourceConnection(connectionDetails, connectionRetryPolicy, commandRetryPolicy, _dataSourceFactory, ownerUri);
}
}
}

6
src/Microsoft.Kusto.ServiceLayer/Connection/IDataSourceConnectionFactory.cs
View File

@@ -3,7 +3,9 @@
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
//
namespace Microsoft.Kusto.ServiceLayer.Connection
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
namespace Microsoft.Kusto.ServiceLayer.Connection
{
/// <summary>
/// Interface for the SQL Connection factory
@@ -13,6 +15,6 @@
/// <summary>
/// Create a new SQL Connection object
/// </summary>
ReliableDataSourceConnection CreateDataSourceConnection(string connectionString, string azureAccountToken, string ownerUri);
ReliableDataSourceConnection CreateDataSourceConnection(ConnectionDetails connectionDetails, string ownerUri);
}
}

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;
}
}
}

4
src/Microsoft.Kusto.ServiceLayer/LanguageServices/ConnectedBindingQueue.cs
View File

@@ -96,9 +96,7 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
try
{
bindingContext.BindingLock.Reset();
string connectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails);
bindingContext.DataSource = _dataSourceFactory.Create(DataSourceType.Kusto, connectionString, connInfo.ConnectionDetails.AzureAccountToken, connInfo.OwnerUri);
bindingContext.DataSource = _dataSourceFactory.Create(DataSourceType.Kusto, connInfo.ConnectionDetails, connInfo.OwnerUri);
bindingContext.BindingTimeout = DefaultBindingTimeout;
bindingContext.IsConnected = true;
}

4
src/Microsoft.Kusto.ServiceLayer/Scripting/Contracts/ScriptingRequest.cs
View File

@@ -24,9 +24,9 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting.Contracts
public string ScriptDestination { get; set; }
/// <summary>
/// Gets or sets connection string of the target database the scripting operation will run against.
/// Gets or sets the target database the scripting operation will run against.
/// </summary>
public string ConnectionString { get; set; }
public string DatabaseName { get; set; }
/// <summary>
/// Gets or sets a list of scripting objects to script.

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

@@ -5,7 +5,6 @@
using System;
using System.Collections.Generic;
using System.Data.SqlClient;
using Microsoft.Kusto.ServiceLayer.Scripting.Contracts;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Microsoft.SqlTools.Utility;
@@ -14,7 +13,6 @@ using Microsoft.SqlServer.Management.Smo;
using Microsoft.SqlServer.Management.SqlScriptPublish;
using Microsoft.SqlServer.Management.Sdk.Sfc;
using System.Diagnostics;
using Kusto.Data;
namespace Microsoft.Kusto.ServiceLayer.Scripting
{
@@ -158,7 +156,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
IEnumerable<ScriptingObject> selectedObjects = new List<ScriptingObject>(this.Parameters.ScriptingObjects);
_serverName = dataSource.ClusterName;
_databaseName = new KustoConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
_databaseName = Parameters.DatabaseName;
UrnCollection urnCollection = new UrnCollection();
foreach (var scriptingObject in selectedObjects)
{

4
src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingListObjectsOperation.cs
View File

@@ -9,7 +9,7 @@ using Microsoft.SqlServer.Management.SqlScriptPublish;
using Microsoft.Kusto.ServiceLayer.Scripting.Contracts;
using Microsoft.SqlTools.Utility;
using System.Diagnostics;
using Kusto.Data;
using Microsoft.Kusto.ServiceLayer.DataSource;
namespace Microsoft.Kusto.ServiceLayer.Scripting
{
@@ -103,7 +103,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
{
try
{
var builder = new KustoConnectionStringBuilder(this.Parameters.ConnectionString);
var builder = DataSourceFactory.CreateConnectionStringBuilder(DataSourceType.Kusto, Parameters.ConnectionString);
}
catch (Exception e)
{

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

@@ -1,296 +0,0 @@
//
// Copyright (c) Microsoft. All rights reserved.
// Licensed under the MIT license. See LICENSE file in the project root for full license information.
//
using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using Kusto.Data;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Microsoft.SqlServer.Management.SqlScriptPublish;
using Microsoft.Kusto.ServiceLayer.Scripting.Contracts;
using Microsoft.SqlTools.Utility;
namespace Microsoft.Kusto.ServiceLayer.Scripting
{
/// <summary>
/// Class to represent an in-progress script operation.
/// </summary>
public sealed class ScriptingScriptOperation : SmoScriptingOperation
{
private int scriptedObjectCount = 0;
private int totalScriptedObjectCount = 0;
private int eventSequenceNumber = 1;
public ScriptingScriptOperation(ScriptingParams parameters, IDataSource dataSource) : base(parameters, dataSource)
{
}
public override void Execute()
{
SqlScriptPublishModel publishModel = null;
try
{
this.CancellationToken.ThrowIfCancellationRequested();
this.ValidateScriptDatabaseParams();
publishModel = BuildPublishModel();
publishModel.ScriptItemsCollected += this.OnPublishModelScriptItemsCollected;
publishModel.ScriptProgress += this.OnPublishModelScriptProgress;
publishModel.ScriptError += this.OnPublishModelScriptError;
publishModel.AllowSystemObjects = true;
ScriptDestination destination = !string.IsNullOrWhiteSpace(this.Parameters.ScriptDestination)
? (ScriptDestination)Enum.Parse(typeof(ScriptDestination), this.Parameters.ScriptDestination)
: ScriptDestination.ToSingleFile;
// SMO is currently hardcoded to produce UTF-8 encoding when running on dotnet core.
ScriptOutputOptions outputOptions = new ScriptOutputOptions
{
SaveFileMode = ScriptFileMode.Overwrite,
SaveFileName = this.Parameters.FilePath,
ScriptDestination = destination,
};
this.CancellationToken.ThrowIfCancellationRequested();
publishModel.GenerateScript(outputOptions);
this.CancellationToken.ThrowIfCancellationRequested();
Logger.Write(
TraceEventType.Verbose,
string.Format(
"Sending script complete notification event for operation {0}, sequence number {1} with total count {2} and scripted count {3}",
this.OperationId,
this.eventSequenceNumber,
this.totalScriptedObjectCount,
this.scriptedObjectCount));
ScriptText = publishModel.RawScript;
this.SendCompletionNotificationEvent(new ScriptingCompleteParams
{
Success = true,
});
}
catch (Exception e)
{
if (e.IsOperationCanceledException())
{
Logger.Write(TraceEventType.Information, string.Format("Scripting operation {0} was canceled", this.OperationId));
this.SendCompletionNotificationEvent(new ScriptingCompleteParams
{
Canceled = true,
});
}
else
{
Logger.Write(TraceEventType.Error, string.Format("Scripting operation {0} failed with exception {1}", this.OperationId, e));
this.SendCompletionNotificationEvent(new ScriptingCompleteParams
{
HasError = true,
ErrorMessage = e.Message,
ErrorDetails = e.ToString(),
});
}
}
finally
{
if (publishModel != null)
{
publishModel.ScriptItemsCollected -= this.OnPublishModelScriptItemsCollected;
publishModel.ScriptProgress -= this.OnPublishModelScriptProgress;
publishModel.ScriptError -= this.OnPublishModelScriptError;
}
}
}
protected override void SendCompletionNotificationEvent(ScriptingCompleteParams parameters)
{
base.SendCompletionNotificationEvent(parameters);
}
protected override void SendPlanNotificationEvent(ScriptingPlanNotificationParams parameters)
{
base.SendPlanNotificationEvent(parameters);
}
protected override void SendProgressNotificationEvent(ScriptingProgressNotificationParams parameters)
{
base.SendProgressNotificationEvent(parameters);
}
protected override void SetCommonEventProperties(ScriptingEventParams parameters)
{
base.SetCommonEventProperties(parameters);
parameters.SequenceNumber = this.eventSequenceNumber;
this.eventSequenceNumber += 1;
}
private SqlScriptPublishModel BuildPublishModel()
{
SqlScriptPublishModel publishModel = new SqlScriptPublishModel(this.Parameters.ConnectionString);
// See if any filtering criteria was specified. If not, we're scripting the entire database. Otherwise, the filtering
// criteria should include the target objects to script.
//
bool hasObjectsSpecified = this.Parameters.ScriptingObjects != null && this.Parameters.ScriptingObjects.Any();
bool hasCriteriaSpecified =
(this.Parameters.IncludeObjectCriteria != null && this.Parameters.IncludeObjectCriteria.Any()) ||
(this.Parameters.ExcludeObjectCriteria != null && this.Parameters.ExcludeObjectCriteria.Any()) ||
(this.Parameters.IncludeSchemas != null && this.Parameters.IncludeSchemas.Any()) ||
(this.Parameters.ExcludeSchemas != null && this.Parameters.ExcludeSchemas.Any()) ||
(this.Parameters.IncludeTypes != null && this.Parameters.IncludeTypes.Any()) ||
(this.Parameters.ExcludeTypes != null && this.Parameters.ExcludeTypes.Any());
bool scriptAllObjects = !hasObjectsSpecified && !hasCriteriaSpecified;
// In the getter for SqlScriptPublishModel.AdvancedOptions, there is some strange logic which will
// cause the SqlScriptPublishModel.AdvancedOptions to get reset and lose all values based the ordering
// of when SqlScriptPublishModel.ScriptAllObjects is set.
//
publishModel.ScriptAllObjects = scriptAllObjects;
if (scriptAllObjects)
{
// Due to the getter logic within publishModel.AdvancedOptions, we explicitly populate the options
// after we determine what objects we are scripting.
//
PopulateAdvancedScriptOptions(this.Parameters.ScriptOptions, publishModel.AdvancedOptions);
return publishModel;
}
IEnumerable<ScriptingObject> selectedObjects = new List<ScriptingObject>();
if (hasCriteriaSpecified)
{
// This is an expensive remote call to load all objects from the database.
//
List<ScriptingObject> allObjects = publishModel.GetDatabaseObjects();
selectedObjects = ScriptingObjectMatcher.Match(
this.Parameters.IncludeObjectCriteria,
this.Parameters.ExcludeObjectCriteria,
this.Parameters.IncludeSchemas,
this.Parameters.ExcludeSchemas,
this.Parameters.IncludeTypes,
this.Parameters.ExcludeTypes,
allObjects);
}
if (hasObjectsSpecified)
{
selectedObjects = selectedObjects.Union(this.Parameters.ScriptingObjects);
}
// Populating advanced options after we select our objects in question, otherwise we lose all
// advanced options. After this call to PopulateAdvancedScriptOptions, DO NOT reference the
// publishModel.AdvancedOptions getter as it will reset the options in the model.
//
PopulateAdvancedScriptOptions(this.Parameters.ScriptOptions, publishModel.AdvancedOptions);
Logger.Write(
TraceEventType.Information,
string.Format(
"Scripting object count {0}, objects: {1}",
selectedObjects.Count(),
string.Join(", ", selectedObjects)));
string server = GetServerNameFromLiveInstance();
string database = new KustoConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
foreach (ScriptingObject scriptingObject in selectedObjects)
{
publishModel.SelectedObjects.Add(scriptingObject.ToUrn(server, database));
}
return publishModel;
}
private void OnPublishModelScriptError(object sender, ScriptEventArgs e)
{
this.CancellationToken.ThrowIfCancellationRequested();
Logger.Write(
TraceEventType.Verbose,
string.Format(
"Sending scripting error progress event, Urn={0}, OperationId={1}, Sequence={2}, Completed={3}, Error={4}",
e.Urn,
this.OperationId,
this.eventSequenceNumber,
e.Completed,
e?.Error?.ToString() ?? "null"));
// Keep scripting...it's a best effort operation.
e.ContinueScripting = true;
this.SendProgressNotificationEvent(new ScriptingProgressNotificationParams
{
ScriptingObject = e.Urn?.ToScriptingObject(),
Status = e.GetStatus(),
CompletedCount = this.scriptedObjectCount,
TotalCount = this.totalScriptedObjectCount,
ErrorMessage = e?.Error?.Message,
ErrorDetails = e?.Error?.ToString(),
});
}
private void OnPublishModelScriptItemsCollected(object sender, ScriptItemsArgs e)
{
this.CancellationToken.ThrowIfCancellationRequested();
List<ScriptingObject> scriptingObjects = e.Urns.Select(urn => urn.ToScriptingObject()).ToList();
this.totalScriptedObjectCount = scriptingObjects.Count;
Logger.Write(
TraceEventType.Verbose,
string.Format(
"Sending scripting plan notification event OperationId={0}, Sequence={1}, Count={2}, Objects: {3}",
this.OperationId,
this.eventSequenceNumber,
this.totalScriptedObjectCount,
string.Join(", ", e.Urns)));
this.SendPlanNotificationEvent(new ScriptingPlanNotificationParams
{
ScriptingObjects = scriptingObjects,
Count = scriptingObjects.Count,
});
}
private void OnPublishModelScriptProgress(object sender, ScriptEventArgs e)
{
this.CancellationToken.ThrowIfCancellationRequested();
if (e.Completed)
{
this.scriptedObjectCount += 1;
}
Logger.Write(
TraceEventType.Verbose,
string.Format(
"Sending progress event, Urn={0}, OperationId={1}, Sequence={2}, Status={3}, Error={4}",
e.Urn,
this.OperationId,
this.eventSequenceNumber,
e.GetStatus(),
e?.Error?.ToString() ?? "null"));
this.SendProgressNotificationEvent(new ScriptingProgressNotificationParams
{
ScriptingObject = e.Urn.ToScriptingObject(),
Status = e.GetStatus(),
CompletedCount = this.scriptedObjectCount,
TotalCount = this.totalScriptedObjectCount,
ErrorMessage = e?.Error?.Message,
ErrorDetails = e?.Error?.ToString(),
});
}
}
}

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

@@ -92,13 +92,13 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
// 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
if (parameters.ConnectionString == null)
if (parameters.DatabaseName == null)
{
ConnectionInfo connInfo;
_connectionService.TryFindConnection(parameters.OwnerUri, out connInfo);
if (connInfo != null)
{
parameters.ConnectionString = ConnectionService.BuildConnectionString(connInfo.ConnectionDetails);
parameters.DatabaseName = connInfo.ConnectionDetails.DatabaseName;
}
else
{
@@ -106,18 +106,15 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
}
}
SmoScriptingOperation operation;
var datasource = _connectionService.GetOrOpenConnection(parameters.OwnerUri, ConnectionType.Default)
.Result.GetUnderlyingConnection();
if (!ShouldCreateScriptAsOperation(parameters))
{
operation = new ScriptingScriptOperation(parameters, datasource);
}
else
{
operation = new ScriptAsScriptingOperation(parameters, _scripter, datasource);
throw new InvalidOperationException("Unable to create script.");
}
var datasource = _connectionService.GetOrOpenConnection(parameters.OwnerUri, ConnectionType.Default)
.Result.GetUnderlyingConnection();
SmoScriptingOperation operation = new ScriptAsScriptingOperation(parameters, _scripter, datasource);
operation.PlanNotification += (sender, e) => requestContext.SendEvent(ScriptingPlanNotificationEvent.Type, e).Wait();
operation.ProgressNotification += (sender, e) => requestContext.SendEvent(ScriptingProgressNotificationEvent.Type, e).Wait();
operation.CompleteNotification += (sender, e) => this.SendScriptingCompleteEvent(requestContext, ScriptingCompleteEvent.Type, e, operation, parameters.ScriptDestination);
@@ -136,17 +133,11 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
// Scripting as operation should be used to script one object.
// Scripting data and scripting to file is not supported by scripting as operation
// To script Select, alter and execute use scripting as operation. The other operation doesn't support those types
if( (parameters.ScriptingObjects != null && parameters.ScriptingObjects.Count == 1 && parameters.ScriptOptions != null
&& parameters.ScriptOptions.TypeOfDataToScript == "SchemaOnly" && parameters.ScriptDestination == "ToEditor") ||
parameters.Operation == ScriptingOperationType.Select || parameters.Operation == ScriptingOperationType.Execute ||
parameters.Operation == ScriptingOperationType.Alter)
{
return true;
}
else
{
return false;
}
return parameters.ScriptingObjects != null && parameters.ScriptingObjects.Count == 1 && parameters.ScriptOptions != null
&& parameters.ScriptOptions.TypeOfDataToScript == "SchemaOnly" && parameters.ScriptDestination == "ToEditor"
|| parameters.Operation == ScriptingOperationType.Select
|| parameters.Operation == ScriptingOperationType.Execute
|| parameters.Operation == ScriptingOperationType.Alter;
}
/// <summary>
@@ -227,7 +218,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
{
disposed = true;
foreach (ScriptingScriptOperation operation in this.ActiveOperations.Values)
foreach (var operation in this.ActiveOperations.Values)
{
operation.Dispose();
}

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

@@ -10,7 +10,6 @@ using System;
using System.Diagnostics;
using System.IO;
using System.Reflection;
using Kusto.Data;
using static Microsoft.SqlServer.Management.SqlScriptPublish.SqlScriptOptions;
namespace Microsoft.Kusto.ServiceLayer.Scripting
@@ -80,19 +79,15 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
protected void ValidateScriptDatabaseParams()
{
try
if (string.IsNullOrWhiteSpace(Parameters.DatabaseName))
{
var builder = new KustoConnectionStringBuilder(this.Parameters.ConnectionString);
throw new ArgumentException(SR.ScriptingParams_ConnectionString_Property_Invalid);
}
catch (Exception e)
{
throw new ArgumentException(SR.ScriptingParams_ConnectionString_Property_Invalid, e);
}
if (this.Parameters.FilePath == null && this.Parameters.ScriptDestination != "ToEditor")
if (Parameters.FilePath == null && this.Parameters.ScriptDestination != "ToEditor")
{
throw new ArgumentException(SR.ScriptingParams_FilePath_Property_Invalid);
}
else if (this.Parameters.FilePath != null && this.Parameters.ScriptDestination != "ToEditor")
if (Parameters.FilePath != null && this.Parameters.ScriptDestination != "ToEditor")
{
if (!Directory.Exists(Path.GetDirectoryName(this.Parameters.FilePath)))
{

12
test/Microsoft.Kusto.ServiceLayer.UnitTests/Connection/ConnectionInfoTests.cs
View File

@@ -35,8 +35,8 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
var connectionInfo = new ConnectionInfo(connectionFactoryMock.Object, "", new ConnectionDetails());
var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
var reliableDataSource = new ReliableDataSourceConnection("", RetryPolicyFactory.NoRetryPolicy,
RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object, "");
var reliableDataSource = new ReliableDataSourceConnection(new ConnectionDetails(), RetryPolicyFactory.NoRetryPolicy,
RetryPolicyFactory.NoRetryPolicy, dataSourceFactoryMock.Object, "");
connectionInfo.AddConnection("ConnectionType", reliableDataSource);
connectionInfo.TryGetConnection("ConnectionType", out var connection);
@@ -59,8 +59,8 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
var connectionInfo = new ConnectionInfo(connectionFactoryMock.Object, "", new ConnectionDetails());
var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
var reliableDataSource = new ReliableDataSourceConnection("", RetryPolicyFactory.NoRetryPolicy,
RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object, "");
var reliableDataSource = new ReliableDataSourceConnection(new ConnectionDetails(), RetryPolicyFactory.NoRetryPolicy,
RetryPolicyFactory.NoRetryPolicy, dataSourceFactoryMock.Object, "");
connectionInfo.AddConnection("ConnectionType", reliableDataSource);
connectionInfo.RemoveConnection("ConnectionType");
@@ -76,8 +76,8 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
var connectionInfo = new ConnectionInfo(connectionFactoryMock.Object, "", new ConnectionDetails());
var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
var reliableDataSource = new ReliableDataSourceConnection("", RetryPolicyFactory.NoRetryPolicy,
RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object, "");
var reliableDataSource = new ReliableDataSourceConnection(new ConnectionDetails(), RetryPolicyFactory.NoRetryPolicy,
RetryPolicyFactory.NoRetryPolicy, dataSourceFactoryMock.Object, "");
connectionInfo.AddConnection("ConnectionType", reliableDataSource);
connectionInfo.RemoveAllConnections();

3
test/Microsoft.Kusto.ServiceLayer.UnitTests/Connection/DataSourceConnectionFactoryTests.cs
View File

@@ -1,4 +1,5 @@
using Microsoft.Kusto.ServiceLayer.Connection;
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Moq;
using NUnit.Framework;
@@ -12,7 +13,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
{
var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
var connectionFactory = new DataSourceConnectionFactory(dataSourceFactoryMock.Object);
var connection = connectionFactory.CreateDataSourceConnection("", "", "");
var connection = connectionFactory.CreateDataSourceConnection(new ConnectionDetails(), "");
Assert.IsNotNull(connection);
}

10
test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/DataSourceFactoryTests.cs
View File

@@ -1,5 +1,6 @@
using System;
using System.Collections.Generic;
using Microsoft.Kusto.ServiceLayer.Connection.Contracts;
using Microsoft.Kusto.ServiceLayer.DataSource;
using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
using Microsoft.Kusto.ServiceLayer.LanguageServices.Completion;
@@ -10,16 +11,19 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource
{
public class DataSourceFactoryTests
{
[TestCase(typeof(ArgumentNullException), "", "AzureAccountToken")]
[TestCase(typeof(ArgumentNullException), "ConnectionString", "")]
[TestCase(typeof(ArgumentException), "ConnectionString", "AzureAccountToken")]
public void Create_Throws_Exceptions_For_InvalidParams(Type exceptionType,
string connectionString,
string azureAccountToken)
{
var dataSourceFactory = new DataSourceFactory();
var connectionDetails = new ConnectionDetails
{
ConnectionString = connectionString,
AccountToken = azureAccountToken
};
Assert.Throws(exceptionType,
() => dataSourceFactory.Create(DataSourceType.None, connectionString, azureAccountToken, ""));
() => dataSourceFactory.Create(DataSourceType.None, connectionDetails, ""));
}
[Test]

4
test/Microsoft.Kusto.ServiceLayer.UnitTests/LanguageServices/ConnectedBindingQueueTests.cs
View File

@@ -92,7 +92,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.LanguageServices
{
var connectionDetails = new ConnectionDetails
{
AzureAccountToken = "AzureAccountToken"
AccountToken = "AzureAccountToken"
};
var connectionFactory = new Mock<IDataSourceConnectionFactory>();
var connectionInfo = new ConnectionInfo(connectionFactory.Object, "ownerUri", connectionDetails);
@@ -100,7 +100,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.LanguageServices
var dataSourceFactory = new Mock<IDataSourceFactory>();
var dataSourceMock = new Mock<IDataSource>();
dataSourceFactory
.Setup(x => x.Create(It.IsAny<DataSourceType>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>()))
.Setup(x => x.Create(It.IsAny<DataSourceType>(), It.IsAny<ConnectionDetails>(), It.IsAny<string>()))
.Returns(dataSourceMock.Object);
var connectedBindingQueue =