DEMO: Add SpannerBatchWriteCommand to the ADO Layer

apis/Google.Cloud.Spanner.Data/Google.Cloud.Spanner.Data.IntegrationTests/SpannerBatchWriteCommandTests.cs

@@ -0,0 +1,128 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Google.Cloud.Spanner.V1;
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Google.Cloud.Spanner.Data.IntegrationTests;
+
+[Collection(nameof(BatchWriteTableFixture))]
+public class SpannerBatchWriteCommandTests
+{
+    private readonly BatchWriteTableFixture _fixture;
+
+    public SpannerBatchWriteCommandTests(BatchWriteTableFixture fixture) =>
+        _fixture = fixture;
+
+    [Fact]
+    public async Task ExecuteAsync_Success()
+    {
+        using var connection = _fixture.GetConnection();
+        var command = connection.CreateBatchWriteCommand();
+        var key1 = Guid.NewGuid().ToString();
+        var key2 = Guid.NewGuid().ToString();
+
+        var insertCommand1 = connection.CreateInsertCommand(_fixture.TableName, new SpannerParameterCollection { { _fixture.KeyName, SpannerDbType.String, key1 } });
+        var insertCommand2 = connection.CreateInsertCommand(_fixture.TableName, new SpannerParameterCollection { { _fixture.KeyName, SpannerDbType.String, key2 } });
+
+        command.Add(new[] { insertCommand1, insertCommand2 });
+
+        var responses = await command.ExecuteAsync().ToListAsync();
+
+        Assert.Single(responses);
+        var response = responses[0];
+        Assert.Equal(0, response.Status.Code); // OK
+        Assert.Single(response.Indexes);
+        Assert.Equal(0, response.Indexes[0]);
+
+        using var snapshot = await connection.BeginTransactionAsync();
+        var readCommand = connection.CreateReadCommand(_fixture.TableName, ReadOptions.FromColumns(_fixture.KeyName), KeySet.All);
+        readCommand.Transaction = snapshot;
+        using var reader = await readCommand.ExecuteReaderAsync();
+        var keys = new HashSet<string>();
+        while(await reader.ReadAsync())
+        {
+            keys.Add(reader.GetString(0));
+        }
+        Assert.Contains(key1, keys);
+        Assert.Contains(key2, keys);
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_PartialFailure()
+    {
+        using var connection = _fixture.GetConnection();
+        var command = connection.CreateBatchWriteCommand();
+        var conflictKey = Guid.NewGuid().ToString();
+        var successKey = Guid.NewGuid().ToString();
+
+        var insertConflict1 = connection.CreateInsertCommand(_fixture.TableName, new SpannerParameterCollection { { _fixture.KeyName, SpannerDbType.String, conflictKey } });
+        var insertConflict2 = connection.CreateInsertCommand(_fixture.TableName, new SpannerParameterCollection { { _fixture.KeyName, SpannerDbType.String, conflictKey } });
+        var insertSuccess = connection.CreateInsertCommand(_fixture.TableName, new SpannerParameterCollection { { _fixture.KeyName, SpannerDbType.String, successKey } });
+
+        // This group will fail due to a primary key conflict.
+        command.Add(new[] { insertConflict1, insertConflict2 });
+        // This group will succeed.
+        command.Add(insertSuccess);
+
+        var responses = await command.ExecuteAsync().ToListAsync();
+
+        Assert.Equal(2, responses.Count);
+        var failedResponse = responses.Single(r => r.Status.Code != 0);
+        var successResponse = responses.Single(r => r.Status.Code == 0);
+
+        Assert.Single(failedResponse.Indexes);
+        Assert.Equal(0, failedResponse.Indexes[0]);
+        Assert.Single(successResponse.Indexes);
+        Assert.Equal(1, successResponse.Indexes[0]);
+
+        using var snapshot = await connection.BeginTransactionAsync();
+        var readCommand = connection.CreateReadCommand(_fixture.TableName, ReadOptions.FromColumns(_fixture.KeyName), KeySet.All);
+        readCommand.Transaction = snapshot;
+        using var reader = await readCommand.ExecuteReaderAsync();
+        var keys = new HashSet<string>();
+        while (await reader.ReadAsync())
+        {
+            keys.Add(reader.GetString(0));
+        }
+        Assert.Contains(successKey, keys);
+        Assert.DoesNotContain(conflictKey, keys);
+    }
+
+    [Fact]
+    public async Task ExecuteAsync_EmptyCommand()
+    {
+        using var connection = _fixture.GetConnection();
+        var command = connection.CreateBatchWriteCommand();
+
+        var responses = await command.ExecuteAsync().ToListAsync();
+
+        Assert.Empty(responses);
+    }
+
+    [Fact]
+    public void Add_InvalidCommandType_Throws()
+    {
+        using var connection = _fixture.GetConnection();
+        var command = connection.CreateBatchWriteCommand();
+        var selectCommand = connection.CreateSelectCommand($"SELECT * FROM {_fixture.TableName}");
+
+        Assert.Throws<ArgumentOutOfRangeException>(() => command.Add(selectCommand));
+    }
+}

apis/Google.Cloud.Spanner.Data/Google.Cloud.Spanner.Data.IntegrationTests/V1/BatchWriteTableFixture.cs

@@ -0,0 +1,83 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Google.Api.Gax.Grpc;
+using Google.Cloud.Spanner.V1;
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+using Xunit;
+
+namespace Google.Cloud.Spanner.Data.IntegrationTests;
+
+[CollectionDefinition(nameof(BatchWriteTableFixture))]
+public class BatchWriteTableFixture : CommonDataTableFixture, ICollectionFixture<BatchWriteTableFixture>, IAsyncLifetime
+{
+    private const int TimeoutSeconds = 60;
+    private SessionPoolManager _poolManager;
+    private SessionPool _pool;
+
+    /// <summary>
+    /// The name of the key column in the table.
+    /// </summary>
+    public readonly string KeyName = "Key";
+
+    public BatchWriteTableFixture() : base("BatchWrite")
+    {
+    }
+
+    /// <inheritdoc/>
+    protected override void CreateTable()
+    {
+        ExecuteDdl($@"CREATE TABLE {this.TableName} (
+                     {this.KeyName} STRING(256),
+                   ) PRIMARY KEY ({this.KeyName})");
+    }
+
+    /// <inheritdoc/>
+    protected override void PopulateTable(bool fresh)
+    {
+        // Do nothing. The BatchWriteTests don't need pre-populated data.
+    }
+
+    /// <summary>
+    /// Initializes the session pool for the fixture.
+    /// </summary>
+    public async Task InitializeAsync()
+    {
+        _poolManager = SessionPoolManager.Create(new SessionPoolOptions());
+        _pool = await _poolManager.AcquireSessionPoolAsync(SpannerClientCreationOptions);
+    }
+
+
+    public Task<PooledSession> GetPooledSessionAsync() =>
+        _pool.AcquireSessionAsync(DatabaseName, new TransactionOptions(), CancellationToken.None);
+
+    /// <summary>
+    /// Gets the call settings used to interact with the database.
+    /// </summary>
+    public CallSettings GetCallSettings()
+    {
+        using var cts = new CancellationTokenSource(TimeSpan.FromSeconds(TimeoutSeconds));
+        using SpannerConnection connection = GetConnection();
+        return connection.CreateCallSettings(settings => settings.BatchWriteSettings, TimeoutSeconds, cts.Token);
+    }
+
+    /// <inheritdoc/>
+    public Task DisposeAsync()
+    {
+        _poolManager.Release(_pool);
+        return Task.CompletedTask;
+    }
+}

apis/Google.Cloud.Spanner.Data/Google.Cloud.Spanner.Data.IntegrationTests/V1/BatchWriteTests.cs

@@ -0,0 +1,136 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Google.Cloud.Spanner.Data.IntegrationTests;
+using Google.Protobuf.WellKnownTypes;
+using System.Threading.Tasks;
+using Xunit;
+using System.Collections.Generic;
+using System;
+using Newtonsoft.Json;
+
+namespace Google.Cloud.Spanner.V1.IntegrationTests;
+
+[CommonTestDiagnostics]
+[Collection(nameof(BatchWriteTableFixture))]
+public class BatchWriteTests
+{
+    private record BatchWriteResponseAnalysis(int successCount, int failureCount);
+
+    private readonly BatchWriteTableFixture _fixture;
+    private const int TimeoutSeconds = 60;
+
+    public BatchWriteTests(BatchWriteTableFixture fixture) =>
+            _fixture = fixture;
+
+    [Fact]
+    public async Task BatchWrite_Success()
+    {
+        // Create two non-conflicting write mutations
+        PooledSession pooledSession = await _fixture.GetPooledSessionAsync();
+        var mutationGroup = new BatchWriteRequest.Types.MutationGroup();
+        mutationGroup.Mutations.AddRange([
+        new Mutation
+        {
+            Insert = new Mutation.Types.Write
+            {
+                Table = _fixture.TableName,
+                Columns = { _fixture.KeyName },
+                Values = { new ListValue { Values = { Value.ForString(Guid.NewGuid().ToString()) } } }
+            }
+        },
+        new Mutation
+        {
+            Insert = new Mutation.Types.Write
+            {
+                Table = _fixture.TableName,
+                Columns = { _fixture.KeyName },
+                Values = { new ListValue { Values = { Value.ForString(Guid.NewGuid().ToString()) } } }
+            }
+        }
+        ]);
+
+        BatchWriteRequest batchWriteRequest = new()
+        {
+            Session = pooledSession.Session.Name,
+            MutationGroups = { mutationGroup }
+        };
+
+        IAsyncEnumerable<BatchWriteResponse> responseStream = pooledSession.BatchWriteAsync(batchWriteRequest, _fixture.GetCallSettings());
+        BatchWriteResponseAnalysis responseAnalysis = await GetResponseAnalysis(responseStream);
+
+        // The single mutation group will succeed as one atomic unit
+        Assert.Equal(0, responseAnalysis.failureCount);
+        Assert.Equal(1, responseAnalysis.successCount);
+    }
+
+    [Fact]
+    public async Task BatchWrite_Failure_Conflict()
+    {
+        // Create a valid and a conflicting write mutation group.
+        PooledSession pooledSession = await _fixture.GetPooledSessionAsync();
+        var mutation =  new Mutation
+        {
+            Insert = new Mutation.Types.Write
+            {
+                Table = _fixture.TableName,
+                Columns = { _fixture.KeyName },
+                Values = { new ListValue { Values = { Value.ForString("Conflict because matching write.") } } },
+            }
+        };
+
+        var mutationGroupNoConflict = new BatchWriteRequest.Types.MutationGroup();
+        var mutationGroupConflict = new BatchWriteRequest.Types.MutationGroup();
+        mutationGroupConflict.Mutations.AddRange([mutation, mutation]);
+        mutationGroupNoConflict.Mutations.AddRange([mutation]);
+
+        BatchWriteRequest batchWriteRequest = new()
+        {
+            Session = pooledSession.Session.Name,
+            MutationGroups = { mutationGroupConflict, mutationGroupNoConflict }
+        };
+
+        // Send the batch write request and parse the response
+        IAsyncEnumerable<BatchWriteResponse> responseStream = pooledSession.BatchWriteAsync(batchWriteRequest, _fixture.GetCallSettings());
+        BatchWriteResponseAnalysis responseAnalysis = await GetResponseAnalysis(responseStream);
+
+        // The mutation group with a conflict will result in one failure
+        Assert.Equal(1, responseAnalysis.failureCount);
+        // The mutation group without a conflict will result in one success
+        Assert.Equal(1, responseAnalysis.successCount);
+    }
+
+    private static async Task<BatchWriteResponseAnalysis> GetResponseAnalysis(IAsyncEnumerable<BatchWriteResponse> responseStream)
+    {
+        var successCount = 0;
+        var failureCount = 0;
+        string result = "";
+        await foreach(var response in responseStream)
+        {
+            // STATUS CODE of 0 is a success
+            if (response.Status.Code == 0)
+            {
+                successCount += response.Indexes.Count;
+            }
+            else
+            {
+                failureCount += response.Indexes.Count;
+            }
+            result += JsonConvert.SerializeObject(response);
+        }
+
+        return new BatchWriteResponseAnalysis(successCount: successCount, failureCount: failureCount);
+    }
+
+}