feat(Spanner.V1): Add support for BatchWrite.

Author: robertvoinescu-work

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 using Google.Api.Gax;
+using Google.Api.Gax.Grpc;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -49,6 +50,28 @@ public async Task<ManagedTransaction> BeginTransactionAsync(TransactionOptions t
         return ManagedTransaction.FromTransactionOptions(_options, session, transactionOptions, singleUse, shared);
     }
 
+	/// <summary>
+	/// Executes a BatchWrite RPC asynchronously, returning a stream of responses.
+	/// </summary>
+    /// <remarks>
+    /// This method modifies the <paramref name="request"/> to include the session name from this managed session.
+    /// </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. When the task completes, the result is the response stream of <see cref="BatchWriteResponse"/> objects.</returns>
+	public async Task<AsyncResponseStream<BatchWriteResponse>> BatchWriteAsync(BatchWriteRequest request, CallSettings callSettings)
+	{
+		GaxPreconditions.CheckNotNull(request, nameof(request));
+        CancellationToken cancellationToken = callSettings.CancellationToken ?? default;
+        var session = await _lifecycleManager.GetFreshSessionAsync(cancellationToken).ConfigureAwait(false);
+
+        // Populate the request with the current session name.
+        request.Session = session.SessionName.ToString();
+
+        var response = _options.Client.BatchWrite(request, callSettings);
+        return response.GetResponseStream();
+    }
+
     /// <summary>
     /// Ensures the underlying Spanner session is fresh.
     /// </summary>

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

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 using Google.Api.Gax;
+using Google.Api.Gax.Grpc;
 using Google.Cloud.Spanner.Common.V1;
 using Google.Cloud.Spanner.V1.Internal.Logging;
 using Google.Protobuf;
@@ -118,6 +119,23 @@ public SessionPoolSegmentStatistics GetSegmentStatisticsSnapshot(SessionPoolSegm
         public SessionPoolSegmentStatistics GetSegmentStatisticsSnapshot(DatabaseName databaseName) =>
             GetSegmentStatisticsSnapshot(SessionPoolSegmentKey.Create(databaseName));
 
+        /// <summary>
+        /// Executes a BatchWrite RPC asynchronously, returning a stream of responses.
+        /// </summary>
+        /// <param name="request">The batch write request. Must not be null.</param>
+        /// <param name="key">The session pool segment key. 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. When the task completes, the result is the response stream of <see cref="BatchWriteResponse"/> objects.</returns>
+        [Obsolete("Use ManagedSession.BatchWriteAsync instead.")]
+        public async Task<AsyncResponseStream<BatchWriteResponse>> BatchWriteAsync(BatchWriteRequest request, SessionPoolSegmentKey key, CallSettings callSettings)
+        {
+            GaxPreconditions.CheckNotNull(request, nameof(request));
+            GaxPreconditions.CheckNotNull(key, nameof(key));
+
+            var managedSession = _managedSessions.GetOrAdd(key, CreateManagedSession);
+            return await managedSession.BatchWriteAsync(request, callSettings).ConfigureAwait(false);
+        }
+
         /// <summary>
         /// Asynchronously acquires a session that will handle transaction creation as needed.
         /// This is equivalent to calling <see cref="AcquireSessionAsync(SessionPoolSegmentKey, TransactionOptions, CancellationToken)"/>
@@ -308,4 +326,4 @@ public Task ShutdownPoolAsync(SessionPoolSegmentKey key, CancellationToken cance
                 .WithDatabaseRole(key.DatabaseRole)
                 .WithTimeout(Options.Timeout));
     }
-}
+}

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.