feat(csharp/src/Drivers/BigQuery): add additional billing and timeout properties and test settings (apache#2566)

davidhcoe · web-flow · commit 15ad9d185e25 · 2025-03-03T15:56:09.000-08:00
- adds the `adbc.bigquery.client.timeout` value for the BigQueryClient
- adds the `adbc.bigquery.billing_project_id` value for the billing
queries against
- adds support for referencing shared values in multi-environment test
configurations
- adds tests for the new settings

---------

Co-authored-by: David Coe &lt;&gt;
diff --git a/csharp/src/Drivers/BigQuery/BigQueryConnection.cs b/csharp/src/Drivers/BigQuery/BigQueryConnection.cs
@@ -73,9 +73,11 @@ public BigQueryConnection(IReadOnlyDictionary<string, string> properties)
         internal BigQueryClient Open()
         {
             string? projectId = null;
+            string? billingProjectId = null;
             string? clientId = null;
             string? clientSecret = null;
             string? refreshToken = null;
+            TimeSpan? clientTimeout = null;
 
             string tokenEndpoint = BigQueryConstants.TokenEndpoint;
 
@@ -87,6 +89,9 @@ internal BigQueryClient Open()
             if (!this.properties.TryGetValue(BigQueryParameters.ProjectId, out projectId))
                 projectId = BigQueryConstants.DetectProjectId;
 
+            // the billing project can be null if it's not specified
+            this.properties.TryGetValue(BigQueryParameters.BillingProjectId, out billingProjectId);
+
             if (this.properties.TryGetValue(BigQueryParameters.IncludePublicProjectId, out string? result))
             {
                 if (!string.IsNullOrEmpty(result))
@@ -128,7 +133,26 @@ internal BigQueryClient Open()
                 this.credential = ApplyScopes(GoogleCredential.FromJson(json));
             }
 
-            BigQueryClient client = BigQueryClient.Create(projectId, this.credential);
+            if (this.properties.TryGetValue(BigQueryParameters.ClientTimeout, out string? timeoutSeconds) &&
+                int.TryParse(timeoutSeconds, out int seconds))
+            {
+                clientTimeout = TimeSpan.FromSeconds(seconds);
+            }
+
+            BigQueryClientBuilder bigQueryClientBuilder = new BigQueryClientBuilder()
+            {
+                ProjectId = projectId,
+                QuotaProject = billingProjectId,
+                GoogleCredential = this.credential
+            };
+
+            BigQueryClient client = bigQueryClientBuilder.Build();
+
+            if (clientTimeout.HasValue)
+            {
+                client.Service.HttpClient.Timeout = clientTimeout.Value;
+            }
+
             this.client = client;
             return client;
         }
@@ -1033,7 +1057,7 @@ private IReadOnlyDictionary<string, string> ParseOptions()
                 BigQueryParameters.UseLegacySQL,
                 BigQueryParameters.LargeDecimalsAsString,
                 BigQueryParameters.LargeResultsDestinationTable,
-                BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes,
+                BigQueryParameters.GetQueryResultsOptionsTimeout,
                 BigQueryParameters.MaxFetchConcurrency
             };
 
diff --git a/csharp/src/Drivers/BigQuery/BigQueryParameters.cs b/csharp/src/Drivers/BigQuery/BigQueryParameters.cs
@@ -23,6 +23,7 @@ namespace Apache.Arrow.Adbc.Drivers.BigQuery
     public class BigQueryParameters
     {
         public const string ProjectId = "adbc.bigquery.project_id";
+        public const string BillingProjectId = "adbc.bigquery.billing_project_id";
         public const string ClientId = "adbc.bigquery.client_id";
         public const string ClientSecret = "adbc.bigquery.client_secret";
         public const string RefreshToken = "adbc.bigquery.refresh_token";
@@ -34,7 +35,8 @@ 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";
+        public const string ClientTimeout = "adbc.bigquery.client.timeout";
+        public const string GetQueryResultsOptionsTimeout = "adbc.bigquery.get_query_results_options.timeout";
         public const string MaxFetchConcurrency = "adbc.bigquery.max_fetch_concurrency";
         public const string IncludePublicProjectId = "adbc.bigquery.include_public_project_id";
     }
diff --git a/csharp/src/Drivers/BigQuery/BigQueryStatement.cs b/csharp/src/Drivers/BigQuery/BigQueryStatement.cs
@@ -57,15 +57,11 @@ public override QueryResult ExecuteQuery()
 
             GetQueryResultsOptions getQueryResultsOptions = new GetQueryResultsOptions();
 
-            if (this.Options?.TryGetValue(BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes, out string? timeoutMinutes) == true)
+            if (this.Options?.TryGetValue(BigQueryParameters.GetQueryResultsOptionsTimeout, out string? timeoutSeconds) == true &&
+                int.TryParse(timeoutSeconds, out int seconds) &&
+                seconds >= 0)
             {
-                if (int.TryParse(timeoutMinutes, out int minutes))
-                {
-                    if (minutes >= 0)
-                    {
-                        getQueryResultsOptions.Timeout = TimeSpan.FromMinutes(minutes);
-                    }
-                }
+                getQueryResultsOptions.Timeout = TimeSpan.FromSeconds(seconds);
             }
 
             BigQueryResults results = job.GetQueryResults(getQueryResultsOptions);
@@ -126,15 +122,11 @@ public override UpdateResult ExecuteUpdate()
             QueryOptions options = ValidateOptions();
             GetQueryResultsOptions getQueryResultsOptions = new GetQueryResultsOptions();
 
-            if (this.Options?.TryGetValue(BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes, out string? timeoutMinutes) == true)
+            if (this.Options?.TryGetValue(BigQueryParameters.GetQueryResultsOptionsTimeout, out string? timeoutSeconds) == true &&
+                int.TryParse(timeoutSeconds, out int seconds) &&
+                seconds >= 0)
             {
-                if (int.TryParse(timeoutMinutes, out int minutes))
-                {
-                    if (minutes >= 0)
-                    {
-                        getQueryResultsOptions.Timeout = TimeSpan.FromMinutes(minutes);
-                    }
-                }
+                getQueryResultsOptions.Timeout = TimeSpan.FromSeconds(seconds);
             }
 
             BigQueryResults result = this.client.ExecuteQuery(
diff --git a/csharp/src/Drivers/BigQuery/readme.md b/csharp/src/Drivers/BigQuery/readme.md
@@ -42,17 +42,23 @@ The following parameters can be used to configure the driver behavior. The param
 
 https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/Google.Cloud.BigQuery.V2.QueryOptions#Google_Cloud_BigQuery_V2_QueryOptions_AllowLargeResults
 
+**adbc.bigquery.billing_project_id**<br>
+&nbsp;&nbsp;&nbsp;&nbsp;The [Project ID](https://cloud.google.com/resource-manager/docs/creating-managing-projects) used for accessing billing BigQuery. If not specified, will default to the detected project ID.
+
 **adbc.bigquery.client_id**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;The OAuth client ID. Required for `user` authentication.
 
 **adbc.bigquery.client_secret**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;The OAuth client secret. Required for `user` authentication.
 
+**adbc.bigquery.client.timeout**<br>
+&nbsp;&nbsp;&nbsp;&nbsp;Optional. Sets the timeout (in seconds) for the BigQueryClient. Similar to a ConnectionTimeout.
+
 **adbc.bigquery.auth_json_credential**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;Required if using `service` authentication. This value is passed to the [GoogleCredential.FromJson](https://cloud.google.com/dotnet/docs/reference/Google.Apis/latest/Google.Apis.Auth.OAuth2.GoogleCredential#Google_Apis_Auth_OAuth2_GoogleCredential_FromJson_System_String) method.
 
 **adbc.bigquery.get_query_results_options.timeout**<br>
-&nbsp;&nbsp;&nbsp;&nbsp;Optional. Sets the timeout (in minutes) for the GetQueryResultsOptions value. If not set, defaults to 5 minutes.
+&nbsp;&nbsp;&nbsp;&nbsp;Optional. Sets the timeout (in seconds) for the GetQueryResultsOptions value. If not set, defaults to 5 minutes. Similar to a CommandTimeout.
 
 **adbc.bigquery.max_fetch_concurrency**<br>
 &nbsp;&nbsp;&nbsp;&nbsp;Optional. Sets the [maxStreamCount](https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.Storage.V1/latest/Google.Cloud.BigQuery.Storage.V1.BigQueryReadClient#Google_Cloud_BigQuery_Storage_V1_BigQueryReadClient_CreateReadSession_System_String_Google_Cloud_BigQuery_Storage_V1_ReadSession_System_Int32_Google_Api_Gax_Grpc_CallSettings_) for the CreateReadSession method. If not set, defaults to 1.
diff --git a/csharp/test/Apache.Arrow.Adbc.Tests/MultiEnvironmentTestConfiguration.cs b/csharp/test/Apache.Arrow.Adbc.Tests/MultiEnvironmentTestConfiguration.cs
@@ -15,7 +15,9 @@
 * limitations under the License.
 */
 
+using System;
 using System.Collections.Generic;
+using System.Text.Json;
 using System.Text.Json.Serialization;
 
 namespace Apache.Arrow.Adbc.Tests
@@ -42,5 +44,11 @@ public abstract class MultiEnvironmentTestConfiguration<T>
         /// </summary>
         [JsonPropertyName("environments")]
         public Dictionary<string, T> Environments { get; set; } = new Dictionary<string, T>();
+
+        /// <summary>
+        /// Values that are shared across environments so they don't need to be repeated.
+        /// </summary>
+        [JsonPropertyName("shared")]
+        public Dictionary<string, string> SharedKeyValuePairs { get; set; } = new Dictionary<string, string>();
     }
 }
diff --git a/csharp/test/Apache.Arrow.Adbc.Tests/MultiEnvironmentTestUtils.cs b/csharp/test/Apache.Arrow.Adbc.Tests/MultiEnvironmentTestUtils.cs
@@ -18,6 +18,7 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.Reflection;
 using System.Text.Json;
 
 namespace Apache.Arrow.Adbc.Tests
@@ -38,7 +39,6 @@ public static T LoadMultiEnvironmentTestConfiguration<T>(string environmentVaria
                     {
                         // use a JSON file for the various settings
                         string json = File.ReadAllText(environmentValue);
-
                         testConfiguration = JsonSerializer.Deserialize<T>(json)!;
                     }
                 }
@@ -60,17 +60,34 @@ public static List<TEnvironment> GetTestEnvironments<TEnvironment>(MultiEnvironm
                 throw new InvalidOperationException("There are no environments configured");
 
             List<TEnvironment> environments = new List<TEnvironment>();
+            string term = "$ref:shared.";
 
             foreach (string environmentName in GetEnvironmentNames(testConfiguration.TestEnvironmentNames))
             {
-                if (testConfiguration.Environments.TryGetValue(environmentName, out TEnvironment? testEnvironment))
+                if (!testConfiguration.Environments.TryGetValue(environmentName, out TEnvironment? testEnvironment) || testEnvironment is null)
+                    continue;
+
+                testEnvironment.Name = environmentName;
+
+                if (testConfiguration.SharedKeyValuePairs.Count > 0)
                 {
-                    if (testEnvironment != null)
+                    foreach (PropertyInfo pi in testEnvironment.GetType().GetProperties())
                     {
-                        testEnvironment.Name = environmentName;
-                        environments.Add(testEnvironment);
+                        if (pi.PropertyType == typeof(string) &&
+                            pi.GetValue(testEnvironment) is string propertyValue &&
+                            propertyValue.StartsWith(term, StringComparison.Ordinal))
+                        {
+                            string lookupKey = propertyValue.AsSpan(term.Length).ToString();
+
+                            if (testConfiguration.SharedKeyValuePairs.TryGetValue(lookupKey, out string? sharedValue))
+                            {
+                                pi.SetValue(testEnvironment, sharedValue);
+                            }
+                        }
                     }
                 }
+
+                environments.Add(testEnvironment);
             }
 
             if (environments.Count == 0)
diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs b/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs
@@ -15,7 +15,6 @@
 * limitations under the License.
 */
 
-using System.Collections.Generic;
 using System.Text.Json.Serialization;
 
 namespace Apache.Arrow.Adbc.Tests.Drivers.BigQuery
@@ -41,6 +40,9 @@ public BigQueryTestEnvironment()
         [JsonPropertyName("projectId")]
         public string? ProjectId { get; set; }
 
+        [JsonPropertyName("billingProjectId")]
+        public string? BillingProjectId { get; set; }
+
         [JsonPropertyName("clientId")]
         public string ClientId { get; set; } = string.Empty;
 
@@ -68,9 +70,30 @@ public BigQueryTestEnvironment()
         [JsonPropertyName("includePublicProjectId")]
         public bool IncludePublicProjectId { get; set; } = false;
 
+        /// <summary>
+        /// Sets the query timeout (in minutes).
+        /// </summary>
         [JsonPropertyName("timeoutMinutes")]
         public int? TimeoutMinutes { get; set; }
 
+        /// <summary>
+        /// Sets the query timeout (in seconds).
+        /// </summary>
+        [JsonPropertyName("queryTimeout")]
+        public int? QueryTimeout { get; set; }
+
+        /// <summary>
+        /// The number of seconds to allow for the HttpClient timeout.
+        /// </summary>
+        [JsonPropertyName("clientTimeout")]
+        public int? ClientTimeout { get; set; }
+
+        /// <summary>
+        /// Indicates if timeout tests should run during this execution.
+        /// </summary>
+        [JsonPropertyName("runTimeoutTests")]
+        public bool RunTimeoutTests { get; set; } = false;
+
         [JsonPropertyName("maxStreamCount")]
         public int? MaxStreamCount { get; set; }
     }
diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs b/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs
@@ -58,6 +58,11 @@ internal static Dictionary<string, string> GetBigQueryParameters(BigQueryTestEnv
                 parameters.Add(BigQueryParameters.ProjectId, testEnvironment.ProjectId!);
             }
 
+            if (!string.IsNullOrEmpty(testEnvironment.BillingProjectId))
+            {
+                parameters.Add(BigQueryParameters.BillingProjectId, testEnvironment.BillingProjectId!);
+            }
+
             if (!string.IsNullOrEmpty(testEnvironment.JsonCredential))
             {
                 parameters.Add(BigQueryParameters.AuthenticationType, BigQueryConstants.ServiceAccountAuthenticationType);
@@ -92,7 +97,17 @@ internal static Dictionary<string, string> GetBigQueryParameters(BigQueryTestEnv
 
             if (testEnvironment.TimeoutMinutes.HasValue)
             {
-                parameters.Add(BigQueryParameters.GetQueryResultsOptionsTimeoutMinutes, testEnvironment.TimeoutMinutes.Value.ToString());
+                int seconds = testEnvironment.TimeoutMinutes.Value * 60;
+                parameters.Add(BigQueryParameters.GetQueryResultsOptionsTimeout, seconds.ToString());
+            }
+            else if (testEnvironment.QueryTimeout.HasValue)
+            {
+                parameters.Add(BigQueryParameters.GetQueryResultsOptionsTimeout, testEnvironment.QueryTimeout.Value.ToString());
+            }
+
+            if (testEnvironment.ClientTimeout.HasValue)
+            {
+                parameters.Add(BigQueryParameters.ClientTimeout, testEnvironment.ClientTimeout.Value.ToString());
             }
 
             if (testEnvironment.MaxStreamCount.HasValue)
diff --git a/csharp/test/Drivers/BigQuery/DriverTests.cs b/csharp/test/Drivers/BigQuery/DriverTests.cs
@@ -18,6 +18,7 @@
 using System;
 using System.Collections.Generic;
 using System.Linq;
+using System.Threading.Tasks;
 using Apache.Arrow.Adbc.Drivers.BigQuery;
 using Apache.Arrow.Adbc.Tests.Metadata;
 using Apache.Arrow.Adbc.Tests.Xunit;
@@ -307,5 +308,43 @@ private AdbcConnection GetAdbcConnection(string? environmentName)
 
             return _configuredConnections[environmentName!];
         }
+
+        /// <summary>
+        /// Validates the ClientTimeout parameter.
+        /// </summary>
+        [SkippableFact, Order(7)]
+        public void ClientTimeoutTest()
+        {
+            foreach (BigQueryTestEnvironment environment in _environments)
+            {
+                if (environment.RunTimeoutTests && environment.ClientTimeout.HasValue)
+                {
+                    AdbcConnection adbcConnection = BigQueryTestingUtils.GetBigQueryAdbcConnection(environment);
+
+                    AdbcStatement statement = adbcConnection.CreateStatement();
+                    statement.SqlQuery = environment.Query;
+
+                    Assert.Throws<TaskCanceledException>(() => { statement.ExecuteQuery(); });
+                }
+            }
+        }
+
+        /// <summary>
+        /// Validates the GetQueryResultsOptionsTimeoutMinutes parameter.
+        /// </summary>
+        [SkippableFact, Order(8)]
+        public void QueryTimeoutTest()
+        {
+            foreach (BigQueryTestEnvironment environment in _environments)
+            {
+                if (environment.RunTimeoutTests && (environment.QueryTimeout.HasValue || environment.TimeoutMinutes.HasValue))
+                {
+                    AdbcConnection adbcConnection = BigQueryTestingUtils.GetBigQueryAdbcConnection(environment);
+                    AdbcStatement statement = adbcConnection.CreateStatement();
+                    statement.SqlQuery = environment.Query;
+                    Assert.Throws<TimeoutException>(() => { statement.ExecuteQuery(); });
+                }
+            }
+        }
     }
 }
diff --git a/csharp/test/Drivers/BigQuery/Resources/bigqueryconfig.json b/csharp/test/Drivers/BigQuery/Resources/bigqueryconfig.json
@@ -1,10 +1,14 @@
 {
     "testEnvironments": [ "", "" ],
+    "shared": {
+        "clientId": "clientId_example",
+        "clientSecret": "clientSecret_example"
+    }
     "environments": {
         "<env1>": {
             "projectId": "",
-            "clientId": "",
-            "clientSecret": "",
+            "clientId": "$ref:shared.clientId",
+            "clientSecret": "$ref:shared.clientSecret",
             "refreshToken": "",
             "maxStreamCount": 1,
             "metadata": {
@@ -17,6 +21,21 @@
             "expectedResults": 0
         },
         "<env2>": {
+            "projectId": "",
+            "clientId": "$ref:shared.clientId",
+            "clientSecret": "$ref:shared.clientSecret",
+            "refreshToken": "",
+            "maxStreamCount": 1,
+            "metadata": {
+                "catalog": "",
+                "schema": "",
+                "table": "",
+                "expectedColumnCount": 0
+            },
+            "query": "",
+            "expectedResults": 0
+        },
+        "<env3>": {
             "projectId": "",
             "clientId": "",
             "clientSecret": "",
diff --git a/csharp/test/Drivers/BigQuery/readme.md b/csharp/test/Drivers/BigQuery/readme.md

Original file line number	Diff line number	Diff line change
`@@ -18,6 +18,7 @@`
`18`	`18`	`using System;`
`19`	`19`	`using System.Collections.Generic;`
`20`	`20`	`using System.IO;`
	`21`	`+using System.Reflection;`
`21`	`22`	`using System.Text.Json;`
`22`	`23`
`23`	`24`	`namespace Apache.Arrow.Adbc.Tests`
`@@ -38,7 +39,6 @@ public static T LoadMultiEnvironmentTestConfiguration<T>(string environmentVaria`
`38`	`39`	`{`
`39`	`40`	`// use a JSON file for the various settings`
`40`	`41`	`string json = File.ReadAllText(environmentValue);`
`41`		`-`
`42`	`42`	`testConfiguration = JsonSerializer.Deserialize<T>(json)!;`
`43`	`43`	`}`
`44`	`44`	`}`
`@@ -60,17 +60,34 @@ public static List<TEnvironment> GetTestEnvironments<TEnvironment>(MultiEnvironm`
`60`	`60`	`throw new InvalidOperationException("There are no environments configured");`
`61`	`61`
`62`	`62`	`List<TEnvironment> environments = new List<TEnvironment>();`
	`63`	`+ string term = "$ref:shared.";`
`63`	`64`
`64`	`65`	`foreach (string environmentName in GetEnvironmentNames(testConfiguration.TestEnvironmentNames))`
`65`	`66`	`{`
`66`		`- if (testConfiguration.Environments.TryGetValue(environmentName, out TEnvironment? testEnvironment))`
	`67`	`+ if (!testConfiguration.Environments.TryGetValue(environmentName, out TEnvironment? testEnvironment) \|\| testEnvironment is null)`
	`68`	`+ continue;`
	`69`	`+`
	`70`	`+ testEnvironment.Name = environmentName;`
	`71`	`+`
	`72`	`+ if (testConfiguration.SharedKeyValuePairs.Count > 0)`
`67`	`73`	`{`
`68`		`- if (testEnvironment != null)`
	`74`	`+ foreach (PropertyInfo pi in testEnvironment.GetType().GetProperties())`
`69`	`75`	`{`
`70`		`- testEnvironment.Name = environmentName;`
`71`		`- environments.Add(testEnvironment);`
	`76`	`+ if (pi.PropertyType == typeof(string) &&`
	`77`	`+ pi.GetValue(testEnvironment) is string propertyValue &&`
	`78`	`+ propertyValue.StartsWith(term, StringComparison.Ordinal))`
	`79`	`+ {`
	`80`	`+ string lookupKey = propertyValue.AsSpan(term.Length).ToString();`
	`81`	`+`
	`82`	`+ if (testConfiguration.SharedKeyValuePairs.TryGetValue(lookupKey, out string? sharedValue))`
	`83`	`+ {`
	`84`	`+ pi.SetValue(testEnvironment, sharedValue);`
	`85`	`+ }`
	`86`	`+ }`
`72`	`87`	`}`
`73`	`88`	`}`
	`89`	`+`
	`90`	`+ environments.Add(testEnvironment);`
`74`	`91`	`}`
`75`	`92`
`76`	`93`	`if (environments.Count == 0)`