From 60c847ff6011ffa446f50b4c723dff353b7be86e Mon Sep 17 00:00:00 2001
From: Justin M <63619224+JustinMDotNet@users.noreply.github.com>
Date: Fri, 17 Sep 2021 12:25:11 -0700
Subject: [PATCH] Changed RefreshAuthToken in ConnectionService and KustoClient
 to follow TryParse pattern. (#1245)

---
 .../Connection/ConnectionService.cs           | 12 ++++++----
 .../DataSource/Kusto/KustoClient.cs           | 24 +++++++++++++++----
 2 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs
index c8ccfdb5..73dd2e53 100644
--- a/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/Connection/ConnectionService.cs
@@ -213,9 +213,13 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
             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 @@ namespace Microsoft.Kusto.ServiceLayer.Connection
 
             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
index fd914217..e9c34ce6 100644
--- a/src/Microsoft.Kusto.ServiceLayer/DataSource/Kusto/KustoClient.cs
+++ b/src/Microsoft.Kusto.ServiceLayer/DataSource/Kusto/KustoClient.cs
@@ -54,9 +54,13 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
             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 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
             };
             
             Initialize(connectionDetails);
+            return true;
         }
 
         private KustoConnectionStringBuilder GetKustoConnectionStringBuilder(DataSourceConnectionDetails connectionDetails)
@@ -163,7 +168,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
                       exception.InnerException is KustoRequestException innerException
                       && innerException.FailureCode == 401) // Unauthorized
             {
-                RefreshAuthToken();
+                if (!TryRefreshAuthToken())
+                {
+                    throw;
+                }
                 retryCount--;
                 return ExecuteQuery(query, cancellationToken, databaseName, retryCount);
             }
@@ -184,7 +192,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
             }
             catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
             {
-                RefreshAuthToken();
+                if (!TryRefreshAuthToken())
+                {
+                    throw;
+                }
                 retryCount--;
                 await ExecuteControlCommandAsync(command, throwOnError, retryCount);
             }
@@ -232,7 +243,10 @@ namespace Microsoft.Kusto.ServiceLayer.DataSource.Kusto
             }
             catch (KustoRequestException exception) when (retryCount > 0 && exception.FailureCode == 401) // Unauthorized
             {
-                RefreshAuthToken();
+                if (!TryRefreshAuthToken())
+                {
+                    throw;
+                }
                 retryCount--;
                 ExecuteControlCommand(command, retryCount);
             }