activerecord-spanner-adapter 2.3.0 → 2.5.0

data/lib/activerecord_spanner_adapter/connection.rb

@@ -6,9 +6,13 @@
 
 require "google/cloud/spanner"
 require "spanner_client_ext"
+require "active_record/connection_adapters/spanner/type_mapping"
 require "activerecord_spanner_adapter/information_schema"
+require_relative "../active_record/connection_adapters/spanner/errors/transaction_mutation_limit_exceeded_error"
 
 module ActiveRecordSpannerAdapter
+  TransactionMutationLimitExceededError = Google::Cloud::Spanner::Errors::TransactionMutationLimitExceededError
+
   class Connection
     attr_reader :instance_id
     attr_reader :database_id
@@ -40,10 +44,16 @@ module ActiveRecordSpannerAdapter
       end
     end
 
+    def native_database_types
+      NATIVE_DATABASE_TYPES
+    end
+
     # Clears the cached information about the underlying information schemas.
     # Call this method if you drop and recreate a database with the same name
     # to prevent the cached information to be used for the new database.
     def self.reset_information_schemas!
+      return unless @database
+
       @information_schemas.each_value do |info_schema|
         info_schema.connection.disconnect!
       end
@@ -159,6 +169,12 @@ module ActiveRecordSpannerAdapter
       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.
     #
@@ -173,18 +189,39 @@ module ActiveRecordSpannerAdapter
     #     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
 
     ##
@@ -193,10 +230,21 @@ module ActiveRecordSpannerAdapter
     #
     # @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
@@ -205,9 +253,43 @@ module ActiveRecordSpannerAdapter
       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(
@@ -215,11 +297,6 @@ module ActiveRecordSpannerAdapter
           )
       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
@@ -250,6 +327,9 @@ module ActiveRecordSpannerAdapter
       end
       raise
     rescue Google::Cloud::Error => e
+      if TransactionMutationLimitExceededError.is_mutation_limit_error? e
+        raise
+      end
       # Check if it was the first statement in a transaction that included a BeginTransaction
       # option in the request. If so, execute an explicit BeginTransaction and then retry the
       # request without the BeginTransaction option.
@@ -274,9 +354,11 @@ module ActiveRecordSpannerAdapter
 
     # Transactions
 
-    def begin_transaction isolation = nil
+    def begin_transaction isolation = nil, **options
       raise "Nested transactions are not allowed" if current_transaction&.active?
-      self.current_transaction = Transaction.new self, isolation || @isolation_level
+      exclude_from_streams = options.fetch :exclude_txn_from_change_streams, false
+      self.current_transaction = Transaction.new self, isolation || @isolation_level,
+                                                 exclude_txn_from_change_streams: exclude_from_streams
       current_transaction.begin
       current_transaction
     end
@@ -343,6 +425,42 @@ module ActiveRecordSpannerAdapter
       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

data/lib/activerecord_spanner_adapter/transaction.rb

@@ -7,14 +7,21 @@
 module ActiveRecordSpannerAdapter
   class Transaction
     attr_reader :state
+    attr_reader :commit_options
+    attr_reader :begin_transaction_selector
+    attr_accessor :exclude_txn_from_change_streams
 
-    def initialize connection, isolation
+
+
+    def initialize connection, isolation, commit_options = nil, exclude_txn_from_change_streams: false
       @connection = connection
       @isolation = isolation
       @committable = ![:read_only, :pdml].include?(isolation) && !isolation.is_a?(Hash)
       @state = :INITIALIZED
       @sequence_number = 0
       @mutations = []
+      @commit_options = commit_options
+      @exclude_txn_from_change_streams = exclude_txn_from_change_streams
     end
 
     def active?
@@ -59,7 +66,8 @@ module ActiveRecordSpannerAdapter
           @begin_transaction_selector = Google::Cloud::Spanner::V1::TransactionSelector.new \
             begin: Google::Cloud::Spanner::V1::TransactionOptions.new(
               read_write: Google::Cloud::Spanner::V1::TransactionOptions::ReadWrite.new,
-              isolation_level: grpc_isolation
+              isolation_level: grpc_isolation,
+              exclude_txn_from_change_streams: @exclude_txn_from_change_streams
             )
         end
         @state = :STARTED
@@ -95,14 +103,23 @@ module ActiveRecordSpannerAdapter
       @sequence_number += 1 if @committable
     end
 
+    # Sets the commit options for this transaction.
+    # This is used to set the options for the commit RPC, such as return_commit_stats and max_commit_delay.
+    def set_commit_options options # rubocop:disable Naming/AccessorMethodName
+      @commit_options = options&.dup
+    end
+
     def commit
       raise "This transaction is not active" unless active?
 
       begin
         # Start a transaction with an explicit BeginTransaction RPC if the transaction only contains mutations.
         force_begin_read_write if @committable && !@mutations.empty? && !@grpc_transaction
-
-        @connection.session.commit_transaction @grpc_transaction, @mutations if @committable && @grpc_transaction
+        if @committable && @grpc_transaction
+          @connection.session.commit_transaction @grpc_transaction,
+                                                 @mutations,
+                                                 commit_options: commit_options
+        end
         @state = :COMMITTED
       rescue Google::Cloud::NotFoundError => e
         if @connection.session_not_found? e
@@ -148,6 +165,10 @@ module ActiveRecordSpannerAdapter
       @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?
 

data/lib/activerecord_spanner_adapter/version.rb

@@ -5,5 +5,5 @@
 # https://opensource.org/licenses/MIT.
 
 module ActiveRecordSpannerAdapter
-  VERSION = "2.3.0".freeze
+  VERSION = "2.5.0".freeze
 end

data/lib/spanner_client_ext.rb

@@ -23,13 +23,14 @@ module Google
       end
 
       class Session
-        def commit_transaction transaction, mutations = []
+        def commit_transaction transaction, mutations = [], commit_options: nil
           ensure_service!
 
           resp = service.commit(
             path,
             mutations,
-            transaction_id: transaction.transaction_id
+            transaction_id: transaction.transaction_id,
+            commit_options: commit_options
           )
           @last_updated_at = Time.now
           Convert.timestamp_to_time resp.commit_timestamp

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-spanner-adapter
 version: !ruby/object:Gem::Version
-  version: 2.3.0
+  version: 2.5.0
 platform: ruby
 authors:
 - Google LLC
@@ -75,16 +75,22 @@ dependencies:
   name: bundler
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
         version: '2.0'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '5.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
       - !ruby/object:Gem::Version
         version: '2.0'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '5.0'
 - !ruby/object:Gem::Dependency
   name: google-style
   requirement: !ruby/object:Gem::Requirement
@@ -147,14 +153,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.2'
+        version: '5.4'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '5.2'
+        version: '5.4'
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -365,6 +371,14 @@ files:
 - examples/snippets/auto-generated-primary-key/db/seeds.rb
 - examples/snippets/auto-generated-primary-key/models/album.rb
 - examples/snippets/auto-generated-primary-key/models/singer.rb
+- examples/snippets/batch-dml/README.md
+- examples/snippets/batch-dml/Rakefile
+- examples/snippets/batch-dml/application.rb
+- examples/snippets/batch-dml/config/database.yml
+- examples/snippets/batch-dml/db/migrate/01_create_tables.rb
+- examples/snippets/batch-dml/db/seeds.rb
+- examples/snippets/batch-dml/models/album.rb
+- examples/snippets/batch-dml/models/singer.rb
 - examples/snippets/bin/create_emulator_instance.rb
 - examples/snippets/bit-reversed-sequence/README.md
 - examples/snippets/bit-reversed-sequence/Rakefile
@@ -538,12 +552,14 @@ files:
 - examples/solidus/README.md
 - lib/active_record/connection_adapters/spanner/column.rb
 - lib/active_record/connection_adapters/spanner/database_statements.rb
+- lib/active_record/connection_adapters/spanner/errors/transaction_mutation_limit_exceeded_error.rb
 - lib/active_record/connection_adapters/spanner/quoting.rb
 - lib/active_record/connection_adapters/spanner/schema_cache.rb
 - lib/active_record/connection_adapters/spanner/schema_creation.rb
 - lib/active_record/connection_adapters/spanner/schema_definitions.rb
 - lib/active_record/connection_adapters/spanner/schema_dumper.rb
 - lib/active_record/connection_adapters/spanner/schema_statements.rb
+- lib/active_record/connection_adapters/spanner/type_mapping.rb
 - lib/active_record/connection_adapters/spanner/type_metadata.rb
 - lib/active_record/connection_adapters/spanner_adapter.rb
 - lib/active_record/tasks/spanner_database_tasks.rb
@@ -551,6 +567,7 @@ files:
 - lib/active_record/type/spanner/bytes.rb
 - lib/active_record/type/spanner/spanner_active_record_converter.rb
 - lib/active_record/type/spanner/time.rb
+- lib/active_record/type/spanner/uuid.rb
 - lib/activerecord-spanner-adapter.rb
 - lib/activerecord_spanner_adapter/base.rb
 - lib/activerecord_spanner_adapter/connection.rb