Changed RefreshAuthToken in ConnectionService and KustoClient to foll…

…ow TryParse pattern. (#1245)
microsoft · Jul 19, 2022 · 1db0237 · 1db0237
1 parent 8356a8c
commit 1db0237
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 9 deletions.
diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs
@@ -213,9 +213,13 @@ public async Task<ConnectionCompleteParams> Connect(ConnectParams connectionPara
             return completeParams;
         }
 
-        internal string RefreshAuthToken(string ownerUri)
+        internal bool TryRefreshAuthToken(string ownerUri, out string token)
         {
-            TryFindConnection(ownerUri, out ConnectionInfo connection);
+            token = string.Empty;
+            if (!TryFindConnection(ownerUri, out ConnectionInfo connection))
+            {
+                return false;
+            }
 
             var requestMessage = new RequestSecurityTokenParams
             {
@@ -227,8 +231,8 @@ internal string RefreshAuthToken(string ownerUri)
 
             var response = _serviceHost.SendRequest(SecurityTokenRequest.Type, requestMessage, true).Result;
             connection.UpdateAuthToken(response.Token);
-
-            return response.Token;
+            token = response.Token;
+            return true;
         }
 
         private void TryCloseConnectionTemporaryConnection(ConnectParams connectionParams, ConnectionInfo connectionInfo)

diff --git a/src/Microsoft.Kusto.ServiceLayer/DataSource/Kusto/KustoClient.cs b/src/Microsoft.Kusto.ServiceLayer/DataSource/Kusto/KustoClient.cs
@@ -54,9 +54,13 @@ private void Initialize(DataSourceConnectionDetails connectionDetails)
             DatabaseName = databaseName?.ToString() ?? "";
         }
 
-        private void RefreshAuthToken()
+        private bool TryRefreshAuthToken()
         {
-            string accountToken = ConnectionService.Instance.RefreshAuthToken(_ownerUri);
+            if (!ConnectionService.Instance.TryRefreshAuthToken(_ownerUri, out string accountToken))
+            {
+                return false;
+            }
+
             _kustoQueryProvider.Dispose();
             _kustoAdminProvider.Dispose();
 
@@ -69,6 +73,7 @@ private void RefreshAuthToken()
             };
 
             Initialize(connectionDetails);
+            return true;
         }
 
         private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(DataSourceConnectionDetails connectionDetails)
@@ -163,7 +168,10 @@ public IDataReader ExecuteQuery(string query, CancellationToken cancellationToke
                       exception.InnerException is KustoRequestException innerException
                       && innerException.FailureCode == 401) // Unauthorized
             {
-                RefreshAuthToken();
+                if (!TryRefreshAuthToken())
+                {
+                    throw;
+                }
                 retryCount--;
                 return ExecuteQuery(query, cancellationToken, databaseName, retryCount);
             }
@@ -184,7 +192,10 @@ public async Task ExecuteControlCommandAsync(string command, bool throwOnError,
             }
             catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
             {
-                RefreshAuthToken();
+                if (!TryRefreshAuthToken())
+                {
+                    throw;
+                }
                 retryCount--;
                 await ExecuteControlCommandAsync(command, throwOnError, retryCount);
             }
@@ -232,7 +243,10 @@ public void ExecuteControlCommand(string command, int retryCount = 1)
             }
             catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
             {
-                RefreshAuthToken();
+                if (!TryRefreshAuthToken())
+                {
+                    throw;
+                }
                 retryCount--;
                 ExecuteControlCommand(command, retryCount);
             }