4094 Kusto AAD Authentication Refactor (#1115)

* 4094 Upgraded Kusto.Data and Kusto.Language to 9.0.3. Refactored KustoServiceLayer to use KustoConnectionStringBuilder instead of Sql.

* 4094 Refactored KustoClient>Initialize to take connectionString as optional. Removed unused properties from ConnectionService>ParseConnectionString

* 4094 Updated Kusto.Data and Kusto.Language nuget packages to 9.0.4. Removed IsDedicatedAdminConnection function in ConnectionService.

* 4094 Fixed unit tests

Packages.props

@@ -20,8 +20,8 @@
 		<PackageReference Update="Microsoft.Data.SqlClient" Version="2.0.0"/>
 		<PackageReference Update="Microsoft.SqlServer.SqlManagementObjects" Version="161.44091.28" />
 		<PackageReference Update="Microsoft.SqlServer.DACFx" Version="150.4946.1-preview" GeneratePathProperty="true" />
-		<PackageReference Update="Microsoft.Azure.Kusto.Data" Version="8.0.2" />
-		<PackageReference Update="Microsoft.Azure.Kusto.Language" Version="8.1.2"/>
+		<PackageReference Update="Microsoft.Azure.Kusto.Data" Version="9.0.4" />
+		<PackageReference Update="Microsoft.Azure.Kusto.Language" Version="9.0.4"/>
 		<PackageReference Update="Microsoft.SqlServer.Assessment"  Version="1.0.280-preview" />
 		<PackageReference Update="Microsoft.SqlServer.Migration.Assessment"  Version="1.0.20201001.204" />
 

src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs

@@ -19,6 +19,7 @@ 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;
@@ -30,7 +31,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
     /// </summary>
     public class ConnectionService
     {
-        private const string AdminConnectionPrefix = "ADMIN:";
         private const string PasswordPlaceholder = "******";
 
         /// <summary>
@@ -389,7 +389,7 @@ 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 SqlConnectionStringBuilder(connection.ConnectionString);
+                    var connectionString = new KustoConnectionStringBuilder(connection.ConnectionString);
                     response.ConnectionSummary = new ConnectionSummary
                     {
                         ServerName = connectionString.DataSource,
@@ -548,19 +548,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             {
                 throw new InvalidOperationException(SR.ConnectionServiceDbErrorDefaultNotConnected(ownerUri));
             }
-
-            if (IsDedicatedAdminConnection(connectionInfo.ConnectionDetails))
+            
+            // Try to get the ReliableDataSourceClient and create if it doesn't already exist
+            if (!connectionInfo.TryGetConnection(connectionType, out connection) && ConnectionType.Default != connectionType)
             {
-                // Since this is a dedicated connection only 1 is allowed at any time. Return the default connection for use in the requested action
-                connection = defaultConnection;
-            }
-            else
-            {
-                // Try to get the ReliableDataSourceClient and create if it doesn't already exist
-                if (!connectionInfo.TryGetConnection(connectionType, out connection) && ConnectionType.Default != connectionType)
-                {
-                    connection = await TryOpenConnectionForConnectionType(ownerUri, connectionType, alwaysPersistSecurity, connectionInfo);
-                }
+                connection = await TryOpenConnectionForConnectionType(ownerUri, connectionType, alwaysPersistSecurity, connectionInfo);
             }
 
             return connection;
@@ -963,18 +955,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             }
         }
 
-        /// <summary>
-        /// Checks if a ConnectionDetails object represents a DAC connection
-        /// </summary>
-        /// <param name="connectionDetails"></param>
-        public static bool IsDedicatedAdminConnection(ConnectionDetails connectionDetails)
-        {
-            Validate.IsNotNull(nameof(connectionDetails), connectionDetails);
-            SqlConnectionStringBuilder builder = CreateConnectionStringBuilder(connectionDetails);
-            string serverName = builder.DataSource;
-            return serverName != null && serverName.StartsWith(AdminConnectionPrefix, StringComparison.OrdinalIgnoreCase);
-        }
-
         /// <summary>
         /// Build a connection string from a connection details instance
         /// </summary>
@@ -988,152 +968,13 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         /// Build a connection string builder a connection details instance
         /// </summary>
         /// <param name="connectionDetails"></param>
-        private static SqlConnectionStringBuilder CreateConnectionStringBuilder(ConnectionDetails connectionDetails)
+        private static KustoConnectionStringBuilder CreateConnectionStringBuilder(ConnectionDetails connectionDetails)
         {
-            SqlConnectionStringBuilder connectionBuilder;
+            var stringBuilder = string.IsNullOrWhiteSpace(connectionDetails.ConnectionString)
+                ? new KustoConnectionStringBuilder(connectionDetails.ServerName, connectionDetails.DatabaseName)
+                : new KustoConnectionStringBuilder(connectionDetails.ConnectionString);
 
-            // If connectionDetails has a connection string already, use it to initialize the connection builder, then override any provided options.
-            // Otherwise use the server name, username, and password from the connection details.
-            if (!string.IsNullOrEmpty(connectionDetails.ConnectionString))
-            {
-                connectionBuilder = new SqlConnectionStringBuilder(connectionDetails.ConnectionString);
-            }
-            else
-            {
-                // add alternate port to data source property if provided
-                string dataSource = !connectionDetails.Port.HasValue
-                    ? connectionDetails.ServerName
-                    : string.Format("{0},{1}", connectionDetails.ServerName, connectionDetails.Port.Value);
-
-                connectionBuilder = new SqlConnectionStringBuilder
-                {
-                    ["Data Source"] = dataSource,
-                    ["User Id"] = connectionDetails.UserName,
-                    ["Password"] = connectionDetails.Password
-                };
-            }
-
-            // Check for any optional parameters
-            if (!string.IsNullOrEmpty(connectionDetails.DatabaseName))
-            {
-                connectionBuilder["Initial Catalog"] = connectionDetails.DatabaseName;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.AuthenticationType))
-            {
-                switch (connectionDetails.AuthenticationType)
-                {
-                    case "Integrated":
-                        connectionBuilder.IntegratedSecurity = true;
-                        break;
-                    case "SqlLogin":
-                        break;
-                    case "AzureMFA":
-                        connectionBuilder.UserID = "";
-                        connectionBuilder.Password = "";
-                        break;
-                    default:
-                        throw new ArgumentException(SR.ConnectionServiceConnStringInvalidAuthType(connectionDetails.AuthenticationType));
-                }
-            }
-            if (connectionDetails.Encrypt.HasValue)
-            {
-                connectionBuilder.Encrypt = connectionDetails.Encrypt.Value;
-            }
-            if (connectionDetails.TrustServerCertificate.HasValue)
-            {
-                connectionBuilder.TrustServerCertificate = connectionDetails.TrustServerCertificate.Value;
-            }
-            if (connectionDetails.PersistSecurityInfo.HasValue)
-            {
-                connectionBuilder.PersistSecurityInfo = connectionDetails.PersistSecurityInfo.Value;
-            }
-            if (connectionDetails.ConnectTimeout.HasValue)
-            {
-                connectionBuilder.ConnectTimeout = connectionDetails.ConnectTimeout.Value;
-            }
-            if (connectionDetails.ConnectRetryCount.HasValue)
-            {
-                connectionBuilder.ConnectRetryCount = connectionDetails.ConnectRetryCount.Value;
-            }
-            if (connectionDetails.ConnectRetryInterval.HasValue)
-            {
-                connectionBuilder.ConnectRetryInterval = connectionDetails.ConnectRetryInterval.Value;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.ApplicationName))
-            {
-                connectionBuilder.ApplicationName = connectionDetails.ApplicationName;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.WorkstationId))
-            {
-                connectionBuilder.WorkstationID = connectionDetails.WorkstationId;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.ApplicationIntent))
-            {
-                ApplicationIntent intent;
-                switch (connectionDetails.ApplicationIntent)
-                {
-                    case "ReadOnly":
-                        intent = ApplicationIntent.ReadOnly;
-                        break;
-                    case "ReadWrite":
-                        intent = ApplicationIntent.ReadWrite;
-                        break;
-                    default:
-                        throw new ArgumentException(SR.ConnectionServiceConnStringInvalidIntent(connectionDetails.ApplicationIntent));
-                }
-                connectionBuilder.ApplicationIntent = intent;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.CurrentLanguage))
-            {
-                connectionBuilder.CurrentLanguage = connectionDetails.CurrentLanguage;
-            }
-            if (connectionDetails.Pooling.HasValue)
-            {
-                connectionBuilder.Pooling = connectionDetails.Pooling.Value;
-            }
-            if (connectionDetails.MaxPoolSize.HasValue)
-            {
-                connectionBuilder.MaxPoolSize = connectionDetails.MaxPoolSize.Value;
-            }
-            if (connectionDetails.MinPoolSize.HasValue)
-            {
-                connectionBuilder.MinPoolSize = connectionDetails.MinPoolSize.Value;
-            }
-            if (connectionDetails.LoadBalanceTimeout.HasValue)
-            {
-                connectionBuilder.LoadBalanceTimeout = connectionDetails.LoadBalanceTimeout.Value;
-            }
-            if (connectionDetails.Replication.HasValue)
-            {
-                connectionBuilder.Replication = connectionDetails.Replication.Value;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.AttachDbFilename))
-            {
-                connectionBuilder.AttachDBFilename = connectionDetails.AttachDbFilename;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.FailoverPartner))
-            {
-                connectionBuilder.FailoverPartner = connectionDetails.FailoverPartner;
-            }
-            if (connectionDetails.MultiSubnetFailover.HasValue)
-            {
-                connectionBuilder.MultiSubnetFailover = connectionDetails.MultiSubnetFailover.Value;
-            }
-            if (connectionDetails.MultipleActiveResultSets.HasValue)
-            {
-                connectionBuilder.MultipleActiveResultSets = connectionDetails.MultipleActiveResultSets.Value;
-            }
-            if (connectionDetails.PacketSize.HasValue)
-            {
-                connectionBuilder.PacketSize = connectionDetails.PacketSize.Value;
-            }
-            if (!string.IsNullOrEmpty(connectionDetails.TypeSystemVersion))
-            {
-                connectionBuilder.TypeSystemVersion = connectionDetails.TypeSystemVersion;
-            }
-            connectionBuilder.Pooling = false;
-
-            return connectionBuilder;
+            return stringBuilder.WithAadUserTokenAuthentication(connectionDetails.AzureAccountToken);
         }
 
         /// <summary>
@@ -1194,38 +1035,15 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
 
         public ConnectionDetails ParseConnectionString(string connectionString)
         {
-            SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder(connectionString);
-            ConnectionDetails details = new ConnectionDetails()
+            var builder = new KustoConnectionStringBuilder(connectionString);
+            return new ConnectionDetails
             {
-                ApplicationIntent = builder.ApplicationIntent.ToString(),
-                ApplicationName = builder.ApplicationName,
-                AttachDbFilename = builder.AttachDBFilename,
-                AuthenticationType = builder.IntegratedSecurity ? "Integrated" : "SqlLogin",
-                ConnectRetryCount = builder.ConnectRetryCount,
-                ConnectRetryInterval = builder.ConnectRetryInterval,
-                ConnectTimeout = builder.ConnectTimeout,
-                CurrentLanguage = builder.CurrentLanguage,
+                ApplicationName = builder.ApplicationNameForTracing,
+                AuthenticationType = "AzureMFA",
                 DatabaseName = builder.InitialCatalog,
-                Encrypt = builder.Encrypt,
-                FailoverPartner = builder.FailoverPartner,
-                LoadBalanceTimeout = builder.LoadBalanceTimeout,
-                MaxPoolSize = builder.MaxPoolSize,
-                MinPoolSize = builder.MinPoolSize,
-                MultipleActiveResultSets = builder.MultipleActiveResultSets,
-                MultiSubnetFailover = builder.MultiSubnetFailover,
-                PacketSize = builder.PacketSize,
-                Password = !builder.IntegratedSecurity ? builder.Password : string.Empty,
-                PersistSecurityInfo = builder.PersistSecurityInfo,
-                Pooling = builder.Pooling,
-                Replication = builder.Replication,
                 ServerName = builder.DataSource,
-                TrustServerCertificate = builder.TrustServerCertificate,
-                TypeSystemVersion = builder.TypeSystemVersion,
                 UserName = builder.UserID,
-                WorkstationId = builder.WorkstationID,
             };
-
-            return details;
         }
 
         /// <summary>

src/Microsoft.Kusto.ServiceLayer/DataSource/DataSourceIntellisense/KustoIntellisenseHelper.cs

@@ -117,8 +117,11 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense
             var query = "let fn = " + parameters + " { };";
             var code = KustoCode.ParseAndAnalyze(query);
             var let = code.Syntax.GetFirstDescendant<LetStatement>();
-            var variable = let.Name.ReferencedSymbol as VariableSymbol;
-            var function = variable.Type as FunctionSymbol;
+            
+             FunctionSymbol function = let.Name.ReferencedSymbol is VariableSymbol variable
+                ? variable.Type as FunctionSymbol
+                : let.Name.ReferencedSymbol as FunctionSymbol;
+
             return function.Signatures[0].Parameters;
         }
 

src/Microsoft.Kusto.ServiceLayer/DataSource/KustoClient.cs

@@ -14,7 +14,6 @@ using Kusto.Data.Exceptions;
 using Kusto.Data.Net.Client;
 using Kusto.Language;
 using Kusto.Language.Editor;
-using Microsoft.Data.SqlClient;
 using Microsoft.Kusto.ServiceLayer.Connection;
 using Microsoft.Kusto.ServiceLayer.DataSource.DataSourceIntellisense;
 using Microsoft.Kusto.ServiceLayer.Utility;
@@ -36,26 +35,21 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// </summary>
         public GlobalState SchemaState { get; private set; }
 
-        public string ClusterName { get; }
+        public string ClusterName { get; private set; }
         public string DatabaseName { get; private set; }
 
         public KustoClient(string connectionString, string azureAccountToken, string ownerUri)
         {
             _ownerUri = ownerUri;
-            ClusterName = GetClusterName(connectionString);
-            
-            var dbName = new SqlConnectionStringBuilder(connectionString).InitialCatalog;
-            SetDatabaseName(dbName);
-            
-            Initialize(ClusterName, DatabaseName, azureAccountToken);
+            Initialize(azureAccountToken, connectionString);
             SchemaState = LoadSchemaState();
         }
 
-        private void SetDatabaseName(string databaseName)
+        private string ParseDatabaseName(string databaseName)
         {
             var regex = new Regex(@"(?<=\().+?(?=\))");
             
-            DatabaseName = regex.IsMatch(databaseName)
+            return regex.IsMatch(databaseName)
                 ? regex.Match(databaseName).Value
                 : databaseName;
         }
@@ -85,9 +79,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 DatabaseName, ClusterName);
         }
 
-        private void Initialize(string clusterName, string databaseName, string azureAccountToken)
+        private void Initialize(string azureAccountToken, string connectionString = "")
         {
-            var stringBuilder = GetKustoConnectionStringBuilder(clusterName, databaseName, azureAccountToken, "", "");
+            var stringBuilder = GetKustoConnectionStringBuilder(azureAccountToken, connectionString);
             _kustoQueryProvider = KustoClientFactory.CreateCslQueryProvider(stringBuilder);
             _kustoAdminProvider = KustoClientFactory.CreateCslAdminProvider(stringBuilder);
         }
@@ -97,68 +91,26 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             string azureAccountToken = ConnectionService.Instance.RefreshAzureToken(_ownerUri);
             _kustoQueryProvider.Dispose();
             _kustoAdminProvider.Dispose();
-            Initialize(ClusterName, DatabaseName, azureAccountToken);
+            Initialize(azureAccountToken);
         }
 
-        /// <summary>
-        /// Extracts the cluster name from the connectionstring. The string looks like the following:
-        /// "Data Source=clustername.kusto.windows.net;User ID=;Password=;Pooling=False;Application Name=azdata-GeneralConnection"
-        /// <summary>
-        /// <param name="connectionString">A connection string coming over the Data management protocol</param>
-        private string GetClusterName(string connectionString)
+        private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(string userToken, string connectionString)
         {
-            var csb = new SqlConnectionStringBuilder(connectionString);
+            ValidationUtils.IsTrue<ArgumentException>(!string.IsNullOrWhiteSpace(userToken),
+                $"the Kusto authentication is not specified - either set {nameof(userToken)}");
 
-            // If there is no https:// prefix, add it
-            Uri uri;
-            if ((Uri.TryCreate(csb.DataSource, UriKind.Absolute, out uri) ||
-                 Uri.TryCreate("https://" + csb.DataSource, UriKind.Absolute, out uri)) &&
-                (uri.Scheme == Uri.UriSchemeHttp || uri.Scheme == Uri.UriSchemeHttps))
-            {
-                return uri.AbsoluteUri;
-            }
-
-            throw new ArgumentException("Expected a URL of the form clustername.kusto.windows.net");
-        }
-
-        private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(string clusterName, string databaseName,
-            string userToken, string applicationClientId, string applicationKey)
-        {
-            ValidationUtils.IsNotNull(clusterName, nameof(clusterName));
-            ValidationUtils.IsTrue<ArgumentException>(
-                !string.IsNullOrWhiteSpace(userToken)
-                || (!string.IsNullOrWhiteSpace(applicationClientId) && !string.IsNullOrWhiteSpace(applicationKey)),
-                $"the Kusto authentication is not specified - either set {nameof(userToken)}, or set {nameof(applicationClientId)} and {nameof(applicationKey)}");
-
-            var kcsb = new KustoConnectionStringBuilder
-            {
-                DataSource = clusterName,
-
-                // Perform federated auth based on the AAD user token, or based on the AAD application client id and key.
-                FederatedSecurity = true
-            };
-
-            if (!string.IsNullOrWhiteSpace(databaseName))
-            {
-                kcsb.InitialCatalog = databaseName;
-            }
-
-            if (!string.IsNullOrWhiteSpace(userToken))
-            {
-                kcsb.UserToken = userToken;
-            }
-
-            if (!string.IsNullOrWhiteSpace(applicationClientId))
-            {
-                kcsb.ApplicationClientId = applicationClientId;
-            }
-
-            if (!string.IsNullOrWhiteSpace(applicationKey))
-            {
-                kcsb.ApplicationKey = applicationKey;
-            }
-
-            return kcsb;
+            var stringBuilder = string.IsNullOrWhiteSpace(connectionString)
+                ? new KustoConnectionStringBuilder(ClusterName, DatabaseName)
+                : new KustoConnectionStringBuilder(connectionString);
+            
+            ClusterName = stringBuilder.DataSource;
+            var databaseName = ParseDatabaseName(stringBuilder.InitialCatalog);
+            DatabaseName = databaseName;
+            stringBuilder.InitialCatalog = databaseName;
+            
+            ValidationUtils.IsNotNull(ClusterName, nameof(ClusterName));
+            
+            return stringBuilder.WithAadUserTokenAuthentication(userToken);
         }
 
         private ClientRequestProperties GetClientRequestProperties(CancellationToken cancellationToken)
@@ -310,7 +262,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
         public void UpdateDatabase(string databaseName)
         {
-            SetDatabaseName(databaseName);
+            DatabaseName = ParseDatabaseName(databaseName);
             SchemaState = LoadSchemaState();
         }
 

src/Microsoft.Kusto.ServiceLayer/DataSource/Metadata/MetadataFactory.cs

@@ -65,12 +65,12 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Metadata
         {
             var databaseDetails = new List<DatabaseInfo>();
 
-            if (typeof(DatabaseMetadata) == clusterDBDetails.FirstOrDefault().GetType())
+            if (clusterDBDetails.FirstOrDefault() is DatabaseMetadata)
             {
                 foreach (var dbDetail in clusterDBDetails)
                 {
                     DatabaseInfo databaseInfo = new DatabaseInfo();
-                    Int64.TryParse(dbDetail.SizeInMB.ToString(), out long sum_OriginalSize);
+                    long.TryParse(dbDetail.SizeInMB, out long sum_OriginalSize);
                     databaseInfo.Options["name"] = dbDetail.Name;
                     databaseInfo.Options["sizeInMB"] = (sum_OriginalSize / (1024 * 1024)).ToString();
                     databaseDetails.Add(databaseInfo);

src/Microsoft.Kusto.ServiceLayer/LanguageServices/LanguageService.cs

@@ -684,11 +684,6 @@ namespace Microsoft.Kusto.ServiceLayer.LanguageServices
         {
             await Task.Run(() =>
             {
-                if (ConnectionService.IsDedicatedAdminConnection(connInfo.ConnectionDetails))
-                {
-                    // Intellisense cannot be run on these connections as only 1 SqlConnection can be opened on them at a time
-                    return;
-                }
                 ScriptParseInfo scriptInfo = GetScriptParseInfo(connInfo.OwnerUri, createIfNotExists: true);
                 if (Monitor.TryEnter(scriptInfo.BuildingMetadataLock, LanguageService.OnConnectionWaitTimeout))
                 {

src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptAsScriptingOperation.cs

@@ -14,6 +14,7 @@ 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
 {
@@ -157,7 +158,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
             IEnumerable<ScriptingObject> selectedObjects = new List<ScriptingObject>(this.Parameters.ScriptingObjects);
 
             _serverName = dataSource.ClusterName;
-            _databaseName = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
+            _databaseName = new KustoConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
             UrnCollection urnCollection = new UrnCollection();
             foreach (var scriptingObject in selectedObjects)
             {

src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingListObjectsOperation.cs

@@ -5,19 +5,11 @@
 
 using System;
 using System.Collections.Generic;
-using System.Data.SqlClient;
-using System.IO;
-using System.Linq;
-using System.Reflection;
-using System.Threading;
-using System.Threading.Tasks;
-using Microsoft.SqlServer.Management.Sdk.Sfc;
 using Microsoft.SqlServer.Management.SqlScriptPublish;
-using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.Kusto.ServiceLayer.Scripting.Contracts;
 using Microsoft.SqlTools.Utility;
-using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 using System.Diagnostics;
+using Kusto.Data;
 
 namespace Microsoft.Kusto.ServiceLayer.Scripting
 {
@@ -111,7 +103,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
         {
             try
             {
-                SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder(this.Parameters.ConnectionString);
+                var builder = new KustoConnectionStringBuilder(this.Parameters.ConnectionString);
             }
             catch (Exception e)
             {

src/Microsoft.Kusto.ServiceLayer/Scripting/ScriptingScriptOperation.cs

@@ -5,9 +5,9 @@
 
 using System;
 using System.Collections.Generic;
-using System.Data.SqlClient;
 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;
@@ -202,7 +202,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                     string.Join(", ", selectedObjects)));
 
             string server = GetServerNameFromLiveInstance();
-            string database = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
+            string database = new KustoConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
 
             foreach (ScriptingObject scriptingObject in selectedObjects)
             {

src/Microsoft.Kusto.ServiceLayer/Scripting/SmoScriptingOperation.cs

@@ -7,10 +7,10 @@ using Microsoft.Kusto.ServiceLayer.Scripting.Contracts;
 using Microsoft.Kusto.ServiceLayer.DataSource;
 using Microsoft.SqlTools.Utility;
 using System;
-using System.Data.SqlClient;
 using System.Diagnostics;
 using System.IO;
 using System.Reflection;
+using Kusto.Data;
 using static Microsoft.SqlServer.Management.SqlScriptPublish.SqlScriptOptions;
 
 namespace Microsoft.Kusto.ServiceLayer.Scripting
@@ -82,7 +82,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
         {
             try
             {
-                SqlConnectionStringBuilder builder = new SqlConnectionStringBuilder(this.Parameters.ConnectionString);
+                var builder = new KustoConnectionStringBuilder(this.Parameters.ConnectionString);
             }
             catch (Exception e)
             {

test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/DataSourceFactoryTests.cs

@@ -61,7 +61,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource
             
             var semanticMarkers = DataSourceFactory.GetDefaultSemanticMarkers(DataSourceType.Kusto, parseInfo, file, queryText);
             
-            Assert.AreNotEqual(0, semanticMarkers.Length);
+            Assert.AreEqual(0, semanticMarkers.Length);
         }
 
         [Test]

test/Microsoft.Kusto.ServiceLayer.UnitTests/DataSource/DataSourceIntellisense/KustoIntellisenseHelperTests.cs

@@ -65,7 +65,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource.DataSourceIntellisen
             var queryText = ".show databases";
             var completionItems = KustoIntellisenseHelper.GetDefaultDiagnostics(parseInfo, scriptFile, queryText);
             
-            Assert.AreEqual(6, completionItems.Length);
+            Assert.AreEqual(0, completionItems.Length);
         }
     }
 }

test/Microsoft.Kusto.ServiceLayer.UnitTests/LanguageServices/ConnectedBindingQueueTests.cs

@@ -86,11 +86,14 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.LanguageServices
             
             Assert.AreEqual("NULL_NULL_NULL_NULL", connectionKey);
         }
-
-        [TestCase(false)]
-        public void AddConnectionContext_Sets_BindingContext(bool needsMetadata)
+        
+        [Test]
+        public void AddConnectionContext_Sets_BindingContext()
         {
-            var connectionDetails = new ConnectionDetails();
+            var connectionDetails = new ConnectionDetails
+            {
+                AzureAccountToken = "AzureAccountToken" 
+            };
             var connectionFactory = new Mock<IDataSourceConnectionFactory>();
             var connectionInfo = new ConnectionInfo(connectionFactory.Object, "ownerUri", connectionDetails);
 
@@ -103,7 +106,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.LanguageServices
             var connectedBindingQueue =
                 new ConnectedBindingQueue(dataSourceFactory.Object);
             var connectionKey =
-                connectedBindingQueue.AddConnectionContext(connectionInfo, needsMetadata, "featureName");
+                connectedBindingQueue.AddConnectionContext(connectionInfo, false, "featureName");
             var bindingContext = connectedBindingQueue.GetOrCreateBindingContext(connectionKey);
             
             Assert.AreEqual(dataSourceMock.Object, bindingContext.DataSource);