feat: Batch DML support (#370)

* feat: Batch DML support

Add support for Batch DML in read/write transactions. Batch DML reduces the number of round-trips to Spanner when executing multiple writes. The big advantage of using Batch DML over mutations, is that Batch DML does support read-your-writes after the batch has been written.

Author: aakashanandg

examples/snippets/batch-dml/README.md

@@ -0,0 +1,13 @@
+# Sample - Batch DML
+
+This example shows how to use [Batch DML](https://cloud.google.com/spanner/docs/dml-tasks#use-batch)
+with the Spanner ActiveRecord adapter.
+
+The sample will automatically start a Spanner Emulator in a docker container and execute the sample
+against that emulator. The emulator will automatically be stopped when the application finishes.
+
+Run the application with the command
+
+```bash
+bundle exec rake run
+```

examples/snippets/batch-dml/Rakefile

@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+require_relative "../config/environment"
+require "sinatra/activerecord/rake"
+
+desc "Sample showing how to use Batch DML with the Spanner ActiveRecord provider."
+task :run do
+  Dir.chdir("..") { sh "bundle exec rake run[batch-dml]" }
+end

examples/snippets/batch-dml/application.rb

@@ -0,0 +1,54 @@
+# Copyright 2025 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+require "io/console"
+require_relative "../config/environment"
+require_relative "models/singer"
+require_relative "models/album"
+
+class Application
+  def self.run
+    first_names = %w[Pete Alice John Ethel Trudy Naomi Wendy Ruben Thomas Elly]
+    last_names = %w[Wendelson Allison Peterson Johnson Henderson Ericsson Aronson Tennet Courtou]
+
+    # Insert 5 new singers using Batch DML.
+    ActiveRecord::Base.transaction do
+      # The Base.dml_batch function starts a DML batch. All DML statements that are
+      # generated by ActiveRecord inside the block that is given will be added to
+      # the current batch. The batch is executed at the end of the block. The data
+      # that has been written is readable after the block ends.
+      ActiveRecord::Base.dml_batch do
+        5.times do
+          Singer.create first_name: first_names.sample, last_name: last_names.sample
+        end
+      end
+      # Data that has been inserted/update using Batch DML can be read in the same
+      # transaction as the one that added/updated the data. This is different from
+      # mutations, as mutations do not support read-your-writes.
+      singers = Singer.all
+      puts "Inserted #{singers.count} singers in one batch"
+    end
+
+    # Batch DML can also be used to update existing data.
+    ActiveRecord::Base.transaction do
+      # Start a DML batch.
+      singers = nil
+      ActiveRecord::Base.dml_batch do
+        # Queries can be executed inside a DML batch.
+        # These are executed directly and do not affect the DML batch.
+        singers = Singer.all
+        singers.each do |singer|
+          singer.picture = Base64.encode64 SecureRandom.alphanumeric(SecureRandom.random_number(10..200))
+          singer.save
+        end
+      end
+      puts "Updated #{singers.count} singers in one batch"
+    end
+  end
+end
+
+
+Application.run

examples/snippets/batch-dml/config/database.yml

@@ -0,0 +1,11 @@
+# Batch DML Example Configuration File
+# This file is used to configure the database connection for the batch DML example.
+development:
+  adapter: spanner
+  emulator_host: localhost:9010
+  project: test-project
+  instance: test-instance
+  database: testdb
+  pool: 5
+  timeout: 5000
+  schema_dump: false

examples/snippets/batch-dml/db/migrate/01_create_tables.rb

@@ -0,0 +1,23 @@
+# Copyright 2025 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+class CreateTables < ActiveRecord::Migration[6.0]
+  def change
+    connection.ddl_batch do
+      create_table :singers do |t|
+        t.string :first_name
+        t.string :last_name
+        t.binary :picture
+      end
+
+      create_table :albums do |t|
+        t.string :title
+        t.numeric :marketing_budget
+        t.references :singer, index: false, foreign_key: true
+      end
+    end
+  end
+end

examples/snippets/batch-dml/db/seeds.rb

@@ -0,0 +1,8 @@
+# Copyright 2025 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+# This file is intentionally kept empty, as this sample
+# does not need any initial data.

examples/snippets/batch-dml/models/album.rb

@@ -0,0 +1,9 @@
+# Copyright 2025 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+class Album < ActiveRecord::Base
+  belongs_to :singer
+end

examples/snippets/batch-dml/models/singer.rb

@@ -0,0 +1,9 @@
+# Copyright 2025 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+class Singer < ActiveRecord::Base
+  has_many :albums
+end

lib/active_record/connection_adapters/spanner/database_statements.rb

@@ -25,9 +25,13 @@ def execute sql, name = nil, binds = []
 
         def internal_exec_query sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false
           result = internal_execute sql, name, binds, prepare: prepare, async: async, allow_retry: allow_retry
-          ActiveRecord::Result.new(
-            result.fields.keys.map(&:to_s), result.rows.map(&:values)
-          )
+          if result
+            ActiveRecord::Result.new(
+              result.fields.keys.map(&:to_s), result.rows.map(&:values)
+            )
+          else
+            ActiveRecord::Result.new [], []
+          end
         end
 
         def internal_execute sql, name = "SQL", binds = [],
@@ -77,14 +81,16 @@ def execute_query_or_dml statement_type, sql, name, binds
                   @connection.execute_query sql,
                                             params: params,
                                             types: types,
-                                            request_options: request_options
+                                            request_options: request_options,
+                                            statement_type: statement_type
                 end
               else
                 @connection.execute_query sql,
                                           params: params,
                                           types: types,
                                           single_use_selector: selector,
-                                          request_options: request_options
+                                          request_options: request_options,
+                                          statement_type: statement_type
               end
             end
           end
@@ -150,9 +156,13 @@ def query sql, name = nil
 
           def exec_query sql, name = "SQL", binds = [], prepare: false # rubocop:disable Lint/UnusedMethodArgument
             result = execute sql, name, binds
-            ActiveRecord::Result.new(
-              result.fields.keys.map(&:to_s), result.rows.map(&:values)
-            )
+            if result.respond_to? :fields
+              ActiveRecord::Result.new(
+                result.fields.keys.map(&:to_s), result.rows.map(&:values)
+              )
+            else
+              ActiveRecord::Result.new [], []
+            end
           end
 
           def sql_for_insert sql, pk, binds
@@ -198,6 +208,12 @@ def update arel, name = nil, binds = []
         alias delete update
 
         def exec_update sql, name = "SQL", binds = []
+          # Check if a DML batch is active on the connection.
+          if @connection.dml_batch?
+            # This call buffers the SQL.
+            execute sql, name, binds
+            return
+          end
           result = execute sql, name, binds
           # Make sure that we consume the entire result stream before trying to get the stats.
           # This is required because the ExecuteStreamingSql RPC is also used for (Partitioned) DML,
@@ -237,7 +253,7 @@ def execute_ddl statements
 
         # Transaction
 
-        def transaction requires_new: nil, isolation: nil, joinable: true, **kwargs, &block # rubocop:disable Metrics/PerceivedComplexity
+        def transaction requires_new: nil, isolation: nil, joinable: true, **kwargs, &block # rubocop:disable Metrics/PerceivedComplexity,Metrics/CyclomaticComplexity
           commit_options = kwargs.delete :commit_options
           exclude_from_streams = kwargs.delete :exclude_txn_from_change_streams
           @_spanner_begin_transaction_options = {
@@ -270,6 +286,9 @@ def transaction requires_new: nil, isolation: nil, joinable: true, **kwargs, &bl
             else
               raise
             end
+          rescue Google::Cloud::AbortedError => err
+            sleep(delay_from_aborted(err) || backoff *= 1.3)
+            retry
           ensure
             # Clean up the instance variable to avoid leaking options.
             @_spanner_begin_transaction_options = nil

lib/active_record/connection_adapters/spanner_adapter.rb

@@ -150,7 +150,8 @@ def spanner_schema_cache
       end
 
       # Spanner Connection API
-      delegate :ddl_batch, :ddl_batch?, :start_batch_ddl, :abort_batch, :run_batch,
+      delegate :dml_batch, :dml_batch?, :start_batch_dml,
+               :ddl_batch, :ddl_batch?, :start_batch_ddl, :abort_batch, :run_batch,
                :isolation_level, :isolation_level=, to: :@connection
 
       def current_spanner_transaction