Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(csharp/src/Drivers/BigQuery): fix support for large results #1507

Merged
merged 7 commits into from
Feb 4, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 28 additions & 6 deletions csharp/src/Drivers/BigQuery/BigQueryConnection.cs
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
*/

using System;
using System.Collections;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Diagnostics;
Expand Down Expand Up @@ -270,6 +271,23 @@
return new BigQueryInfoArrowStream(StandardSchemas.GetObjectsSchema, dataArrays);
}

/// <summary>
/// Executes the query using the BigQueryClient.
/// </summary>
/// <param name="sql">The query to execute.</param>
/// <param name="parameters">Parameters to include.</param>
/// <param name="queryOptions">Additional query options.</param>
/// <param name="resultsOptions">Additional result options.</param>
/// <returns></returns>
/// <remarks>
/// Can later add logging or metrics around query calls.
/// </remarks>
private BigQueryResults? ExecuteQuery(string sql, IEnumerable<BigQueryParameter>? parameters, QueryOptions? queryOptions = null, GetQueryResultsOptions? resultsOptions = null)
{
BigQueryResults? result = this.client?.ExecuteQuery(sql, parameters, queryOptions, resultsOptions);
return result;
}

private List<IArrowArray> GetCatalogs(
GetObjectsDepth depth,
string catalogPattern,
Expand Down Expand Up @@ -405,7 +423,7 @@
}
}

BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
BigQueryResults? result = ExecuteQuery(query, parameters: null);

if (result != null)
{
Expand Down Expand Up @@ -481,7 +499,7 @@
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)
{
Expand Down Expand Up @@ -570,7 +588,7 @@
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)
{
Expand Down Expand Up @@ -631,7 +649,7 @@

StringArray.Builder constraintColumnNamesBuilder = new StringArray.Builder();

BigQueryResults? result = this.client?.ExecuteQuery(query, parameters: null);
BigQueryResults? result = ExecuteQuery(query, parameters: null);

if (result != null)
{
Expand Down Expand Up @@ -661,7 +679,7 @@
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)
{
Expand Down Expand Up @@ -788,7 +806,7 @@
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<Field> fields = new List<Field>();

Expand Down Expand Up @@ -991,7 +1009,7 @@
Open();
}

BigQueryStatement statement = new BigQueryStatement(this.client, this.credential);

Check warning on line 1012 in csharp/src/Drivers/BigQuery/BigQueryConnection.cs

View workflow job for this annotation

GitHub Actions / C# ubuntu-latest

Possible null reference argument for parameter 'client' in 'BigQueryStatement.BigQueryStatement(BigQueryClient client, GoogleCredential credential)'.

Check warning on line 1012 in csharp/src/Drivers/BigQuery/BigQueryConnection.cs

View workflow job for this annotation

GitHub Actions / C# ubuntu-latest

Possible null reference argument for parameter 'client' in 'BigQueryStatement.BigQueryStatement(BigQueryClient client, GoogleCredential credential)'.

Check warning on line 1012 in csharp/src/Drivers/BigQuery/BigQueryConnection.cs

View workflow job for this annotation

GitHub Actions / C# windows-2019

Possible null reference argument for parameter 'client' in 'BigQueryStatement.BigQueryStatement(BigQueryClient client, GoogleCredential credential)'.

Check warning on line 1012 in csharp/src/Drivers/BigQuery/BigQueryConnection.cs

View workflow job for this annotation

GitHub Actions / C# windows-2019

Possible null reference argument for parameter 'client' in 'BigQueryStatement.BigQueryStatement(BigQueryClient client, GoogleCredential credential)'.

Check warning on line 1012 in csharp/src/Drivers/BigQuery/BigQueryConnection.cs

View workflow job for this annotation

GitHub Actions / C# macos-latest

Possible null reference argument for parameter 'client' in 'BigQueryStatement.BigQueryStatement(BigQueryClient client, GoogleCredential credential)'.

Check warning on line 1012 in csharp/src/Drivers/BigQuery/BigQueryConnection.cs

View workflow job for this annotation

GitHub Actions / C# macos-latest

Possible null reference argument for parameter 'client' in 'BigQueryStatement.BigQueryStatement(BigQueryClient client, GoogleCredential credential)'.
statement.Options = ParseOptions();
return statement;
}
Expand All @@ -1014,6 +1032,10 @@
{
options[keyValuePair.Key] = keyValuePair.Value;
}
if (keyValuePair.Key == BigQueryParameters.LargeResultsDestinationTable)
{
options[keyValuePair.Key] = keyValuePair.Value;
}
}

return new ReadOnlyDictionary<string, string>(options);
Expand Down
1 change: 1 addition & 0 deletions csharp/src/Drivers/BigQuery/BigQueryParameters.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand Down
31 changes: 31 additions & 0 deletions csharp/src/Drivers/BigQuery/BigQueryStatement.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -185,6 +186,36 @@
{
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;
Expand Down Expand Up @@ -341,7 +372,7 @@
{
if (this.readers == null)
{
return null;

Check warning on line 375 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# ubuntu-latest

Possible null reference return.

Check warning on line 375 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# ubuntu-latest

Possible null reference return.

Check warning on line 375 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# windows-2019

Possible null reference return.

Check warning on line 375 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# windows-2019

Possible null reference return.

Check warning on line 375 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# macos-latest

Possible null reference return.

Check warning on line 375 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# macos-latest

Possible null reference return.
}

while (true)
Expand All @@ -351,7 +382,7 @@
if (!this.readers.MoveNext())
{
Dispose(); // TODO: Remove this line
return null;

Check warning on line 385 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# ubuntu-latest

Possible null reference return.

Check warning on line 385 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# ubuntu-latest

Possible null reference return.

Check warning on line 385 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# windows-2019

Possible null reference return.

Check warning on line 385 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# windows-2019

Possible null reference return.

Check warning on line 385 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# macos-latest

Possible null reference return.

Check warning on line 385 in csharp/src/Drivers/BigQuery/BigQueryStatement.cs

View workflow job for this annotation

GitHub Actions / C# macos-latest

Possible null reference return.
}
this.reader = this.readers.Current;
}
Expand Down
4 changes: 4 additions & 0 deletions csharp/src/Drivers/BigQuery/readme.md
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,10 @@ https://cloud.google.com/dotnet/docs/reference/Google.Cloud.BigQuery.V2/latest/G
**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.large_results_destination_table**<br>
&nbsp;&nbsp;&nbsp;&nbsp;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**<br>
&nbsp;&nbsp;&nbsp;&nbsp;The [Project ID](https://cloud.google.com/resource-manager/docs/creating-managing-projects) used for accessing BigQuery.

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,6 @@
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\..\src\Apache.Arrow.Adbc\Apache.Arrow.Adbc.csproj" />
<ProjectReference Include="..\..\..\src\arrow\csharp\src\Apache.Arrow\Apache.Arrow.csproj" />
<ProjectReference Include="..\..\..\src\Client\Apache.Arrow.Adbc.Client.csproj" />
<ProjectReference Include="..\..\..\src\Drivers\BigQuery\Apache.Arrow.Adbc.Drivers.BigQuery.csproj" />
<ProjectReference Include="..\..\Apache.Arrow.Adbc.Tests\Apache.Arrow.Adbc.Tests.csproj" />
Expand Down
10 changes: 10 additions & 0 deletions csharp/test/Drivers/BigQuery/BigQueryTestConfiguration.cs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@ namespace Apache.Arrow.Adbc.Tests.Drivers.BigQuery
/// </summary>
internal class BigQueryTestConfiguration : TestConfiguration
{
public BigQueryTestConfiguration()
{
AllowLargeResults = false;
}

[JsonPropertyName("projectId")]
public string ProjectId { get; set; }

Expand All @@ -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; }
}
}
10 changes: 10 additions & 0 deletions csharp/test/Drivers/BigQuery/BigQueryTestingUtils.cs
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,16 @@ internal static Dictionary<string,string> 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;
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
</ItemGroup>
<ItemGroup>
<ProjectReference Include="..\..\..\src\Apache.Arrow.Adbc\Apache.Arrow.Adbc.csproj" />
<ProjectReference Include="..\..\..\src\arrow\csharp\src\Apache.Arrow\Apache.Arrow.csproj" />
<ProjectReference Include="..\..\..\src\Client\Apache.Arrow.Adbc.Client.csproj" />
<ProjectReference Include="..\..\..\src\Drivers\FlightSql\Apache.Arrow.Adbc.Drivers.FlightSql.csproj" />
<ProjectReference Include="..\..\Apache.Arrow.Adbc.Tests\Apache.Arrow.Adbc.Tests.csproj" />
Expand Down
Loading