successfully resetting Entra token during long query

Author: David Coe

csharp/src/Drivers/BigQuery/BigQueryConnection.cs

@@ -18,6 +18,7 @@
 using System;
 using System.Collections.Generic;
 using System.Collections.ObjectModel;
+using System.Diagnostics;
 using System.Linq;
 using System.Net.Http;
 using System.Text;
@@ -65,7 +66,6 @@ public BigQueryConnection(IReadOnlyDictionary<string, string> properties)
             this.properties = new ReadOnlyDictionary<string, string>(modifiedProperties);
             this.httpClient = new HttpClient();
         }
-
         public Func<Task>? UpdateToken { get; set; }
 
         internal BigQueryClient? Client { get; private set; }
@@ -1225,6 +1225,8 @@ private string Sanitize(string? input)
         {
             try
             {
+                Debug.WriteLine($"{DateTime.Now.ToString()} - Trading {entraAccessToken}");
+
                 var requestBody = new
                 {
                     scope = BigQueryConstants.EntraIdScope,
@@ -1245,6 +1247,8 @@ private string Sanitize(string? input)
 
                 BigQueryStsTokenResponse? bigQueryTokenResponse = JsonSerializer.Deserialize<BigQueryStsTokenResponse>(responseBody);
 
+                Debug.WriteLine($"{DateTime.Now.ToString()} - Traded for Google token {bigQueryTokenResponse?.AccessToken}");
+
                 return bigQueryTokenResponse?.AccessToken;
             }
             catch (Exception ex)

csharp/src/Drivers/BigQuery/BigQueryStatement.cs

@@ -59,7 +59,6 @@ public BigQueryStatement(BigQueryConnection bigQueryConnection)
 
         private GoogleCredential Credential => this.bigQueryConnection.Credential ?? throw new AdbcException("Credential cannot be null");
 
-
         public override QueryResult ExecuteQuery()
         {
             //Func<Task<QueryResult>> func = () => ExecuteQueryInternalAsync();
@@ -74,6 +73,8 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
 
             BigQueryJob job = await this.Client.CreateQueryJobAsync(SqlQuery, null, queryOptions);
 
+            JobReference jobReference = job.Reference;
+
             GetQueryResultsOptions getQueryResultsOptions = new GetQueryResultsOptions();
 
             if (this.Options?.TryGetValue(BigQueryParameters.GetQueryResultsOptionsTimeout, out string? timeoutSeconds) == true &&
@@ -83,16 +84,13 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
                 getQueryResultsOptions.Timeout = TimeSpan.FromSeconds(seconds);
             }
 
-            Task<BigQueryResults> resultsTask = job.GetQueryResultsAsync(getQueryResultsOptions);
-
             DateTime start = DateTime.Now;
 
-            Func<Task<bool>> func = () => Task.Run(() =>
+            Func<Task<BigQueryJob>> checkJobStatus = async () =>
             {
                 while (true)
                 {
-                    var jobWithStatus = this.Client.GetJob(job.Reference);
-                    Debug.WriteLine($"Job state is {jobWithStatus.State}");
+                    var jobWithStatus = await this.Client.GetJobAsync(jobReference);
 
                     if (jobWithStatus.State == JobState.Done)
                     {
@@ -104,49 +102,57 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
                         DateTime end = DateTime.Now;
                         TimeSpan duration = end - start;
                         Debug.WriteLine($"Done at {end.ToString()} after {duration.TotalMinutes}");
-                        return true;
+                        return jobWithStatus;
                     }
                 }
-            });
+            };
 
             Debug.WriteLine($"Starting ExecuteWithRetriesAsync at {start.ToString()}");
 
-            await AdbcRetryManager.ExecuteWithRetriesAsync<bool>(this, func);
+            await AdbcRetryManager.ExecuteWithRetriesAsync<BigQueryJob>(this, checkJobStatus);
 
             Debug.WriteLine($"Getting results at {DateTime.Now.ToString()}");
 
-            BigQueryResults results = resultsTask.Result;
+            Func<Task<BigQueryResults>> getJobResults = async () =>
+            {
+                // if the authentication token was reset, then we need a new job with the latest token
+                BigQueryJob completedJob = await this.Client.GetJobAsync(jobReference);
+                return await completedJob.GetQueryResultsAsync();
+            };
+
+            BigQueryResults results = await AdbcRetryManager.ExecuteWithRetriesAsync(this, getJobResults);
 
             Debug.WriteLine($"Results received at {DateTime.Now.ToString()}");
 
             TokenProtectedReadClientManger clientMgr = new TokenProtectedReadClientManger(this.Credential);
             clientMgr.UpdateToken = () => Task.Run(() =>
             {
+                Debug.WriteLine($"TokenProtectedReadClientManger updating token at {DateTime.Now.ToString()}");
                 this.bigQueryConnection.SetCredential();
                 clientMgr.UpdateCredential(this.Credential);
             });
 
-            if (results.TableReference == null)
+            if (results?.TableReference == null)
             {
                 // To get the results of all statements in a multi-statement query, enumerate the child jobs and call jobs.getQueryResults on each of them.
                 // Related public docs: https://cloud.google.com/bigquery/docs/multi-statement-queries#get_all_executed_statements
                 ListJobsOptions listJobsOptions = new ListJobsOptions();
-                listJobsOptions.ParentJobId = results.JobReference.JobId;
-                PagedEnumerable<JobList, BigQueryJob> joblist = this.Client.ListJobs(listJobsOptions);
+                listJobsOptions.ParentJobId = results?.JobReference.JobId;
+                PagedAsyncEnumerable<JobList, BigQueryJob> joblist = this.Client.ListJobsAsync(listJobsOptions);
                 BigQueryJob firstQueryJob = new BigQueryJob(this.Client, job.Resource);
-                foreach (BigQueryJob childJob in joblist)
+                await foreach (BigQueryJob childJob in joblist)
                 {
-                    var tempJob = this.Client.GetJob(childJob.Reference);
+                    var tempJob = await this.Client.GetJobAsync(childJob.Reference);
                     var query = tempJob.Resource?.Configuration?.Query;
                     if (query != null && query.DestinationTable != null && query.DestinationTable.ProjectId != null && query.DestinationTable.DatasetId != null && query.DestinationTable.TableId != null)
                     {
                         firstQueryJob = tempJob;
                     }
                 }
-                results = firstQueryJob.GetQueryResults();
+                results = await firstQueryJob.GetQueryResultsAsync();
             }
 
-            if (results.TableReference == null)
+            if (results?.TableReference == null)
             {
                 throw new AdbcException("There is no query statement");
             }
@@ -156,6 +162,7 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
             string table = $"projects/{results.TableReference.ProjectId}/datasets/{results.TableReference.DatasetId}/tables/{results.TableReference.TableId}";
 
             int maxStreamCount = 1;
+
             if (this.Options?.TryGetValue(BigQueryParameters.MaxFetchConcurrency, out string? maxStreamCountString) == true)
             {
                 if (int.TryParse(maxStreamCountString, out int count))
@@ -166,28 +173,34 @@ private async Task<QueryResult> ExecuteQueryInternalAsync()
                     }
                 }
             }
+
             ReadSession rs = new ReadSession { Table = table, DataFormat = DataFormat.Arrow };
             ReadSession rrs = clientMgr.ReadClient.CreateReadSession("projects/" + results.TableReference.ProjectId, rs, maxStreamCount);
 
             long totalRows = results.TotalRows == null ? -1L : (long)results.TotalRows.Value;
 
+            Debug.WriteLine($"Starting to read streams at {DateTime.Now.ToString()}");
+
             var readers = rrs.Streams
                              .Select(s => ReadChunkWithRetries(clientMgr, s.Name))
                              .Where(chunk => chunk != null)
                              .Cast<IArrowReader>();
 
+            Debug.WriteLine($"Creating Arrow stream at {DateTime.Now.ToString()}");
+
             IArrowArrayStream stream = new MultiArrowReader(TranslateSchema(results.Schema), readers);
 
+            Debug.WriteLine($"Returning results at {DateTime.Now.ToString()}");
+
             return new QueryResult(totalRows, stream);
         }
 
         public override UpdateResult ExecuteUpdate()
         {
-            Func<Task<UpdateResult>> func = () => Task.Run(() => ExecuteUpdateInternal());
-            return AdbcRetryManager.ExecuteWithRetriesAsync<UpdateResult>(this, func).Result;
+           return ExecuteUpdateInternalAsync().GetAwaiter().GetResult();
         }
 
-        private UpdateResult ExecuteUpdateInternal()
+        private async Task<UpdateResult> ExecuteUpdateInternalAsync()
         {
             QueryOptions options = ValidateOptions();
             GetQueryResultsOptions getQueryResultsOptions = new GetQueryResultsOptions();
@@ -199,7 +212,7 @@ private UpdateResult ExecuteUpdateInternal()
                 getQueryResultsOptions.Timeout = TimeSpan.FromSeconds(seconds);
             }
 
-            BigQueryResults result = this.Client.ExecuteQuery(
+            BigQueryResults result = await this.Client.ExecuteQueryAsync(
                 SqlQuery,
                 parameters: null,
                 queryOptions: options,

csharp/test/Drivers/BigQuery/AuthenticationTests.cs

@@ -46,7 +46,7 @@ public AuthenticationTests(ITestOutputHelper outputHelper)
         }
 
         /// <summary>
-        /// Validates if the Entra token can sign in and refresh.
+        /// Validates if the Entra token can sign in.
         /// </summary>
         [SkippableFact, Order(1)]
         public void CanSignInWithEntraToken()
@@ -70,8 +70,11 @@ public void CanSignInWithEntraToken()
             Tests.DriverTests.CanExecuteQuery(queryResult, environment.ExpectedResultsCount, environment.Name);
         }
 
+        /// <summary>
+        /// Validates if the Entra token can sign in and refresh.
+        /// </summary>
         [SkippableFact, Order(1)]
-        public void CanRefreshToken()
+        public void CanSignInWithAndRefreshEntraToken()
         {
             BigQueryTestEnvironment? environment = _environments.Where(x => x.AuthenticationType == BigQueryConstants.EntraIdAuthenticationType).FirstOrDefault();
             Assert.NotNull(environment);
@@ -86,10 +89,20 @@ public void CanRefreshToken()
                 _outputHelper.WriteLine("Successfully set a new token");
             });
 
+            // create a query that takes 75 minutes because Entra tokens typically expire in 60 minutes
             AdbcStatement statement = connection.CreateStatement();
-            statement.SqlQuery = "SELECT COUNT(*) FROM UNNEST(GENERATE_ARRAY(1, 1000000, 1)) AS a, UNNEST(GENERATE_ARRAY(1, 1000000, 1)) AS b WHERE RAND() < 0.0000001";
+            statement.SqlQuery = @"
+                DECLARE end_time TIMESTAMP;
+                SET end_time = TIMESTAMP_ADD(CURRENT_TIMESTAMP(), INTERVAL 75 MINUTE);
+
+                WHILE CURRENT_TIMESTAMP() < end_time DO
+                END WHILE;
+
+                SELECT 'Query completed after 75 minutes' AS result;";
 
             QueryResult queryResult = statement.ExecuteQuery();
+
+            _outputHelper.WriteLine($"Retrieve query result with {queryResult.RowCount} rows");
         }
 
         private string GetAccessToken(BigQueryTestEnvironment environment)
@@ -107,6 +120,8 @@ private string GetAccessToken(BigQueryTestEnvironment environment)
             TokenRequestContext requestContext = new TokenRequestContext(environment.EntraConfiguration.Scopes, claims: claimJson);
             AccessToken accessToken = credential.GetToken(requestContext);
 
+            _outputHelper.WriteLine($"Access token is {accessToken.Token}");
+
             return accessToken.Token;
         }
     }