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/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];
-
-            Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
+            try
             {
-                var minimalQuery = codeBlock.Service.GetMinimalText(MinimalTextKind.RemoveLeadingWhitespaceAndComments);
-                
-                try
+                Parallel.ForEach(script.Blocks, (codeBlock, state, index) =>
                 {
+                    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
-                {
-                    throw new DataSourceUnauthorizedException(exception);
-                }
-            });
-
-            return new KustoResultsReader(origReaders);
+                });
+                
+                return new KustoResultsReader(origReaders);
+            }
+            catch (AggregateException exception) 
+                when (retryCount > 0 &&
+                      exception.InnerException is KustoRequestException innerException 
+                      && 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);
             }
         }