feat: Spanner.V1 BatchWrite

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

@@ -0,0 +1,82 @@
+// 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(TransactionOptions transactionOptions = default) =>
+        _pool.AcquireSessionAsync(DatabaseName, transactionOptions ?? 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,167 @@
+// 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);
+    }
+
+    [Fact]
+    public async Task BatchWrite_Failure_TransactionModeNotNone()
+    {
+        // Any transaction associated with the session will cause the batchwrite to throw an exception
+        var options = new TransactionOptions()
+        {
+            ReadWrite = new TransactionOptions.Types.ReadWrite()
+        };
+        PooledSession pooledSession = await _fixture.GetPooledSessionAsync(options);
+
+        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()) } } }
+            }
+        }]);
+
+        BatchWriteRequest batchWriteRequest = new()
+        {
+            Session = pooledSession.Session.Name,
+            MutationGroups = { mutationGroup }
+        };
+
+        Assert.Throws<InvalidOperationException>(() =>
+                pooledSession.BatchWriteAsync(batchWriteRequest, _fixture.GetCallSettings()));
+    }
+
+    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);
+    }
+}

apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/Internal/ExecuteHelper.cs

@@ -141,4 +141,4 @@ async Task<TResult> ImplAsync()
             }
         }
     }
-}
+}

apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/PooledSession.cs

@@ -18,6 +18,7 @@
 using Google.Protobuf;
 using Google.Protobuf.WellKnownTypes;
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -578,7 +579,7 @@
                 skipTransactionCreation: request.Mutations.Count == 0, // If there are only mutations we won't have a transaction but we need one.
                 callSettings?.CancellationToken ?? default);
 
-            void SetCommandTransaction(TransactionSelector transactionSelector) 
+            void SetCommandTransaction(TransactionSelector transactionSelector)
             {
                 switch (transactionSelector.SelectorCase)
                 {
@@ -850,6 +851,44 @@
             Transaction GetInlinedTransaction(ExecuteBatchDmlResponse response) => response?.ResultSets?.FirstOrDefault()?.Metadata?.Transaction;
         }
 
+        /// <summary>
+        /// Executes a BatchWrite RPC asynchronously, returning a stream of responses.
+        /// </summary>
+        /// <remarks>
+        /// This operation can only be run when the session is configured to use no transaction.
+        /// </remarks>
+        /// <param name="request">The batch write request. Must not be null.</param>
+        /// <param name="callSettings">If not null, applies overrides to this RPC call.</param>
+        /// <returns>A task representing the asynchronous operation, with a result of the stream of <see cref="BatchWriteResponse"/> objects.</returns>
+        public Task<AsyncResponseStream<BatchWriteResponse>> BatchWriteAsync(BatchWriteRequest request, CallSettings callSettings)
+        {
+            CheckNotDisposed();
+            GaxPreconditions.CheckNotNull(request, nameof(request));
+
+            request.SessionAsSessionName = SessionName;
+
+            return ExecuteMaybeWithTransactionSelectorAsync(
+                transactionSelectorSetter: SetCommandTransaction,
+                commandAsync: async () =>
+                {
+                    var call = Client.BatchWrite(request, callSettings);
+                    var stream = call.GetResponseStream();
+
+                    // Update refresh time before returning the stream to ensure the session isn't reclaimed prematurely.
+                    UpdateRefreshTime();
+
+                    return stream;
+                },
+                inlinedTransactionExtractor: null,
+                skipTransactionCreation: true,
+                cancellationToken: callSettings?.CancellationToken ?? default);
+
+            void SetCommandTransaction(TransactionSelector transactionSelector) =>
+                    throw new InvalidOperationException($"{nameof(BatchWriteAsync)} BatchWrite does not support explicitly set transactions, it won't be executed within an active transaction."
+                    + $"The sessions most be acquired with {nameof(TransactionOptions.ModeCase)} = {nameof(ModeOneofCase.None)} or {nameof(BatchWriteAsync)} should complete execution before any other commands that"
+                    + "would create a transaction are executed.");
+        }
+
         private void MaybeApplyDirectedReadOptions(IReadOrQueryRequest request)
         {
             if (TransactionMode == ModeOneofCase.ReadOnly // Directed reads apply only to single use or read only transactions. Single use are read only.
@@ -876,6 +915,22 @@
             UpdateRefreshTime();
         }
 
+        private async IAsyncEnumerable<TResponse> StreamResponsesAndRecordSuccessAsync<TResponse>(AsyncResponseStream<TResponse> responseStream)
+        {
+            try
+            {
+                while (await responseStream.MoveNextAsync().WithSessionExpiryChecking(Session).ConfigureAwait(false))
+                {
+                    UpdateRefreshTime();
+                    yield return responseStream.Current;
+                }
+            }
+            finally
+            {
+                await responseStream.DisposeAsync().ConfigureAwait(false);
+            }
+        }
+
         /// <summary>
         /// Updates the refresh time for the session based on the current time. This should
         /// be called when a successful RPC is made with the associated session.

apis/Google.Cloud.Spanner.V1/Google.Cloud.Spanner.V1/SpannerClientPartial.cs

@@ -166,6 +166,13 @@ partial void Modify_PartitionReadRequest(ref PartitionReadRequest request, ref C
             ApplyRequestIdHeader(ref settings);
         }
 
+        partial void Modify_BatchWriteRequest(ref BatchWriteRequest request, ref CallSettings settings)
+        {
+            ApplyResourcePrefixHeaderFromSession(ref settings, request.Session);
+            MaybeApplyRouteToLeaderHeader(ref settings);
+            ApplyRequestIdHeader(ref settings);
+        }
+
         internal static void ApplyResourcePrefixHeaderFromDatabase(ref CallSettings settings, string resource)
         {
             // If we haven't been given a resource name, just leave the request as it is.
@@ -203,4 +210,4 @@ internal void ApplyRequestIdHeader(ref CallSettings settings)
             settings = settings.MergedWith(newSettings);
         }
     }
-}
+}