diff --git a/examples/snippets/batch-dml/README.md b/examples/snippets/batch-dml/README.md new file mode 100644 index 00000000..ade98932 --- /dev/null +++ b/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 +``` \ No newline at end of file diff --git a/examples/snippets/batch-dml/Rakefile b/examples/snippets/batch-dml/Rakefile new file mode 100644 index 00000000..1242a4dd --- /dev/null +++ b/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 diff --git a/examples/snippets/batch-dml/application.rb b/examples/snippets/batch-dml/application.rb new file mode 100644 index 00000000..1f2754b4 --- /dev/null +++ b/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 diff --git a/examples/snippets/batch-dml/config/database.yml b/examples/snippets/batch-dml/config/database.yml new file mode 100644 index 00000000..84bb3096 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/examples/snippets/batch-dml/db/migrate/01_create_tables.rb b/examples/snippets/batch-dml/db/migrate/01_create_tables.rb new file mode 100644 index 00000000..632540e5 --- /dev/null +++ b/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 diff --git a/examples/snippets/batch-dml/db/seeds.rb b/examples/snippets/batch-dml/db/seeds.rb new file mode 100644 index 00000000..721cd42d --- /dev/null +++ b/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. diff --git a/examples/snippets/batch-dml/models/album.rb b/examples/snippets/batch-dml/models/album.rb new file mode 100644 index 00000000..4412d5d8 --- /dev/null +++ b/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 diff --git a/examples/snippets/batch-dml/models/singer.rb b/examples/snippets/batch-dml/models/singer.rb new file mode 100644 index 00000000..d34342e8 --- /dev/null +++ b/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 diff --git a/lib/active_record/connection_adapters/spanner/database_statements.rb b/lib/active_record/connection_adapters/spanner/database_statements.rb index 42cf4165..3b184af4 100644 --- a/lib/active_record/connection_adapters/spanner/database_statements.rb +++ b/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 diff --git a/lib/active_record/connection_adapters/spanner_adapter.rb b/lib/active_record/connection_adapters/spanner_adapter.rb index 0b8e5466..463e12fb 100644 --- a/lib/active_record/connection_adapters/spanner_adapter.rb +++ b/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 diff --git a/lib/activerecord_spanner_adapter/base.rb b/lib/activerecord_spanner_adapter/base.rb index ec045321..7380d049 100644 --- a/lib/activerecord_spanner_adapter/base.rb +++ b/lib/activerecord_spanner_adapter/base.rb @@ -49,6 +49,10 @@ def self.buffered_mutations? spanner_adapter? && connection&.current_spanner_transaction&.isolation == :buffered_mutations end + def self.dml_batch(&) + connection.dml_batch(&) + end + def self._should_use_standard_insert_record? values !(buffered_mutations? || (primary_key && values.is_a?(Hash))) || !spanner_adapter? end diff --git a/lib/activerecord_spanner_adapter/connection.rb b/lib/activerecord_spanner_adapter/connection.rb index f6e9fa56..d219ffa5 100644 --- a/lib/activerecord_spanner_adapter/connection.rb +++ b/lib/activerecord_spanner_adapter/connection.rb @@ -164,6 +164,12 @@ def ddl_batch? false end + # Returns true if this connection is currently executing a DML batch, and otherwise false. + def dml_batch? + return true if @dml_batch + false + end + ## # Starts a manual DDL batch. The batch must be ended by calling either run_batch or abort_batch. # @@ -178,18 +184,39 @@ def ddl_batch? # raise # end def start_batch_ddl - if @ddl_batch - raise Google::Cloud::FailedPreconditionError, "A DDL batch is already active on this connection" + if @ddl_batch && @dml_batch + raise Google::Cloud::FailedPreconditionError, "Batch is already active on this connection" end @ddl_batch = [] end + ## + # Starts a manual DML batch. The batch must be ended by calling either run_batch or abort_batch. + # + # @example + # begin + # connection.start_batch_dml + # connection.execute_query "insert into `Users` (Id, Name) VALUES (1, 'Test 1')" + # connection.execute_query "insert into `Users` (Id, Name) VALUES (2, 'Test 2')" + # connection.run_batch + # rescue StandardError + # connection.abort_batch + # raise + # end + def start_batch_dml + if @ddl_batch || @dml_batch + raise Google::Cloud::FailedPreconditionError, "A batch is already active on this connection" + end + @dml_batch = [] + end + ## # Aborts the current batch on this connection. This is a no-op if there is no batch on this connection. # # @see start_batch_ddl def abort_batch @ddl_batch = nil + @dml_batch = nil end ## @@ -198,10 +225,21 @@ def abort_batch # # @see start_batch_ddl def run_batch - unless @ddl_batch + unless @ddl_batch || @dml_batch raise Google::Cloud::FailedPreconditionError, "There is no batch active on this connection" end # Just return if the batch is empty. + return true if @ddl_batch&.empty? || @dml_batch&.empty? + begin + if @ddl_batch + run_ddl_batch + else + run_dml_batch + end + end + end + + def run_ddl_batch return true if @ddl_batch.empty? begin execute_ddl_statements @ddl_batch, nil, true @@ -210,9 +248,43 @@ def run_batch end end + def run_dml_batch + return true if @dml_batch.empty? + begin + # Execute the DML statements in the batch. + execute_dml_statements_in_batch @dml_batch + ensure + @dml_batch = nil + end + end + + ## + # Executes a set of DML statements as one batch. This method raises an error if no block is given. + def dml_batch + raise Google::Cloud::FailedPreconditionError, "No block given for the DML batch" unless block_given? + begin + start_batch_dml + yield + run_batch + rescue StandardError + abort_batch + raise + ensure + @dml_batch = nil + end + end + # DQL, DML Statements - def execute_query sql, params: nil, types: nil, single_use_selector: nil, request_options: nil + def execute_query sql, params: nil, types: nil, single_use_selector: nil, request_options: nil, statement_type: nil + # Clear the transaction from the previous statement. + unless current_transaction&.active? + self.current_transaction = nil + end + if statement_type == :dml && dml_batch? + @dml_batch.push({ sql: sql, params: params, types: types }) + return + end if params converted_params, types = Google::Cloud::Spanner::Convert.to_input_params_and_types( @@ -220,11 +292,6 @@ def execute_query sql, params: nil, types: nil, single_use_selector: nil, reques ) end - # Clear the transaction from the previous statement. - unless current_transaction&.active? - self.current_transaction = nil - end - selector = transaction_selector || single_use_selector execute_sql_request sql, converted_params, types, selector, request_options end @@ -353,6 +420,42 @@ def execute_ddl_statements statements, operation_id, wait_until_done job.done? end + # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity + def execute_dml_statements_in_batch statements + selector = transaction_selector + response = session.batch_update selector, current_transaction&.next_sequence_number do |batch| + statements.each do |statement| + batch.batch_update statement[:sql], params: statement[:params], types: statement[:types] + end + end + batch_update_results = Google::Cloud::Spanner::BatchUpdateResults.new response + first_res = response.result_sets.first + current_transaction.grpc_transaction = first_res.metadata.transaction \ + if current_transaction && first_res&.metadata&.transaction + batch_update_results.row_counts + rescue Google::Cloud::AbortedError + # Mark the current transaction as aborted to prevent any unnecessary further requests on the transaction. + current_transaction&.mark_aborted + raise + rescue Google::Cloud::Spanner::BatchUpdateError => e + # Check if the status is ABORTED, and if it is, just propagate an AbortedError. + if e.cause&.code == GRPC::Core::StatusCodes::ABORTED + current_transaction&.mark_aborted + raise e.cause + end + + # Check if the request returned a transaction or not. + # BatchDML is capable of returning BOTH an error and a transaction. + if current_transaction && !current_transaction.grpc_transaction? && selector&.begin&.read_write + selector = create_transaction_after_failed_first_statement e + retry + end + # It was not the first statement, or it returned a transaction, so propagate the error. + raise + end + # rubocop:enable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity + + ## # Retrieves the delay value from Google::Cloud::AbortedError or # GRPC::Aborted diff --git a/lib/activerecord_spanner_adapter/transaction.rb b/lib/activerecord_spanner_adapter/transaction.rb index c2763c60..34986ef2 100644 --- a/lib/activerecord_spanner_adapter/transaction.rb +++ b/lib/activerecord_spanner_adapter/transaction.rb @@ -165,6 +165,10 @@ def grpc_transaction= grpc @grpc_transaction = Google::Cloud::Spanner::Transaction.from_grpc grpc, @connection.session end + def grpc_transaction? + @grpc_transaction if @grpc_transaction + end + def transaction_selector return unless active? diff --git a/test/activerecord_spanner_mock_server/spanner_active_record_with_mock_server_test.rb b/test/activerecord_spanner_mock_server/spanner_active_record_with_mock_server_test.rb index c3005752..8ef962cd 100644 --- a/test/activerecord_spanner_mock_server/spanner_active_record_with_mock_server_test.rb +++ b/test/activerecord_spanner_mock_server/spanner_active_record_with_mock_server_test.rb @@ -14,6 +14,7 @@ module MockServerTests CommitRequest = Google::Cloud::Spanner::V1::CommitRequest ExecuteSqlRequest = Google::Cloud::Spanner::V1::ExecuteSqlRequest BeginTransactionRequest = Google::Cloud::Spanner::V1::BeginTransactionRequest + ExecuteBatchDmlRequest = Google::Cloud::Spanner::V1::ExecuteBatchDmlRequest class SpannerActiveRecordMockServerTest < BaseSpannerMockServerTest VERSION_7_1_0 = Gem::Version.create('7.1.0') @@ -148,6 +149,179 @@ def test_upsert assert_equal "id", mutation.insert_or_update.columns[2] end + def test_insert_batch_dml + sql1 = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Alice', 'Ecila')" + sql2 = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Pete', 'Etep')" + @mock.put_statement_result sql1, StatementResult.new(1) + @mock.put_statement_result sql2, StatementResult.new(1) + + singer1 = { first_name: "Alice", last_name: "Ecila" } + singer2 = { first_name: "Pete", last_name: "Etep" } + ActiveRecord::Base.transaction do + ActiveRecord::Base.dml_batch do + Singer.insert(singer1) + Singer.insert(singer2) + end + end + + # This test should use BatchDML instead of individual ExecuteSqlRequests. + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql1 } + assert_equal 0, execute_requests.length + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql2 } + assert_equal 0, execute_requests.length + batch_requests = @mock.requests.select { |req| req.is_a?(ExecuteBatchDmlRequest) } + assert_equal 1, batch_requests.length + assert_equal 2, batch_requests.first.statements.length + assert_equal sql1, batch_requests.first.statements[0].sql + assert_equal sql2, batch_requests.first.statements[1].sql + + # There should be only one transaction with no mutations. + commit_requests = @mock.requests.select { |req| req.is_a?(CommitRequest) } + assert_equal 1, commit_requests.length + mutations = commit_requests[0].mutations + assert_equal 0, mutations.length + end + + def test_update_batch_dml + num_rows = 4 + select_sql = "SELECT `singers`.* FROM `singers`" + @mock.put_statement_result select_sql, MockServerTests::create_random_singers_result(num_rows) + update_sql = "UPDATE `singers` SET `picture` = @p1 WHERE `singers`.`id` = @p2" + @mock.put_statement_result update_sql, StatementResult.new(1) + + ActiveRecord::Base.transaction do + ActiveRecord::Base.dml_batch do + Singer.all.each do |singer| + singer.picture = Base64.encode64(SecureRandom.alphanumeric(SecureRandom.random_number(10..200))) + singer.save + end + end + end + + # This test should use BatchDML instead of individual ExecuteSqlRequests. + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == update_sql } + assert_equal 0, execute_requests.length + batch_requests = @mock.requests.select { |req| req.is_a?(ExecuteBatchDmlRequest) } + assert_equal 1, batch_requests.length + assert_equal num_rows, batch_requests.first.statements.length + + # There should be only one transaction with no mutations. + commit_requests = @mock.requests.select { |req| req.is_a?(CommitRequest) } + assert_equal 1, commit_requests.length + mutations = commit_requests[0].mutations + assert_equal 0, mutations.length + end + + def test_insert_batch_dml_error + sql1 = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Alice', 'Ecila')" + sql2 = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Pete', 'Etep')" + error = GRPC::BadStatus.new GRPC::Core::StatusCodes::INVALID_ARGUMENT, "Invalid value" + @mock.put_statement_result sql1, StatementResult.new(1) + @mock.put_statement_result sql2, StatementResult.new(error) + + singer1 = { first_name: "Alice", last_name: "Ecila" } + singer2 = { first_name: "Pete", last_name: "Etep" } + ActiveRecord::Base.transaction do + err = assert_raises Google::Cloud::Spanner::BatchUpdateError do + ActiveRecord::Base.dml_batch do + Singer.insert(singer1) + Singer.insert(singer2) + end + end + # Batch DML returns an error with an array that contains the update counts + # of the successful statements. In this case, only the first statement + # succeeded. + assert_equal 1, err.row_counts.length + end + + # This test should use BatchDML instead of individual ExecuteSqlRequests. + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql1 } + assert_equal 0, execute_requests.length + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql2 } + assert_equal 0, execute_requests.length + batch_requests = @mock.requests.select { |req| req.is_a?(ExecuteBatchDmlRequest) } + # The BatchDML request is NOT retried, because it returns the first result, a transaction AND an error. + assert_equal 1, batch_requests.length + assert_equal 2, batch_requests.first.statements.length + assert_equal sql1, batch_requests.first.statements[0].sql + assert_equal sql2, batch_requests.first.statements[1].sql + + # There should be only one transaction with no mutations. + commit_requests = @mock.requests.select { |req| req.is_a?(CommitRequest) } + assert_equal 1, commit_requests.length + mutations = commit_requests[0].mutations + assert_equal 0, mutations.length + end + + def test_insert_batch_dml_error_on_first + sql = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Alice', 'Ecila')" + error = GRPC::BadStatus.new GRPC::Core::StatusCodes::INVALID_ARGUMENT, "Invalid value" + @mock.put_statement_result sql, StatementResult.new(error) + + singer1 = { first_name: "Alice", last_name: "Ecila" } + singer2 = { first_name: "Pete", last_name: "Etep" } + ActiveRecord::Base.transaction do + err = assert_raises Google::Cloud::Spanner::BatchUpdateError do + ActiveRecord::Base.dml_batch do + Singer.insert(singer1) + Singer.insert(singer2) + end + end + # Batch DML returns an error with an array that contains the update counts + # of the successful statements. In this case, non succeeded. + assert_equal 0, err.row_counts.length + end + + # This test should use BatchDML instead of individual ExecuteSqlRequests. + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql } + assert_equal 0, execute_requests.length + batch_requests = @mock.requests.select { |req| req.is_a?(ExecuteBatchDmlRequest) } + # The BatchDML request is retried, because it does not return any results or a transaction. + assert_equal 2, batch_requests.length + end + + def test_insert_batch_dml_aborted + sql1 = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Alice', 'Ecila')" + sql2 = "INSERT OR IGNORE INTO `singers` (`first_name`,`last_name`) VALUES ('Pete', 'Etep')" + error = GRPC::BadStatus.new GRPC::Core::StatusCodes::ABORTED, "Transaction aborted" + @mock.put_statement_result sql1, StatementResult.new(1) + @mock.put_statement_result sql2, StatementResult.new(error) + + singer1 = { first_name: "Alice", last_name: "Ecila" } + singer2 = { first_name: "Pete", last_name: "Etep" } + first = true + ActiveRecord::Base.transaction do + # Update the results of the second statement from ABORTED to success on the second attempt. + if first + first = false + else + @mock.put_statement_result sql2, StatementResult.new(1) + end + ActiveRecord::Base.dml_batch do + Singer.insert(singer1) + Singer.insert(singer2) + end + end + + # This test should use BatchDML instead of individual ExecuteSqlRequests. + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql1 } + assert_equal 0, execute_requests.length + execute_requests = @mock.requests.select { |req| req.is_a?(ExecuteSqlRequest) && req.sql == sql2 } + assert_equal 0, execute_requests.length + batch_requests = @mock.requests.select { |req| req.is_a?(ExecuteBatchDmlRequest) } + # The BatchDML request is retried, because the transaction is aborted. + assert_equal 2, batch_requests.length + assert_equal 2, batch_requests[1].statements.length + assert_equal sql1, batch_requests[1].statements[0].sql + assert_equal sql2, batch_requests[1].statements[1].sql + + # There should be only one transaction with no mutations. + commit_requests = @mock.requests.select { |req| req.is_a?(CommitRequest) } + assert_equal 1, commit_requests.length + mutations = commit_requests[0].mutations + assert_equal 0, mutations.length + end + def test_selects_all_singers_without_transaction sql = "SELECT `singers`.* FROM `singers`" @mock.put_statement_result sql, MockServerTests::create_random_singers_result(4) diff --git a/test/mock_server/spanner_mock_server.rb b/test/mock_server/spanner_mock_server.rb index 94014529..5dbf5c61 100644 --- a/test/mock_server/spanner_mock_server.rb +++ b/test/mock_server/spanner_mock_server.rb @@ -104,7 +104,30 @@ def do_execute_sql request, streaming def execute_batch_dml request, _unused_call @requests << request - raise GRPC::BadStatus.new GRPC::Core::StatusCodes::UNIMPLEMENTED, "Not yet implemented" + validate_session request.session + created_transaction = do_create_transaction request.session if request.transaction&.begin + transaction_id = created_transaction&.id || request.transaction&.id + validate_transaction request.session, transaction_id if transaction_id && transaction_id != "" + + status = Google::Rpc::Status.new + response = Google::Cloud::Spanner::V1::ExecuteBatchDmlResponse.new + first = true + request.statements.each do |stmt| + result = get_statement_result(stmt.sql).clone + if result.result_type == StatementResult::EXCEPTION + status.code = result.result.code + status.message = result.result.message + break + end + if first + response.result_sets << result.result(created_transaction) + first = false + else + response.result_sets << result.result + end + end + response.status = status + response end def read request, _unused_call