3490 Kusto Connection Refresh Fix (#1085)

* 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

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

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

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

@@ -6,7 +6,6 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
-using System.Globalization;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -268,23 +267,22 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             return completeParams;
         }
 
-        internal void RefreshAzureToken(string ownerUri)
+        internal string RefreshAzureToken(string ownerUri)
         {
-            ConnectionInfo existingConnection = OwnerToConnectionMap[ownerUri];
+            TryFindConnection(ownerUri, out ConnectionInfo connection);
-            
+
             var requestMessage = new RequestSecurityTokenParams
             {
-                AccountId = existingConnection.ConnectionDetails.GetOptionValue("azureAccount", string.Empty),
+                AccountId = connection.ConnectionDetails.GetOptionValue("azureAccount", string.Empty),
-                Authority = existingConnection.ConnectionDetails.GetOptionValue("azureTenantId", string.Empty),
+                Authority = connection.ConnectionDetails.GetOptionValue("azureTenantId", string.Empty),
                 Provider = "Azure",
                 Resource = "SQL"
             };
 
             var response = Instance.ServiceHost.SendRequest(SecurityTokenRequest.Type, requestMessage, true).Result;
-            existingConnection.UpdateAzureToken(response.Token);
+            connection.UpdateAzureToken(response.Token);
-            
+
-            existingConnection.TryGetConnection(ConnectionType.Query, out var reliableDataSourceConnection);
+            return response.Token;
-            reliableDataSourceConnection.GetUnderlyingConnection().UpdateAzureToken(response.Token);
         }
 
         private void TryCloseConnectionTemporaryConnection(ConnectParams connectionParams, ConnectionInfo connectionInfo)
@@ -452,7 +450,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
                 string connectionString = BuildConnectionString(connectionInfo.ConnectionDetails);
 
                 // create a sql connection instance
-                connection = connectionInfo.Factory.CreateDataSourceConnection(connectionString, connectionInfo.ConnectionDetails.AzureAccountToken);
+                connection = connectionInfo.Factory.CreateDataSourceConnection(connectionString, connectionInfo.ConnectionDetails.AzureAccountToken, connectionInfo.OwnerUri);
                 connectionInfo.AddConnection(connectionParams.Type, connection);
 
                 // Add a cancellation token source so that the connection OpenAsync() can be cancelled
@@ -786,9 +784,10 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             {
                 throw new Exception(SR.ConnectionServiceListDbErrorNotConnected(owner));
             }
-            ConnectionDetails connectionDetails = info.ConnectionDetails.Clone();
 
-            IDataSource dataSource = OpenDataSourceConnection(info);
+            info.TryGetConnection(ConnectionType.Default, out ReliableDataSourceConnection connection);
+            IDataSource dataSource = connection.GetUnderlyingConnection();
+            
             DataSourceObjectMetadata objectMetadata = MetadataFactory.CreateClusterMetadata(info.ConnectionDetails.ServerName);
 
             ListDatabasesResponse response = new ListDatabasesResponse();
@@ -813,7 +812,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             ServiceHost = serviceHost;
             _dataSourceConnectionFactory = dataSourceConnectionFactory;
             _dataSourceFactory = dataSourceFactory;
-            
             connectedQueues.AddOrUpdate("Default", connectedBindingQueue, (key, old) => connectedBindingQueue);
             LockedDatabaseManager.ConnectionService = this;
 
@@ -1262,7 +1260,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
                                 string connectionString = BuildConnectionString(info.ConnectionDetails);
 
                                 // create a sql connection instance
-                                ReliableDataSourceConnection connection = info.Factory.CreateDataSourceConnection(connectionString, info.ConnectionDetails.AzureAccountToken);
+                                ReliableDataSourceConnection connection = info.Factory.CreateDataSourceConnection(connectionString, info.ConnectionDetails.AzureAccountToken, ownerUri);
                                 connection.Open();
                                 info.AddConnection(key, connection);
                             }
@@ -1358,35 +1356,6 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             }
         }
 
-        /// <summary>
-        /// Create and open a new SqlConnection from a ConnectionInfo object
-        /// Note: we need to audit all uses of this method to determine why we're
-        /// bypassing normal ConnectionService connection management
-        /// </summary>
-        /// <param name="connInfo">The connection info to connect with</param>
-        /// <param name="featureName">A plaintext string that will be included in the application name for the connection</param>
-        /// <returns>A SqlConnection created with the given connection info</returns>
-        private IDataSource OpenDataSourceConnection(ConnectionInfo connInfo, string featureName = null)
-        {
-            try
-            {
-                // generate connection string
-                string connectionString = BuildConnectionString(connInfo.ConnectionDetails);
-
-                // TODOKusto: Pass in type of DataSource needed to make this generic. Hard coded to Kusto right now.
-                return _dataSourceFactory.Create(DataSourceType.Kusto, connectionString, connInfo.ConnectionDetails.AzureAccountToken);
-            }
-            catch (Exception ex)
-            {
-                string error = string.Format(CultureInfo.InvariantCulture,
-                    "Failed opening a DataSource of type {0}: error:{1} inner:{2} stacktrace:{3}",
-                    DataSourceType.Kusto, ex.Message, ex.InnerException != null ? ex.InnerException.Message : string.Empty, ex.StackTrace);
-                Logger.Write(TraceEventType.Error, error);
-            }
-
-            return null;
-        }
-
         public static void EnsureConnectionIsOpen(ReliableDataSourceConnection conn, bool forceReopen = false)
         {
             // verify that the connection is open

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

@@ -28,11 +28,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         /// <summary>
         /// Creates a new SqlConnection object
         /// </summary>
-        public ReliableDataSourceConnection CreateDataSourceConnection(string connectionString, string azureAccountToken)
+        public ReliableDataSourceConnection CreateDataSourceConnection(string connectionString, string azureAccountToken, string ownerUri)
         {
             RetryPolicy connectionRetryPolicy = RetryPolicyFactory.CreateDefaultConnectionRetryPolicy();
             RetryPolicy commandRetryPolicy = RetryPolicyFactory.CreateDefaultConnectionRetryPolicy();
-            return new ReliableDataSourceConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken, _dataSourceFactory);
+            return new ReliableDataSourceConnection(connectionString, connectionRetryPolicy, commandRetryPolicy, azureAccountToken, _dataSourceFactory, ownerUri);
         }
     }
 }

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

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

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

@@ -89,8 +89,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
         public abstract string GenerateExecuteFunctionScript(string functionName);
 
-        public abstract void UpdateAzureToken(string azureToken);
-
         /// <inheritdoc/>
         public DataSourceType DataSourceType { get; protected set; }
         

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

@@ -13,7 +13,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
     [Export(typeof(IDataSourceFactory))]
     public class DataSourceFactory : IDataSourceFactory
     {
-        public IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken)
+        public IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken, string ownerUri)
         {
             ValidationUtils.IsArgumentNotNullOrWhiteSpace(connectionString, nameof(connectionString));
             ValidationUtils.IsArgumentNotNullOrWhiteSpace(azureAccountToken, nameof(azureAccountToken));
@@ -22,7 +22,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             {
                 case DataSourceType.Kusto:
                 {
-                    var kustoClient = new KustoClient(connectionString, azureAccountToken);
+                    var kustoClient = new KustoClient(connectionString, azureAccountToken, ownerUri);
                     return new KustoDataSource(kustoClient);
                 }
 

src/Microsoft.Kusto.ServiceLayer/DataSource/Exceptions/DataSourceUnauthorizedException.cs

@@ -1,11 +0,0 @@
-using System;
-
-namespace Microsoft.Kusto.ServiceLayer.DataSource.Exceptions
-{
-    public class DataSourceUnauthorizedException : Exception
-    {
-        public DataSourceUnauthorizedException(Exception ex) : base (ex.Message, ex)
-        {
-        }
-    }
-}

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

@@ -111,11 +111,5 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// <param name="functionName"></param>
         /// <returns></returns>
         string GenerateExecuteFunctionScript(string functionName);
-
-        /// <summary>
-        /// Updates Azure Token
-        /// </summary>
-        /// <param name="azureToken"></param>
-        void UpdateAzureToken(string azureToken);
     }
 }

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

@@ -2,6 +2,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
     public interface IDataSourceFactory
     {
-        IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken);
+        IDataSource Create(DataSourceType dataSourceType, string connectionString, string azureAccountToken, string ownerUri);
     }
 }

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

@@ -17,9 +17,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         
         string DatabaseName { get; }
         
-        void UpdateAzureToken(string azureAccountToken);
+        IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null, int retryCount = 1);
-        
-        IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null);
 
         /// <summary>
         /// Executes a query or command against a kusto cluster and returns a sequence of result row instances.
@@ -37,7 +35,8 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// Executes a Kusto control command.
         /// </summary>
         /// <param name="command">The command.</param>
-        void ExecuteControlCommand(string command);
+        /// <param name="retryCount"></param>
+        void ExecuteControlCommand(string command, int retryCount = 1);
 
         void UpdateDatabase(string databaseName);
         

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

@@ -14,14 +14,16 @@ 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.DataSource.Exceptions;
 using Microsoft.Kusto.ServiceLayer.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.DataSource
 {
     public class KustoClient : IKustoClient
     {
+        private readonly string _ownerUri;
+
         [DebuggerBrowsable(DebuggerBrowsableState.Never)]
         private ICslAdminProvider _kustoAdminProvider;
 
@@ -36,8 +38,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         public string ClusterName { get; }
         public string DatabaseName { get; private set; }
 
-        public KustoClient(string connectionString, string azureAccountToken)
+        public KustoClient(string connectionString, string azureAccountToken, string ownerUri)
         {
+            _ownerUri = ownerUri;
             ClusterName = GetClusterName(connectionString);
             var databaseName = new SqlConnectionStringBuilder(connectionString).InitialCatalog;
             Initialize(ClusterName, databaseName, azureAccountToken);
@@ -75,8 +78,9 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             _kustoAdminProvider = KustoClientFactory.CreateCslAdminProvider(stringBuilder);
         }
 
-        public void UpdateAzureToken(string azureAccountToken)
+        private void RefreshAzureToken()
         {
+            string azureAccountToken = ConnectionService.Instance.RefreshAzureToken(_ownerUri);
             _kustoQueryProvider.Dispose();
             _kustoAdminProvider.Dispose();
             Initialize(ClusterName, DatabaseName, azureAccountToken);
@@ -162,7 +166,7 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             }
         }
 
-        public IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null)
+        public IDataReader ExecuteQuery(string query, CancellationToken cancellationToken, string databaseName = null, int retryCount = 1)
         {
             ValidationUtils.IsArgumentNotNullOrWhiteSpace(query, nameof(query));
 
@@ -175,29 +179,33 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
 
             var script = CodeScript.From(query, GlobalState.Default);
             IDataReader[] origReaders = new IDataReader[script.Blocks.Count];
-
+            try
-            Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
             {
-                var minimalQuery = codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
+                Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
-                
-                try
                 {
+                    var minimalQuery =
+                        codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
                     IDataReader origReader = _kustoQueryProvider.ExecuteQuery(
                         KustoQueryUtils.IsClusterLevelQuery(minimalQuery) ? "" : databaseName,
                         minimalQuery,
                         clientRequestProperties);
-                    
+
                     origReaders[index] = origReader;
-                }
+                });
-                catch (KustoRequestException exception) when (exception.FailureCode == 401) // Unauthorized
+                
-                {
+                return new KustoResultsReader(origReaders);
-                    throw new DataSourceUnauthorizedException(exception);
+            }
-                }
+            catch (AggregateException exception) 
-            });
+                when (retryCount > 0 &&
-
+                      exception.InnerException is KustoRequestException innerException 
-            return new KustoResultsReader(origReaders);
+                      && innerException.FailureCode == 401) // Unauthorized
+            {
+                RefreshAzureToken();
+                retryCount--;
+                return ExecuteQuery(query, cancellationToken, databaseName, retryCount);
+            }
         }
-
+        
         /// <summary>
         /// Executes a query or command against a kusto cluster and returns a sequence of result row instances.
         /// </summary>
@@ -211,10 +219,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
                 var tableReader = results[WellKnownDataSet.PrimaryResult].Single().TableData.CreateDataReader();
                 return new ObjectReader<T>(tableReader);
             }
-            catch (DataSourceUnauthorizedException)
-            {
-                throw;
-            }
             catch (Exception) when (!throwOnError)
             {
                 return null;
@@ -243,12 +247,22 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
         /// Executes a Kusto control command.
         /// </summary>
         /// <param name="command">The command.</param>
-        public void ExecuteControlCommand(string command)
+        /// <param name="retryCount"></param>
+        public void ExecuteControlCommand(string command, int retryCount = 1)
         {
             ValidationUtils.IsArgumentNotNullOrWhiteSpace(command, nameof(command));
 
-            using (var adminOutput = _kustoAdminProvider.ExecuteControlCommand(command, null))
+            try
             {
+                using (var adminOutput = _kustoAdminProvider.ExecuteControlCommand(command, null))
+                {
+                }
+            }
+            catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
+            {
+                RefreshAzureToken();
+                retryCount--;
+                ExecuteControlCommand(command, retryCount);
             }
         }
 

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

@@ -796,10 +796,6 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource
             return string.IsNullOrWhiteSpace(objectName) ? databaseName : $"{databaseName}.{objectName}";
         }
 
-        public override void UpdateAzureToken(string azureToken)
-        {
-            _kustoClient.UpdateAzureToken(azureToken);
-        }
         #endregion
     }
 }

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

@@ -42,8 +42,9 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         private readonly Guid _azureSessionId = Guid.NewGuid();
 
         private readonly string _connectionString;
-        private readonly string _azureAccountToken;
+        private string _azureAccountToken;
         private readonly IDataSourceFactory _dataSourceFactory;
+        private readonly string _ownerUri;
 
         /// <summary>
         /// Initializes a new instance of the ReliableKustoClient class with a given connection string
@@ -55,13 +56,15 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         /// <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)
+            RetryPolicy commandRetryPolicy, string azureAccountToken, IDataSourceFactory dataSourceFactory, string ownerUri)
         {
             _connectionString = connectionString;
             _azureAccountToken = azureAccountToken;
             _dataSourceFactory = dataSourceFactory;
-            _dataSource = dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccountToken);
+            _ownerUri = ownerUri;
+            _dataSource = dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccountToken, ownerUri);
             
             _connectionRetryPolicy = connectionRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
             _commandRetryPolicy = commandRetryPolicy ?? RetryPolicyFactory.CreateNoRetryPolicy();
@@ -190,7 +193,7 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             {
                 _connectionRetryPolicy.ExecuteAction(() =>
                 {
-                    _dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, _connectionString, _azureAccountToken);
+                    _dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, _connectionString, _azureAccountToken, _ownerUri);
                 });
             }
         }
@@ -247,6 +250,11 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
         {
             get { return _dataSource.DatabaseName; }
         }
+
+        public void UpdateAzureToken(string token)
+        {
+            _azureAccountToken = token;
+        }
     }
 }
 

src/Microsoft.Kusto.ServiceLayer/HostLoader.cs

@@ -12,7 +12,6 @@ using Microsoft.Kusto.ServiceLayer.Admin;
 using Microsoft.Kusto.ServiceLayer.Metadata;
 using Microsoft.Kusto.ServiceLayer.Connection;
 using Microsoft.Kusto.ServiceLayer.DataSource;
-using Microsoft.Kusto.ServiceLayer.DataSource.Metadata;
 using Microsoft.Kusto.ServiceLayer.LanguageServices;
 using Microsoft.Kusto.ServiceLayer.QueryExecution;
 using Microsoft.Kusto.ServiceLayer.Scripting;
@@ -91,7 +90,7 @@ namespace Microsoft.Kusto.ServiceLayer
             QueryExecutionService.Instance.InitializeService(serviceHost);
             serviceProvider.RegisterSingleService(QueryExecutionService.Instance);
 
-            ScriptingService.Instance.InitializeService(serviceHost, scripter, dataSourceFactory);
+            ScriptingService.Instance.InitializeService(serviceHost, scripter, ConnectionService.Instance);
             serviceProvider.RegisterSingleService(ScriptingService.Instance);
 
             AdminService.Instance.InitializeService(serviceHost, ConnectionService.Instance);

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

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

src/Microsoft.Kusto.ServiceLayer/QueryExecution/Batch.cs

@@ -15,7 +15,6 @@ using Microsoft.Kusto.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.Kusto.ServiceLayer.QueryExecution.DataStorage;
 using Microsoft.SqlTools.Utility;
 using System.Globalization;
-using Microsoft.Kusto.ServiceLayer.DataSource.Exceptions;
 
 namespace Microsoft.Kusto.ServiceLayer.QueryExecution
 {
@@ -66,8 +65,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         /// </summary>
         private readonly bool getFullColumnSchema;
 
-        private int _retryCount;
-
         #endregion
 
         internal Batch(string batchText, SelectionData selection, int ordinalId,
@@ -88,7 +85,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             this.outputFileFactory = outputFileFactory;
             specialAction = new SpecialAction();
             BatchExecutionCount = executionCount > 0 ? executionCount : 1;
-            _retryCount = 1;
             this.getFullColumnSchema = getFullColumnSchema;
         }
 
@@ -252,7 +248,7 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
         public async Task Execute(ReliableDataSourceConnection conn, CancellationToken cancellationToken)
         {
             // Sanity check to make sure we haven't already run this batch
-            if (HasExecuted && _retryCount < 0)
+            if (HasExecuted)
             {
                 throw new InvalidOperationException("Batch has already executed.");
             }
@@ -267,12 +263,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
             {
                 await DoExecute(conn, cancellationToken);
             }
-            catch (DataSourceUnauthorizedException)
-            {
-                // Rerun the query once if unauthorized
-                _retryCount--;
-                throw;
-            }
             catch (TaskCanceledException)
             {
                 // Cancellation isn't considered an error condition

src/Microsoft.Kusto.ServiceLayer/QueryExecution/Query.cs

@@ -16,7 +16,6 @@ using Microsoft.SqlTools.Utility;
 using Microsoft.SqlTools.ServiceLayer.BatchParser.ExecutionEngineCode;
 using System.Collections.Generic;
 using System.Diagnostics;
-using Microsoft.Kusto.ServiceLayer.DataSource.Exceptions;
 using Microsoft.Kusto.ServiceLayer.Utility;
 
 namespace Microsoft.Kusto.ServiceLayer.QueryExecution
@@ -391,11 +390,6 @@ namespace Microsoft.Kusto.ServiceLayer.QueryExecution
                     await QueryCompleted(this);
                 }
             }
-            catch (DataSourceUnauthorizedException)
-            {
-                ConnectionService.Instance.RefreshAzureToken(editorConnection.OwnerUri);
-                await ExecuteInternal();
-            }
             catch (Exception e)
             {
                 HasErrored = true;

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

@@ -24,20 +24,15 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
     {
         private readonly IScripter _scripter;
         private static readonly Dictionary<string, SqlServerVersion> scriptCompatibilityMap = LoadScriptCompatibilityMap();
+        private string _serverName;
+        private string _databaseName;
 
-        public ScriptAsScriptingOperation(ScriptingParams parameters, string azureAccountToken, IScripter scripter, IDataSourceFactory dataSourceFactory) : base(parameters, dataSourceFactory)
+        public ScriptAsScriptingOperation(ScriptingParams parameters, IScripter scripter, IDataSource datasource) :
+            base(parameters, datasource)
         {
-            DataSource = _dataSourceFactory.Create(DataSourceType.Kusto, this.Parameters.ConnectionString,
-                azureAccountToken);
             _scripter = scripter;
         }
 
-        internal IDataSource DataSource { get; set; }
-
-        private string serverName;
-        private string databaseName;
-        private bool disconnectAtDispose = false;
-
         public override void Execute()
         {
             try
@@ -49,7 +44,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                 this.CancellationToken.ThrowIfCancellationRequested();
                 string resultScript = string.Empty;
                 
-                UrnCollection urns = CreateUrns(DataSource);
+                UrnCollection urns = CreateUrns(_dataSource);
                 ScriptingOptions options = new ScriptingOptions();
                 SetScriptBehavior(options);
                 ScriptAsOptions scriptAsOptions = new ScriptAsOptions(this.Parameters.ScriptOptions);
@@ -65,12 +60,12 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                 switch (this.Parameters.Operation)
                 {
                     case ScriptingOperationType.Select:
-                        resultScript = GenerateScriptSelect(DataSource, urns);
+                        resultScript = GenerateScriptSelect(_dataSource, urns);
                         break;
                     
                     case ScriptingOperationType.Alter:
                     case ScriptingOperationType.Execute:
-                        resultScript = GenerateScriptForFunction(DataSource);
+                        resultScript = GenerateScriptForFunction(_dataSource);
                         break;
                 }
 
@@ -118,13 +113,6 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                     });
                 }
             }
-            finally
-            {
-                if (disconnectAtDispose && DataSource != null)
-                {
-                    DataSource.Dispose();
-                }
-            }
         }
 
         private string GenerateScriptSelect(IDataSource dataSource, UrnCollection urns)
@@ -168,8 +156,8 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
         {
             IEnumerable<ScriptingObject> selectedObjects = new List<ScriptingObject>(this.Parameters.ScriptingObjects);
 
-            serverName = dataSource.ClusterName;
+            _serverName = dataSource.ClusterName;
-            databaseName = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
+            _databaseName = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
             UrnCollection urnCollection = new UrnCollection();
             foreach (var scriptingObject in selectedObjects)
             {
@@ -178,7 +166,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                     // TODO: get the default schema
                     scriptingObject.Schema = "dbo";
                 }
-                urnCollection.Add(scriptingObject.ToUrn(serverName, databaseName));
+                urnCollection.Add(scriptingObject.ToUrn(_serverName, _databaseName));
             }
             return urnCollection;
         }

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

@@ -20,18 +20,15 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
     /// </summary>
     public sealed class ScriptingScriptOperation : SmoScriptingOperation
     {
-
         private int scriptedObjectCount = 0;
 
         private int totalScriptedObjectCount = 0;
 
         private int eventSequenceNumber = 1;
 
-        private string azureAccessToken;
+        public ScriptingScriptOperation(ScriptingParams parameters, IDataSource dataSource) : base(parameters, dataSource)
-
-        public ScriptingScriptOperation(ScriptingParams parameters, string azureAccessToken, IDataSourceFactory dataSourceFactory) : base(parameters, dataSourceFactory)
         {
-            this.azureAccessToken = azureAccessToken;
+            
         }
 
         public override void Execute()
@@ -204,7 +201,7 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
                     selectedObjects.Count(),
                     string.Join(", ", selectedObjects)));
 
-            string server = GetServerNameFromLiveInstance(this.Parameters.ConnectionString, this.azureAccessToken);
+            string server = GetServerNameFromLiveInstance();
             string database = new SqlConnectionStringBuilder(this.Parameters.ConnectionString).InitialCatalog;
 
             foreach (ScriptingObject scriptingObject in selectedObjects)

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

@@ -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>
+        /// <param name="scripter"></param>
-        public void InitializeService(ServiceHost serviceHost, IScripter scripter, IDataSourceFactory dataSourceFactory)
+        /// <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();

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

@@ -20,14 +20,13 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
     /// </summary>
     public abstract class SmoScriptingOperation : ScriptingOperation
     {
-        protected readonly IDataSourceFactory _dataSourceFactory;
+        protected readonly IDataSource _dataSource;
         private bool _disposed;
 
-        protected SmoScriptingOperation(ScriptingParams parameters, IDataSourceFactory dataSourceFactory)
+        protected SmoScriptingOperation(ScriptingParams parameters, IDataSource datasource)
         {
-            _dataSourceFactory = dataSourceFactory;
+            _dataSource = datasource;
             Validate.IsNotNull("parameters", parameters);
-
             this.Parameters = parameters;
         }
 
@@ -73,17 +72,10 @@ namespace Microsoft.Kusto.ServiceLayer.Scripting
             parameters.OperationId = this.OperationId;
         }
 
-        protected string GetServerNameFromLiveInstance(string connectionString, string azureAccessToken)
+        protected string GetServerNameFromLiveInstance()
         {
-            string serverName = string.Empty;
+            Logger.Write(TraceEventType.Verbose, string.Format("Resolved server name '{0}'", _dataSource.ClusterName));
-            
+            return _dataSource.ClusterName;
-            using(var dataSource = _dataSourceFactory.Create(DataSourceType.Kusto, connectionString, azureAccessToken))
-            {
-                serverName = dataSource.ClusterName;
-            }
-
-            Logger.Write(TraceEventType.Verbose, string.Format("Resolved server name '{0}'", serverName));
-            return serverName;
         }
 
         protected void ValidateScriptDatabaseParams()

test/Microsoft.Kusto.ServiceLayer.UnitTests/Connection/ConnectionInfoTests.cs

@@ -36,7 +36,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
 
             var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
             var reliableDataSource = new ReliableDataSourceConnection("", RetryPolicyFactory.NoRetryPolicy,
-                RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object);
+                RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object, "");
             connectionInfo.AddConnection("ConnectionType", reliableDataSource);
 
             connectionInfo.TryGetConnection("ConnectionType", out var connection);
@@ -60,7 +60,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
 
             var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
             var reliableDataSource = new ReliableDataSourceConnection("", RetryPolicyFactory.NoRetryPolicy,
-                RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object);
+                RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object, "");
             connectionInfo.AddConnection("ConnectionType", reliableDataSource);
 
             connectionInfo.RemoveConnection("ConnectionType");
@@ -77,7 +77,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.Connection
 
             var dataSourceFactoryMock = new Mock<IDataSourceFactory>();
             var reliableDataSource = new ReliableDataSourceConnection("", RetryPolicyFactory.NoRetryPolicy,
-                RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object);
+                RetryPolicyFactory.NoRetryPolicy, "", dataSourceFactoryMock.Object, "");
             connectionInfo.AddConnection("ConnectionType", reliableDataSource);
 
             connectionInfo.RemoveAllConnections();

test/Microsoft.Kusto.ServiceLayer.UnitTests/Connection/DataSourceConnectionFactoryTests.cs

@@ -12,7 +12,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("", "", "");
             
             Assert.IsNotNull(connection);
         }

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

@@ -19,7 +19,7 @@ namespace Microsoft.Kusto.ServiceLayer.UnitTests.DataSource
         {
             var dataSourceFactory = new DataSourceFactory();
             Assert.Throws(exceptionType,
-                () => dataSourceFactory.Create(DataSourceType.None, connectionString, azureAccountToken));
+                () => dataSourceFactory.Create(DataSourceType.None, connectionString, azureAccountToken, ""));
         }
 
         [Test]

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

@@ -97,7 +97,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>()))
+                .Setup(x => x.Create(It.IsAny<DataSourceType>(), It.IsAny<string>(), It.IsAny<string>(), It.IsAny<string>()))
                 .Returns(dataSourceMock.Object);
 
             var connectedBindingQueue =