Skip to content
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
73 changes: 62 additions & 11 deletions csharp/src/AdbcDrivers.BigQuery/BigQueryStatement.cs
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,41 @@ public override void SetOption(string key, string value)

switch (key)
{
case AdbcOptions.Ingest.TargetCatalog:
_isBulkIngest = true;
_ingestTargetCatalog = value;
break;
case AdbcOptions.Ingest.TargetDbSchema:
_isBulkIngest = true;
_ingestTargetDbSchema = value;
break;
case AdbcOptions.Ingest.TargetTable:
_isBulkIngest = true;
_ingestTargetTable = value;
break;
case AdbcOptions.Ingest.Mode:
BulkIngestMode ingestMode = value switch
{
AdbcOptions.IngestMode.Create => BulkIngestMode.Create,
AdbcOptions.IngestMode.Append => BulkIngestMode.Append,
AdbcOptions.IngestMode.Replace => BulkIngestMode.Replace,
AdbcOptions.IngestMode.CreateAppend => BulkIngestMode.CreateAppend,
_ => throw new AdbcException($"Unsupported bulk ingest mode: {value}", AdbcStatusCode.InvalidArgument),
};
_isBulkIngest = true;
_ingestMode = ingestMode;
break;
case AdbcOptions.Ingest.Temporary:
switch (value)
{
case AdbcOptions.Enabled:
throw AdbcException.NotImplemented("Temporary table bulk ingest is not supported for BigQuery");
case AdbcOptions.Disabled:
break;
default:
throw new AdbcException($"Unsupported value for {AdbcOptions.Ingest.Temporary}: {value}", AdbcStatusCode.InvalidArgument);
}
break;
case AdbcOptions.Telemetry.TraceParent:
SetTraceParent(string.IsNullOrWhiteSpace(value) ? null : value);
break;
Expand Down Expand Up @@ -844,12 +879,13 @@ private async Task<UpdateResult> ExecuteUpdateInternalAsync()

return await this.TraceActivityAsync(async activity =>
{
GetQueryResultsOptions getQueryResultsOptions = new GetQueryResultsOptions();

TimeSpan? queryResultsTimeout = GetEffectiveQueryResultsTimeout();
TimeSpan pollTimeout = queryResultsTimeout ?? TimeSpan.FromSeconds(BigQueryConstants.DefaultQueryResultsTimeoutSeconds);
PollSettings pollSettings = new PollSettings(
Expiration.FromTimeout(pollTimeout),
TimeSpan.FromSeconds(1));
if (queryResultsTimeout.HasValue)
{
getQueryResultsOptions.Timeout = queryResultsTimeout.Value;
activity?.AddBigQueryParameterTag(BigQueryParameters.GetQueryResultsOptionsTimeout, (int)queryResultsTimeout.Value.TotalSeconds);
}

Expand Down Expand Up @@ -881,20 +917,35 @@ private async Task<UpdateResult> ExecuteUpdateInternalAsync()
};
}

Task<BigQueryResults> getQueryResultsAsyncFunc()
Task<BigQueryJob> createJobAsyncFunc()
{
return ExecuteCancellableJobAsync(context, activity, async (context, jobActivity) =>
{
context.Job = await this.Client.CreateQueryJobAsync(SqlQuery, null, updateQueryOptions, context.CancellationToken).ConfigureAwait(false);
jobActivity?.AddEvent("getqueryresultsasync_started", [new("job.id", context.Job.Reference.JobId)]);
BigQueryResults results = await context.Job.GetQueryResultsAsync(getQueryResultsOptions, context.CancellationToken).ConfigureAwait(false);
jobActivity?.AddEvent("getqueryresultsasync_completed", GetJobStatistics(jobActivity, context.Job));
return context.Job;
}, ClassName + "." + nameof(ExecuteUpdateInternalAsync) + "." + nameof(BigQueryClient.CreateQueryJobAsync));
}

return results;
}, ClassName + "." + nameof(ExecuteUpdateInternalAsync) + "." + nameof(BigQueryJob.GetQueryResultsAsync));
BigQueryJob job = await ExecuteWithRetriesAsync(createJobAsyncFunc, activity, context.CancellationToken).ConfigureAwait(false);
context.Job = job;

Task<BigQueryJob> pollJobAsyncFunc()
{
return ExecuteCancellableJobAsync(context, activity, async (context, jobActivity) =>
{
context.Job = await Client.GetJobAsync(job.Reference, cancellationToken: context.CancellationToken).ConfigureAwait(false);
jobActivity?.AddEvent("polluntilcompletedasync_started", [new("job.id", context.Job.Reference.JobId)]);
context.Job = await context.Job.PollUntilCompletedAsync(
pollSettings: pollSettings,
cancellationToken: context.CancellationToken).ConfigureAwait(false);
context.Job.ThrowOnFatalError();
jobActivity?.AddEvent("polluntilcompletedasync_completed", GetJobStatistics(jobActivity, context.Job));

return context.Job;
}, ClassName + "." + nameof(ExecuteUpdateInternalAsync) + "." + nameof(BigQueryJob.PollUntilCompletedAsync));
}
BigQueryResults? result = await ExecuteWithRetriesAsync(getQueryResultsAsyncFunc, activity, context.CancellationToken);
long updatedRows = result?.NumDmlAffectedRows.HasValue == true ? result.NumDmlAffectedRows.Value : -1L;
BigQueryJob completedJob = await ExecuteWithRetriesAsync(pollJobAsyncFunc, activity, context.CancellationToken).ConfigureAwait(false);
long updatedRows = completedJob.Resource.Statistics?.Query?.NumDmlAffectedRows ?? -1L;

activity?.AddTag(SemanticConventions.Db.Response.ReturnedRows, updatedRows);
return new UpdateResult(updatedRows);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@ public sealed class BigQueryMockServer : IDisposable
private readonly ConcurrentDictionary<string, Table> _tables = new();
private readonly ConcurrentDictionary<string, bool> _sessions = new();
private readonly ConcurrentQueue<string> _executedQueries = new();
private int _queryResultsRequestCount;

/// <summary>
/// The REST API endpoint as host:port (e.g., "127.0.0.1:12345").
Expand All @@ -64,6 +65,11 @@ public sealed class BigQueryMockServer : IDisposable
/// </summary>
public IReadOnlyList<string> ExecutedQueries => _executedQueries.ToArray();

/// <summary>
/// The number of requests made to the query-results endpoint.
/// </summary>
public int QueryResultsRequestCount => _queryResultsRequestCount;

/// <summary>
/// The mock gRPC service for configuring Storage Read API responses.
/// </summary>
Expand Down Expand Up @@ -192,6 +198,8 @@ private void MapRestRoutes(WebApplication app)
JobId = jobId,
ProjectId = projectId,
Status = "DONE",
StatementType = queryText?.TrimStart().StartsWith("UPDATE", StringComparison.OrdinalIgnoreCase) == true ? "UPDATE" : "SELECT",
NumDmlAffectedRows = queryText?.TrimStart().StartsWith("UPDATE", StringComparison.OrdinalIgnoreCase) == true ? 2 : null,
};

// If CreateSession is requested, generate a new session ID
Expand Down Expand Up @@ -232,6 +240,7 @@ private void MapRestRoutes(WebApplication app)
// GET /bigquery/v2/projects/{projectId}/queries/{jobId} - Get query results
app.MapGet("/bigquery/v2/projects/{projectId}/queries/{jobId}", async (HttpContext ctx, string projectId, string jobId) =>
{
Interlocked.Increment(ref _queryResultsRequestCount);
if (!_jobs.TryGetValue(jobId, out var mockJob))
{
ctx.Response.StatusCode = 404;
Expand Down Expand Up @@ -354,7 +363,8 @@ private static Job CreateJobResource(MockJob mockJob)
EndTime = now,
Query = new JobStatistics2
{
StatementType = "SELECT",
StatementType = mockJob.StatementType,
NumDmlAffectedRows = mockJob.NumDmlAffectedRows,
TotalBytesProcessed = 0,
TotalBytesBilled = 0,
}
Expand Down Expand Up @@ -419,6 +429,8 @@ private class MockJob
public string ProjectId { get; set; } = string.Empty;
public string Status { get; set; } = "DONE";
public string? SessionId { get; set; }
public string StatementType { get; set; } = "SELECT";
public long? NumDmlAffectedRows { get; set; }
}
}
}
191 changes: 191 additions & 0 deletions csharp/test/AdbcDrivers.BigQuery.Tests/MockServer/MockServerTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
#if NET8_0_OR_GREATER

using System.Collections.Generic;
using System.Linq;
using Apache.Arrow;
using Apache.Arrow.Adbc;
using Apache.Arrow.Types;
Expand Down Expand Up @@ -158,6 +159,196 @@ public async System.Threading.Tasks.Task CanBulkIngestAppendToTable()
// Verify the table was created in the REST API
// (CreateAppend mode should create it since it didn't exist)
}

[Theory]
[InlineData(AdbcOptions.IngestMode.Create, false)]
[InlineData(AdbcOptions.IngestMode.Append, true)]
[InlineData(AdbcOptions.IngestMode.Replace, true)]
[InlineData(AdbcOptions.IngestMode.CreateAppend, false)]
public void CanBulkIngestThroughStatementOptions(string mode, bool createTableFirst)
{
using var mockServer = new BigQueryMockServer();

const string projectId = "mock-project";
const string datasetId = "test_dataset";
string tableId = $"option_ingest_{mode.Substring(mode.LastIndexOf('.') + 1)}";
var parameters = new Dictionary<string, string>
{
{ BigQueryParameters.ProjectId, projectId },
{ BigQueryParameters.AuthenticationType, BigQueryConstants.MockAuthenticationType },
{ BigQueryParameters.TestRestEndpoint, mockServer.RestEndpoint },
{ BigQueryParameters.TestStorageEndpoint, mockServer.GrpcEndpoint },
};

using var driver = new BigQueryDriver();
using AdbcDatabase database = driver.Open(parameters);
using AdbcConnection connection = database.Connect(new Dictionary<string, string>());
using RecordBatch batch = CreateBatch();
var existingStreamNames = mockServer.WriteService.Streams.Keys.ToHashSet();

if (createTableFirst)
{
using AdbcStatement create = connection.BulkIngest(projectId, datasetId, tableId, BulkIngestMode.Create, false);
create.Bind(batch, batch.Schema);
create.ExecuteUpdate();
existingStreamNames = mockServer.WriteService.Streams.Keys.ToHashSet();
}

using AdbcStatement statement = connection.CreateStatement();
statement.SetOption(AdbcOptions.Ingest.TargetCatalog, projectId);
statement.SetOption(AdbcOptions.Ingest.TargetDbSchema, datasetId);
statement.SetOption(AdbcOptions.Ingest.TargetTable, tableId);
statement.SetOption(AdbcOptions.Ingest.Temporary, AdbcOptions.Disabled);
statement.SetOption(AdbcOptions.Ingest.Mode, mode);
statement.Bind(batch, batch.Schema);

UpdateResult result = statement.ExecuteUpdate();

Assert.Equal(3, result.AffectedRows);
Assert.Empty(mockServer.ExecutedQueries);
string writeStreamName = Assert.Single(mockServer.WriteService.Streams.Keys.Except(existingStreamNames));
var writeStream = mockServer.WriteService.Streams[writeStreamName];
Assert.True(writeStream.Finalized);
Assert.Single(writeStream.RecordBatches);
}

[Fact]
public void BulkIngestThroughStatementOptionsRejectsTemporaryTable()
{
using var mockServer = new BigQueryMockServer();
var parameters = new Dictionary<string, string>
{
{ BigQueryParameters.ProjectId, "mock-project" },
{ BigQueryParameters.AuthenticationType, BigQueryConstants.MockAuthenticationType },
{ BigQueryParameters.TestRestEndpoint, mockServer.RestEndpoint },
{ BigQueryParameters.TestStorageEndpoint, mockServer.GrpcEndpoint },
};

using var driver = new BigQueryDriver();
using AdbcDatabase database = driver.Open(parameters);
using AdbcConnection connection = database.Connect(new Dictionary<string, string>());
using AdbcStatement statement = connection.CreateStatement();

AdbcException exception = Assert.Throws<AdbcException>(
() => statement.SetOption(AdbcOptions.Ingest.Temporary, AdbcOptions.Enabled));

Assert.Equal(AdbcStatusCode.NotImplemented, exception.Status);
}

[Fact]
public void DisabledTemporaryOptionDoesNotChangeExecutionMode()
{
using var mockServer = new BigQueryMockServer();
var parameters = new Dictionary<string, string>
{
{ BigQueryParameters.ProjectId, "mock-project" },
{ BigQueryParameters.AuthenticationType, BigQueryConstants.MockAuthenticationType },
{ BigQueryParameters.TestRestEndpoint, mockServer.RestEndpoint },
{ BigQueryParameters.TestStorageEndpoint, mockServer.GrpcEndpoint },
};

using var driver = new BigQueryDriver();
using AdbcDatabase database = driver.Open(parameters);
using AdbcConnection connection = database.Connect(new Dictionary<string, string>());
using AdbcStatement statement = connection.CreateStatement();
statement.SetOption(AdbcOptions.Ingest.Temporary, AdbcOptions.Disabled);
statement.SqlQuery = "UPDATE test_table SET value = 1";

UpdateResult result = statement.ExecuteUpdate();

Assert.Equal(2, result.AffectedRows);
}

[Fact]
public void InvalidIngestModeDoesNotChangeExecutionMode()
{
using var mockServer = new BigQueryMockServer();
var parameters = new Dictionary<string, string>
{
{ BigQueryParameters.ProjectId, "mock-project" },
{ BigQueryParameters.AuthenticationType, BigQueryConstants.MockAuthenticationType },
{ BigQueryParameters.TestRestEndpoint, mockServer.RestEndpoint },
{ BigQueryParameters.TestStorageEndpoint, mockServer.GrpcEndpoint },
};

using var driver = new BigQueryDriver();
using AdbcDatabase database = driver.Open(parameters);
using AdbcConnection connection = database.Connect(new Dictionary<string, string>());
using AdbcStatement statement = connection.CreateStatement();
Assert.Throws<AdbcException>(() => statement.SetOption(AdbcOptions.Ingest.Mode, "invalid"));
statement.SqlQuery = "UPDATE test_table SET value = 1";

UpdateResult result = statement.ExecuteUpdate();

Assert.Equal(2, result.AffectedRows);
}

[Fact]
public void DropTableExecuteUpdateDoesNotRequestQueryResults()
{
using var mockServer = new BigQueryMockServer();
var parameters = new Dictionary<string, string>
{
{ BigQueryParameters.ProjectId, "mock-project" },
{ BigQueryParameters.AuthenticationType, BigQueryConstants.MockAuthenticationType },
{ BigQueryParameters.TestRestEndpoint, mockServer.RestEndpoint },
{ BigQueryParameters.TestStorageEndpoint, mockServer.GrpcEndpoint },
};

using var driver = new BigQueryDriver();
using AdbcDatabase database = driver.Open(parameters);
using AdbcConnection connection = database.Connect(new Dictionary<string, string>());
using AdbcStatement statement = connection.CreateStatement();
statement.SqlQuery = "DROP TABLE IF EXISTS `mock-project.test_dataset.test_table`";

UpdateResult result = statement.ExecuteUpdate();

Assert.Equal(-1, result.AffectedRows);
Assert.Equal(0, mockServer.QueryResultsRequestCount);
Assert.Single(mockServer.ExecutedQueries);
}

[Fact]
public void DmlExecuteUpdateReturnsAffectedRowsWithoutRequestingQueryResults()
{
using var mockServer = new BigQueryMockServer();
var parameters = new Dictionary<string, string>
{
{ BigQueryParameters.ProjectId, "mock-project" },
{ BigQueryParameters.AuthenticationType, BigQueryConstants.MockAuthenticationType },
{ BigQueryParameters.TestRestEndpoint, mockServer.RestEndpoint },
{ BigQueryParameters.TestStorageEndpoint, mockServer.GrpcEndpoint },
};

using var driver = new BigQueryDriver();
using AdbcDatabase database = driver.Open(parameters);
using AdbcConnection connection = database.Connect(new Dictionary<string, string>());
using AdbcStatement statement = connection.CreateStatement();
statement.SqlQuery = "UPDATE test_table SET value = 1";

UpdateResult result = statement.ExecuteUpdate();

Assert.Equal(2, result.AffectedRows);
Assert.Equal(0, mockServer.QueryResultsRequestCount);
}

private static RecordBatch CreateBatch()
{
var schema = new Schema(new[]
{
new Field("id", Int64Type.Default, nullable: false),
new Field("name", StringType.Default, nullable: true),
}, null);

return new RecordBatch(
schema,
new IArrowArray[]
{
new Int64Array.Builder().Append(1).Append(2).Append(3).Build(),
new StringArray.Builder().Append("Alice").Append("Bob").Append("Charlie").Build(),
},
3);
}
}
}

Expand Down