diff --git a/csharp/src/Drivers/BigQuery/BigQueryConnection.cs b/csharp/src/Drivers/BigQuery/BigQueryConnection.cs index ef68ef4e34..a06b72eb0e 100644 --- a/csharp/src/Drivers/BigQuery/BigQueryConnection.cs +++ b/csharp/src/Drivers/BigQuery/BigQueryConnection.cs @@ -16,6 +16,7 @@ */ using System; +using System.Collections; using System.Collections.Generic; using System.Collections.ObjectModel; using System.Diagnostics; @@ -270,6 +271,23 @@ public override IArrowArrayStream GetObjects( return new BigQueryInfoArrowStream(StandardSchemas.GetObjectsSchema, dataArrays); } + /// + /// Executes the query using the BigQueryClient. + /// + /// The query to execute. + /// Parameters to include. + /// Additional query options. + /// Additional result options. + /// + /// + /// Can later add logging or metrics around query calls. + /// + private BigQueryResults? ExecuteQuery(string sql, IEnumerable? parameters, QueryOptions? queryOptions = null, GetQueryResultsOptions? resultsOptions = null) + { + BigQueryResults? result = this.client?.ExecuteQuery(sql, parameters, queryOptions, resultsOptions); + return result; + } + private List GetCatalogs( GetObjectsDepth depth, string catalogPattern, @@ -405,7 +423,7 @@ private StructArray GetTableSchemas( } } - BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null); + BigQueryResults? result = ExecuteQuery(query, parameters: null); if (result != null) { @@ -481,7 +499,7 @@ private StructArray GetColumnSchema( query = string.Concat(query, string.Format("AND column_name LIKE '{0}'", Sanitize(columnNamePattern))); } - BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null); + BigQueryResults? result = ExecuteQuery(query, parameters: null); if (result != null) { @@ -570,7 +588,7 @@ private StructArray GetConstraintSchema( string query = string.Format("SELECT * FROM `{0}`.`{1}`.INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE table_name = '{2}'", Sanitize(catalog), Sanitize(dbSchema), Sanitize(table)); - BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null); + BigQueryResults? result = ExecuteQuery(query, parameters: null); if (result != null) { @@ -631,7 +649,7 @@ private StringArray GetConstraintColumnNames( StringArray.Builder constraintColumnNamesBuilder = new StringArray.Builder(); - BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null); + BigQueryResults? result = ExecuteQuery(query, parameters: null); if (result != null) { @@ -661,7 +679,7 @@ private StructArray GetConstraintsUsage( string query = string.Format("SELECT * FROM `{0}`.`{1}`.INFORMATION_SCHEMA.CONSTRAINT_COLUMN_USAGE WHERE constraint_name = '{2}'", Sanitize(catalog), Sanitize(dbSchema), Sanitize(constraintName)); - BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null); + BigQueryResults? result = ExecuteQuery(query, parameters: null); if (result != null) { @@ -788,7 +806,7 @@ public override Schema GetTableSchema(string catalog, string dbSchema, string ta string query = string.Format("SELECT * FROM `{0}`.`{1}`.INFORMATION_SCHEMA.COLUMNS WHERE table_name = '{2}'", Sanitize(catalog), Sanitize(dbSchema), Sanitize(tableName)); - BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null); + BigQueryResults? result = ExecuteQuery(query, parameters: null); List fields = new List(); @@ -1014,6 +1032,10 @@ private IReadOnlyDictionary ParseOptions() { options[keyValuePair.Key] = keyValuePair.Value; } + if (keyValuePair.Key == BigQueryParameters.LargeResultsDestinationTable) + { + options[keyValuePair.Key] = keyValuePair.Value; + } } return new ReadOnlyDictionary(options); diff --git a/csharp/src/Drivers/BigQuery/BigQueryParameters.cs b/csharp/src/Drivers/BigQuery/BigQueryParameters.cs index ff05036dc4..ff57ea8faa 100644 --- a/csharp/src/Drivers/BigQuery/BigQueryParameters.cs +++ b/csharp/src/Drivers/BigQuery/BigQueryParameters.cs @@ -29,6 +29,7 @@ public class BigQueryParameters public const string AuthenticationType = "adbc.bigquery.auth_type"; public const string JsonCredential = "adbc.bigquery.auth_json_credential"; public const string AllowLargeResults = "adbc.bigquery.allow_large_results"; + public const string LargeResultsDestinationTable = "adbc.bigquery.large_results_destination_table"; public const string UseLegacySQL = "adbc.bigquery.use_legacy_sql"; public const string LargeDecimalsAsString = "adbc.bigquery.large_decimals_as_string"; public const string Scopes = "adbc.bigquery.scopes"; diff --git a/csharp/src/Drivers/BigQuery/BigQueryStatement.cs b/csharp/src/Drivers/BigQuery/BigQueryStatement.cs index 389e462e55..f4f21740a6 100644 --- a/csharp/src/Drivers/BigQuery/BigQueryStatement.cs +++ b/csharp/src/Drivers/BigQuery/BigQueryStatement.cs @@ -27,6 +27,7 @@ using Apache.Arrow.Ipc; using Apache.Arrow.Types; using Google.Apis.Auth.OAuth2; +using Google.Apis.Bigquery.v2.Data; using Google.Cloud.BigQuery.Storage.V1; using Google.Cloud.BigQuery.V2; using TableFieldSchema = Google.Apis.Bigquery.v2.Data.TableFieldSchema; @@ -185,6 +186,36 @@ static IArrowReader ReadChunk(BigQueryReadClient readClient, string streamName) { options.AllowLargeResults = true ? keyValuePair.Value.ToLower().Equals("true") : false; } + if (keyValuePair.Key == BigQueryParameters.LargeResultsDestinationTable) + { + string destinationTable = keyValuePair.Value; + + if (!destinationTable.Contains(".")) + throw new InvalidOperationException($"{BigQueryParameters.LargeResultsDestinationTable} is invalid"); + + string projectId = string.Empty; + string datasetId = string.Empty; + string tableId = string.Empty; + + string[] segments = destinationTable.Split('.'); + + if(segments.Length != 3) + throw new InvalidOperationException($"{BigQueryParameters.LargeResultsDestinationTable} cannot be parsed"); + + projectId = segments[0]; + datasetId = segments[1]; + tableId = segments[2]; + + if(string.IsNullOrEmpty(projectId.Trim()) || string.IsNullOrEmpty(datasetId.Trim()) || string.IsNullOrEmpty(tableId.Trim())) + throw new InvalidOperationException($"{BigQueryParameters.LargeResultsDestinationTable} contains invalid values"); + + options.DestinationTable = new TableReference() + { + ProjectId = projectId, + DatasetId = datasetId, + TableId = tableId + }; + } if (keyValuePair.Key == BigQueryParameters.UseLegacySQL) { options.UseLegacySql = true ? keyValuePair.Value.ToLower().Equals("true") : false; diff --git a/csharp/src/Drivers/BigQuery/readme.md b/csharp/src/Drivers/BigQuery/readme.md index 611374c585..847513c485 100644 --- a/csharp/src/Drivers/BigQuery/readme.md +++ b/csharp/src/Drivers/BigQuery/readme.md @@ -49,6 +49,10 @@ https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/G **adbc.bigquery.auth_json_credential**
    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.large_results_destination_table**
+    Sets the [DestinationTable](https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/Google.Cloud.BigQuery.V2.QueryOptions#Google_Cloud_BigQuery_V2_QueryOptions_DestinationTable) value of the QueryOptions if configured. Expects the format to be `{projectId}.{datasetId}.{tableId}` to set the corresponding values in the [TableReference](https://github.com/googleapis/google-api-dotnet-client/blob/6c415c73788b848711e47c6dd33c2f93c76faf97/Src/Generated/Google.Apis.Bigquery.v2/Google.Apis.Bigquery.v2.cs#L9348) class. + + **adbc.bigquery.project_id**
    The [Project ID](https://cloud.google.com/resource-manager/docs/creating-managing-projects) used for accessing BigQuery. diff --git a/csharp/test/Drivers/BigQuery/Apache.Arrow.Adbc.Tests.Drivers.BigQuery.csproj b/csharp/test/Drivers/BigQuery/Apache.Arrow.Adbc.Tests.Drivers.BigQuery.csproj index 74712cdc2f..e2d7921490 100644 --- a/csharp/test/Drivers/BigQuery/Apache.Arrow.Adbc.Tests.Drivers.BigQuery.csproj +++ b/csharp/test/Drivers/BigQuery/Apache.Arrow.Adbc.Tests.Drivers.BigQuery.csproj @@ -13,7 +13,6 @@ - diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs b/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs index a6b0ab6b07..a7a5622929 100644 --- a/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs +++ b/csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs @@ -24,6 +24,11 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.BigQuery /// internal class BigQueryTestConfiguration : TestConfiguration { + public BigQueryTestConfiguration() + { + AllowLargeResults = false; + } + [JsonPropertyName("projectId")] public string ProjectId { get; set; } @@ -42,5 +47,10 @@ internal class BigQueryTestConfiguration : TestConfiguration [JsonPropertyName("jsonCredential")] public string JsonCredential { get; set; } + [JsonPropertyName("allowLargeResults")] + public bool AllowLargeResults { get; set; } + + [JsonPropertyName("largeResultsDestinationTable")] + public string LargeResultsDestinationTable { get; set; } } } diff --git a/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs b/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs index 874ccac791..ffca233cc5 100644 --- a/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs +++ b/csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs @@ -74,6 +74,16 @@ internal static Dictionary GetBigQueryParameters(BigQueryTestConf parameters.Add(BigQueryParameters.Scopes, testConfiguration.Scopes); } + if(testConfiguration.AllowLargeResults) + { + parameters.Add(BigQueryParameters.AllowLargeResults, testConfiguration.AllowLargeResults.ToString()); + } + + if(!string.IsNullOrEmpty(testConfiguration.LargeResultsDestinationTable)) + { + parameters.Add(BigQueryParameters.LargeResultsDestinationTable, testConfiguration.LargeResultsDestinationTable); + } + return parameters; } diff --git a/csharp/test/Drivers/FlightSql/Apache.Arrow.Adbc.Tests.Drivers.FlightSql.csproj b/csharp/test/Drivers/FlightSql/Apache.Arrow.Adbc.Tests.Drivers.FlightSql.csproj index 9036f5acfd..7a45c99b00 100644 --- a/csharp/test/Drivers/FlightSql/Apache.Arrow.Adbc.Tests.Drivers.FlightSql.csproj +++ b/csharp/test/Drivers/FlightSql/Apache.Arrow.Adbc.Tests.Drivers.FlightSql.csproj @@ -15,7 +15,6 @@ -