diff --git a/csharp/src/AdbcDrivers.BigQuery/BigQueryStatement.cs b/csharp/src/AdbcDrivers.BigQuery/BigQueryStatement.cs index ac95fc4..5fc1c22 100644 --- a/csharp/src/AdbcDrivers.BigQuery/BigQueryStatement.cs +++ b/csharp/src/AdbcDrivers.BigQuery/BigQueryStatement.cs @@ -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; @@ -844,12 +879,13 @@ private async Task 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); } @@ -881,20 +917,35 @@ private async Task ExecuteUpdateInternalAsync() }; } - Task getQueryResultsAsyncFunc() + Task 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 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); diff --git a/csharp/test/AdbcDrivers.BigQuery.MockServer/BigQueryMockServer.cs b/csharp/test/AdbcDrivers.BigQuery.MockServer/BigQueryMockServer.cs index c47c92a..b87e5a4 100644 --- a/csharp/test/AdbcDrivers.BigQuery.MockServer/BigQueryMockServer.cs +++ b/csharp/test/AdbcDrivers.BigQuery.MockServer/BigQueryMockServer.cs @@ -46,6 +46,7 @@ public sealed class BigQueryMockServer : IDisposable private readonly ConcurrentDictionary _tables = new(); private readonly ConcurrentDictionary _sessions = new(); private readonly ConcurrentQueue _executedQueries = new(); + private int _queryResultsRequestCount; /// /// The REST API endpoint as host:port (e.g., "127.0.0.1:12345"). @@ -64,6 +65,11 @@ public sealed class BigQueryMockServer : IDisposable /// public IReadOnlyList ExecutedQueries => _executedQueries.ToArray(); + /// + /// The number of requests made to the query-results endpoint. + /// + public int QueryResultsRequestCount => _queryResultsRequestCount; + /// /// The mock gRPC service for configuring Storage Read API responses. /// @@ -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 @@ -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; @@ -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, } @@ -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; } } } } diff --git a/csharp/test/AdbcDrivers.BigQuery.Tests/MockServer/MockServerTests.cs b/csharp/test/AdbcDrivers.BigQuery.Tests/MockServer/MockServerTests.cs index e28c1c4..2f1ab53 100644 --- a/csharp/test/AdbcDrivers.BigQuery.Tests/MockServer/MockServerTests.cs +++ b/csharp/test/AdbcDrivers.BigQuery.Tests/MockServer/MockServerTests.cs @@ -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; @@ -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 + { + { 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()); + 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 + { + { 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()); + using AdbcStatement statement = connection.CreateStatement(); + + AdbcException exception = Assert.Throws( + () => 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 + { + { 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()); + 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 + { + { 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()); + using AdbcStatement statement = connection.CreateStatement(); + Assert.Throws(() => 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 + { + { 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()); + 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 + { + { 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()); + 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); + } } }