feat(csharp/src/Drivers/BigQuery): add support for configurable query…

… timeouts (#2043) Adds a new option to support specifying the timeout value for a query --------- Co-authored-by: David Coe <coedavid@umich.edu>
apache · Jul 31, 2024 · bf31b1f · bf31b1f
1 parent 8fb7c2f
commit bf31b1f
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 16 deletions.
diff --git a/csharp/src/Drivers/BigQuery/BigQueryConnection.cs b/csharp/src/Drivers/BigQuery/BigQueryConnection.cs
@@ -997,23 +997,19 @@ private IReadOnlyDictionary<string, string> ParseOptions()
         {
             Dictionary<string, string> options = new Dictionary<string, string>();
 
-            foreach (KeyValuePair<string, string> keyValuePair in this.properties)
+            string[] statementOptions = new string[] {
+                BigQueryParameters.AllowLargeResults,
+                BigQueryParameters.UseLegacySQL,
+                BigQueryParameters.LargeDecimalsAsString,
+                BigQueryParameters.LargeResultsDestinationTable,
+                BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes
+            };
+
+            foreach (string key in statementOptions)
             {
-                if (keyValuePair.Key == BigQueryParameters.AllowLargeResults)
-                {
-                    options[keyValuePair.Key] = keyValuePair.Value;
-                }
-                if (keyValuePair.Key == BigQueryParameters.UseLegacySQL)
-                {
-                    options[keyValuePair.Key] = keyValuePair.Value;
-                }
-                if (keyValuePair.Key == BigQueryParameters.LargeDecimalsAsString)
-                {
-                    options[keyValuePair.Key] = keyValuePair.Value;
-                }
-                if (keyValuePair.Key == BigQueryParameters.LargeResultsDestinationTable)
+                if (properties.TryGetValue(key, out string? value))
                 {
-                    options[keyValuePair.Key] = keyValuePair.Value;
+                    options[key] = value;
                 }
             }
 

diff --git a/csharp/src/Drivers/BigQuery/BigQueryParameters.cs b/csharp/src/Drivers/BigQuery/BigQueryParameters.cs
@@ -34,6 +34,7 @@ public class BigQueryParameters
         public const string LargeDecimalsAsString = "adbc.bigquery.large_decimals_as_string";
         public const string Scopes = "adbc.bigquery.scopes";
         public const string IncludeConstraintsWithGetObjects = "adbc.bigquery.include_constraints_getobjects";
+        public const string GetQueryResultsOptionsTimeoutMinutes = "adbc.bigquery.get_query_results_options.timeout";
     }
 
     /// <summary>

diff --git a/csharp/src/Drivers/BigQuery/BigQueryStatement.cs b/csharp/src/Drivers/BigQuery/BigQueryStatement.cs
@@ -55,7 +55,21 @@ public override QueryResult ExecuteQuery()
         {
             QueryOptions? queryOptions = ValidateOptions();
             BigQueryJob job = this.client.CreateQueryJob(SqlQuery, null, queryOptions);
-            BigQueryResults results = job.GetQueryResults();
+
+            GetQueryResultsOptions getQueryResultsOptions = new GetQueryResultsOptions();
+
+            if (this.Options?.TryGetValue(BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes, out string? timeoutMinutes) == true)
+            {
+                if (int.TryParse(timeoutMinutes, out int minutes))
+                {
+                    if (minutes >= 0)
+                    {
+                        getQueryResultsOptions.Timeout = TimeSpan.FromMinutes(minutes);
+                    }
+                }
+            }
+
+            BigQueryResults results = job.GetQueryResults(getQueryResultsOptions);
 
             BigQueryReadClientBuilder readClientBuilder = new BigQueryReadClientBuilder();
             readClientBuilder.Credential = this.credential;

diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs b/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs
@@ -56,5 +56,8 @@ public BigQueryTestConfiguration()
 
         [JsonPropertyName("includeTableConstraints")]
         public bool IncludeTableConstraints { get; set; }
+
+        [JsonPropertyName("timeoutMinutes")]
+        public int? TimeoutMinutes { get; set; }
     }
 }
diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs b/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs
@@ -86,6 +86,11 @@ internal static Dictionary<string, string> GetBigQueryParameters(BigQueryTestCon
                 parameters.Add(BigQueryParameters.LargeResultsDestinationTable, testConfiguration.LargeResultsDestinationTable);
             }
 
+            if (testConfiguration.TimeoutMinutes.HasValue)
+            {
+                parameters.Add(BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes, testConfiguration.TimeoutMinutes.Value.ToString());
+            }
+
             return parameters;
         }