activerecord-spanner-adapter 0.6.0 → 1.0.1

data/examples/snippets/partitioned-dml/README.md

@@ -0,0 +1,16 @@
+# Sample - Partitioned DML
+
+This example shows how to use Partitioned DML with the Spanner ActiveRecord adapter.
+
+See https://cloud.google.com/spanner/docs/dml-partitioned for more information on Partitioned DML.
+
+## Running the Sample
+
+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
+```

data/examples/snippets/partitioned-dml/Rakefile

@@ -0,0 +1,13 @@
+# Copyright 2021 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 work with Partitioned DML in ActiveRecord."
+task :run do
+  Dir.chdir("..") { sh "bundle exec rake run[partitioned-dml]" }
+end

data/examples/snippets/partitioned-dml/application.rb

@@ -0,0 +1,48 @@
+# Copyright 2021 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
+    singer_count = Singer.all.count
+    album_count = Album.all.count
+    puts ""
+    puts "Singers in the database: #{singer_count}"
+    puts "Albums in the database: #{album_count}"
+
+    puts ""
+    puts "Deleting all albums in the database using Partitioned DML"
+    # Note that a Partitioned DML transaction can contain ONLY ONE DML statement.
+    # If we want to delete all data in two different tables, we need to do so in two different PDML transactions.
+    Album.transaction isolation: :pdml do
+      count = Album.delete_all
+      puts "Deleted #{count} albums"
+    end
+
+    puts ""
+    puts "Deleting all singers in the database using Partitioned DML"
+    Singer.transaction isolation: :pdml do
+      count = Singer.delete_all
+      puts "Deleted #{count} singers"
+    end
+
+    singer_count = Singer.all.count
+    album_count = Album.all.count
+    puts ""
+    puts "Singers in the database: #{singer_count}"
+    puts "Albums in the database: #{album_count}"
+
+    puts ""
+    puts "Press any key to end the application"
+    STDIN.getch
+  end
+end
+
+Application.run

data/examples/snippets/partitioned-dml/config/database.yml

@@ -0,0 +1,8 @@
+development:
+  adapter: spanner
+  emulator_host: localhost:9010
+  project: test-project
+  instance: test-instance
+  database: testdb
+  pool: 5
+  timeout: 5000

data/examples/snippets/partitioned-dml/db/migrate/01_create_tables.rb

@@ -0,0 +1,21 @@
+# Copyright 2021 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, limit: 100
+        t.string :last_name, limit: 200, null: false
+      end
+
+      create_table :albums do |t|
+        t.string :title
+        t.references :singer, index: false, foreign_key: true
+      end
+    end
+  end
+end

data/examples/snippets/partitioned-dml/db/schema.rb

@@ -0,0 +1,26 @@
+# This file is auto-generated from the current state of the database. Instead
+# of editing this file, please use the migrations feature of Active Record to
+# incrementally modify your database, and then regenerate this schema definition.
+#
+# This file is the source Rails uses to define your schema when running `bin/rails
+# db:schema:load`. When creating a new database, `bin/rails db:schema:load` tends to
+# be faster and is potentially less error prone than running all of your
+# migrations from scratch. Old migrations may fail to apply correctly if those
+# migrations use external dependencies or application code.
+#
+# It's strongly recommended that you check this file into your version control system.
+
+ActiveRecord::Schema.define(version: 1) do
+
+  create_table "albums", id: { limit: 8 }, force: :cascade do |t|
+    t.string "title"
+    t.integer "singer_id", limit: 8
+  end
+
+  create_table "singers", id: { limit: 8 }, force: :cascade do |t|
+    t.string "first_name", limit: 100
+    t.string "last_name", limit: 200, null: false
+  end
+
+  add_foreign_key "albums", "singers"
+end

data/examples/snippets/partitioned-dml/db/seeds.rb

@@ -0,0 +1,29 @@
+# Copyright 2021 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.rb"
+require_relative "../models/singer"
+require_relative "../models/album"
+
+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]
+
+adjectives = %w[daily happy blue generous cooked bad open]
+nouns = %w[windows potatoes bank street tree glass bottle]
+
+# This ensures all the records are inserted using one read/write transaction that will use mutations instead of DML.
+ActiveRecord::Base.transaction isolation: :buffered_mutations do
+  singers = []
+  5.times do
+    singers << Singer.create(first_name: first_names.sample, last_name: last_names.sample)
+  end
+
+  albums = []
+  20.times do
+    singer = singers.sample
+    albums << Album.create(title: "#{adjectives.sample} #{nouns.sample}", singer: singer)
+  end
+end

data/examples/snippets/partitioned-dml/models/album.rb

@@ -0,0 +1,9 @@
+# Copyright 2021 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

data/examples/snippets/partitioned-dml/models/singer.rb

@@ -0,0 +1,9 @@
+# Copyright 2021 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

data/lib/active_record/connection_adapters/spanner_adapter.rb

@@ -198,7 +198,7 @@ module ActiveRecord
 
       def register_array_types m
         m.register_type %r{^ARRAY<BOOL>}i, Type::Spanner::Array.new(Type::Boolean.new)
-        m.register_type %r{^ARRAY<BYTES\((MAX|d+)\)>}i, Type::Spanner::Array.new(Type::Binary.new)
+        m.register_type %r{^ARRAY<BYTES\((MAX|d+)\)>}i, Type::Spanner::Array.new(ActiveRecord::Type::Spanner::Bytes.new)
         m.register_type %r{^ARRAY<DATE>}i, Type::Spanner::Array.new(Type::Date.new)
         m.register_type %r{^ARRAY<FLOAT64>}i, Type::Spanner::Array.new(Type::Float.new)
         m.register_type %r{^ARRAY<NUMERIC>}i, Type::Spanner::Array.new(Type::Decimal.new)

data/lib/active_record/tasks/spanner_database_tasks.rb

@@ -18,7 +18,7 @@ module ActiveRecord
         @connection.create_database
       rescue Google::Cloud::Error => error
         if error.instance_of? Google::Cloud::AlreadyExistsError
-          raise ActiveRecord::Tasks::DatabaseAlreadyExists
+          raise ActiveRecord::DatabaseAlreadyExists
         end
 
         raise error

data/lib/active_record/type/spanner/array.rb

@@ -15,15 +15,29 @@ module ActiveRecord
           @element_type = element_type
         end
 
-        def serialize value
+        def cast value
           return super if value.nil?
-          return super unless @element_type.is_a? Type::Decimal
           return super unless value.respond_to? :map
 
-          # Convert a decimal (NUMERIC) array to a String array to prevent it from being encoded as a FLOAT64 array.
           value.map do |v|
-            next if v.nil?
-            v.to_s
+            @element_type.cast v
+          end
+        end
+
+        def serialize value
+          return super if value.nil?
+          return super unless value.respond_to? :map
+
+          if @element_type.is_a? ActiveRecord::Type::Decimal
+            # Convert a decimal (NUMERIC) array to a String array to prevent it from being encoded as a FLOAT64 array.
+            value.map do |v|
+              next if v.nil?
+              v.to_s
+            end
+          else
+            value.map do |v|
+              @element_type.serialize v
+            end
           end
         end
       end

data/lib/activerecord_spanner_adapter/base.rb

@@ -15,7 +15,17 @@ module ActiveRecord
     # Creates an object (or multiple objects) and saves it to the database. This method will use mutations instead
     # of DML if there is no active transaction, or if the active transaction has been created with the option
     # isolation: :buffered_mutations.
+    def self.create! attributes = nil, &block
+      return super unless spanner_adapter?
+      return super if active_transaction?
+
+      transaction isolation: :buffered_mutations do
+        return super
+      end
+    end
+
     def self.create attributes = nil, &block
+      return super unless spanner_adapter?
       return super if active_transaction?
 
       transaction isolation: :buffered_mutations do
@@ -23,8 +33,16 @@ module ActiveRecord
       end
     end
 
+    def self.spanner_adapter?
+      connection.adapter_name == "spanner"
+    end
+
+    def self.buffered_mutations?
+      spanner_adapter? && connection&.current_spanner_transaction&.isolation == :buffered_mutations
+    end
+
     def self._insert_record values
-      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+      return super unless buffered_mutations?
 
       primary_key = self.primary_key
       primary_key_value = nil
@@ -48,7 +66,7 @@ module ActiveRecord
           values: [grpc_values.list_value]
         )
       )
-      Base.connection.current_spanner_transaction.buffer mutation
+      connection.current_spanner_transaction.buffer mutation
 
       primary_key_value
     end
@@ -56,6 +74,7 @@ module ActiveRecord
     # Deletes all records of this class. This method will use mutations instead of DML if there is no active
     # transaction, or if the active transaction has been created with the option isolation: :buffered_mutations.
     def self.delete_all
+      return super unless spanner_adapter?
       return super if active_transaction?
 
       transaction isolation: :buffered_mutations do
@@ -72,7 +91,8 @@ module ActiveRecord
     # of DML if there is no active transaction, or if the active transaction has been created with the option
     # isolation: :buffered_mutations.
     def update attributes
-      return super if Base.active_transaction?
+      return super unless self.class.spanner_adapter?
+      return super if self.class.active_transaction?
 
       transaction isolation: :buffered_mutations do
         return super
@@ -83,7 +103,8 @@ module ActiveRecord
     # of DML if there is no active transaction, or if the active transaction has been created with the option
     # isolation: :buffered_mutations.
     def destroy
-      return super if Base.active_transaction?
+      return super unless self.class.spanner_adapter?
+      return super if self.class.active_transaction?
 
       transaction isolation: :buffered_mutations do
         return super
@@ -110,7 +131,7 @@ module ActiveRecord
     private_class_method :_create_grpc_values_for_insert
 
     def _update_row attribute_names, attempted_action = "update"
-      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+      return super unless self.class.buffered_mutations?
 
       if locking_enabled?
         _execute_version_check attempted_action
@@ -129,7 +150,7 @@ module ActiveRecord
           values: [grpc_values.list_value]
         )
       )
-      Base.connection.current_spanner_transaction.buffer mutation
+      self.class.connection.current_spanner_transaction.buffer mutation
       1 # Affected rows
     end
 
@@ -161,13 +182,13 @@ module ActiveRecord
     end
 
     def destroy_row
-      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+      return super unless self.class.buffered_mutations?
 
       _delete_row
     end
 
     def _delete_row
-      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+      return super unless self.class.buffered_mutations?
       if locking_enabled?
         _execute_version_check "destroy"
       end
@@ -186,7 +207,7 @@ module ActiveRecord
           key_set: { keys: [list_value] }
         )
       )
-      Base.connection.current_spanner_transaction.buffer mutation
+      self.class.connection.current_spanner_transaction.buffer mutation
       1 # Affected rows
     end
 
@@ -210,7 +231,7 @@ module ActiveRecord
               "WHERE `#{self.class.primary_key}` = @id AND `#{locking_column}` = @lock_version"
       params = { "id" => id_in_database, "lock_version" => previous_lock_value }
       param_types = { "id" => :INT64, "lock_version" => :INT64 }
-      locked_row = Base.connection.raw_connection.execute_query sql, params: params, types: param_types
+      locked_row = self.class.connection.raw_connection.execute_query sql, params: params, types: param_types
       raise ActiveRecord::StaleObjectError.new(self, attempted_action) unless locked_row.rows.any?
     end
   end

data/lib/activerecord_spanner_adapter/connection.rb

@@ -208,27 +208,53 @@ module ActiveRecordSpannerAdapter
         self.current_transaction = nil
       end
 
-      begin
-        session.execute_query \
-          sql,
-          params: converted_params,
-          types: types,
-          transaction: transaction_selector || single_use_selector,
-          seqno: (current_transaction&.next_sequence_number)
-      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::NotFoundError => e
-        if session_not_found?(e) || transaction_not_found?(e)
-          reset!
-          # Force a retry of the entire transaction if this statement was executed as part of a transaction.
-          # Otherwise, just retry the statement itself.
-          raise_aborted_err if current_transaction&.active?
-          retry
-        end
-        raise
+      selector = transaction_selector || single_use_selector
+      execute_sql_request sql, converted_params, types, selector
+    end
+
+    def execute_sql_request sql, converted_params, types, selector
+      res = session.execute_query \
+        sql,
+        params: converted_params,
+        types: types,
+        transaction: selector,
+        seqno: (current_transaction&.next_sequence_number)
+      current_transaction.grpc_transaction = res.metadata.transaction \
+          if current_transaction && res&.metadata&.transaction
+      res
+    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::NotFoundError => e
+      if session_not_found?(e) || transaction_not_found?(e)
+        reset!
+        # Force a retry of the entire transaction if this statement was executed as part of a transaction.
+        # Otherwise, just retry the statement itself.
+        raise_aborted_err if current_transaction&.active?
+        retry
+      end
+      raise
+    rescue Google::Cloud::Error => e
+      # 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.
+      if current_transaction && selector&.begin&.read_write
+        selector = create_transaction_after_failed_first_statement e
+        retry
       end
+      # It was not the first statement, so propagate the error.
+      raise
+    end
+
+    # Creates a transaction using a BeginTransaction RPC. This is used if the first statement of a
+    # transaction fails, as that also means that no transaction id was returned.
+    def create_transaction_after_failed_first_statement original_error
+      transaction = current_transaction.force_begin_read_write
+      Google::Spanner::V1::TransactionSelector.new id: transaction.transaction_id
+    rescue Google::Cloud::Error
+      # Raise the original error if the BeginTransaction RPC also fails.
+      raise original_error
     end
 
     # Transactions

data/lib/activerecord_spanner_adapter/information_schema.rb

@@ -62,7 +62,8 @@ module ActiveRecordSpannerAdapter
     end
 
     def table_columns table_name, column_name: nil
-      sql = +"SELECT COLUMN_NAME, SPANNER_TYPE, IS_NULLABLE, COLUMN_DEFAULT, ORDINAL_POSITION"
+      sql = +"SELECT COLUMN_NAME, SPANNER_TYPE, IS_NULLABLE,"
+      sql << " CAST(COLUMN_DEFAULT AS STRING) AS COLUMN_DEFAULT, ORDINAL_POSITION"
       sql << " FROM INFORMATION_SCHEMA.COLUMNS"
       sql << " WHERE TABLE_NAME=%<table_name>s"
       sql << " AND COLUMN_NAME=%<column_name>s" if column_name

data/lib/activerecord_spanner_adapter/transaction.rb

@@ -30,28 +30,37 @@ module ActiveRecordSpannerAdapter
       @mutations << mutation
     end
 
+    # Begins the transaction.
+    #
+    # Read-only and PDML transactions are started by executing a BeginTransaction RPC.
+    # Read/write transactions are not really started by this method, and instead a
+    # transaction selector is prepared that will be included with the first statement
+    # on the transaction.
     def begin
       raise "Nested transactions are not allowed" if @state != :INITIALIZED
       begin
-        @grpc_transaction =
-          case @isolation
-          when Hash
-            if @isolation[:timestamp]
-              @connection.session.create_snapshot timestamp: @isolation[:timestamp]
-            elsif @isolation[:staleness]
-              @connection.session.create_snapshot staleness: @isolation[:staleness]
-            elsif @isolation[:strong]
-              @connection.session.create_snapshot strong: true
-            else
-              raise "Invalid snapshot argument: #{@isolation}"
-            end
-          when :read_only
-            @connection.session.create_snapshot strong: true
-          when :pdml
-            @connection.session.create_pdml
+        case @isolation
+        when Hash
+          if @isolation[:timestamp]
+            @grpc_transaction = @connection.session.create_snapshot timestamp: @isolation[:timestamp]
+          elsif @isolation[:staleness]
+            @grpc_transaction = @connection.session.create_snapshot staleness: @isolation[:staleness]
+          elsif @isolation[:strong]
+            @grpc_transaction = @connection.session.create_snapshot strong: true
           else
-            @connection.session.create_transaction
+            raise "Invalid snapshot argument: #{@isolation}"
           end
+        when :read_only
+          @grpc_transaction = @connection.session.create_snapshot strong: true
+        when :pdml
+          @grpc_transaction = @connection.session.create_pdml
+        else
+          @begin_transaction_selector = Google::Spanner::V1::TransactionSelector.new \
+            begin: Google::Spanner::V1::TransactionOptions.new(
+              read_write: Google::Spanner::V1::TransactionOptions::ReadWrite.new
+            )
+
+        end
         @state = :STARTED
       rescue Google::Cloud::NotFoundError => e
         if @connection.session_not_found? e
@@ -66,6 +75,12 @@ module ActiveRecordSpannerAdapter
       end
     end
 
+    # Forces a BeginTransaction RPC for a read/write transaction. This is used by a
+    # connection if the first statement of a transaction failed.
+    def force_begin_read_write
+      @grpc_transaction = @connection.session.create_transaction
+    end
+
     def next_sequence_number
       @sequence_number += 1 if @committable
     end
@@ -74,7 +89,10 @@ module ActiveRecordSpannerAdapter
       raise "This transaction is not active" unless active?
 
       begin
-        @connection.session.commit_transaction @grpc_transaction, @mutations if @committable
+        # 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
         @state = :COMMITTED
       rescue Google::Cloud::NotFoundError => e
         if @connection.session_not_found? e
@@ -93,7 +111,7 @@ module ActiveRecordSpannerAdapter
     def rollback
       # Allow rollback after abort and/or a failed commit.
       raise "This transaction is not active" unless active? || @state == :FAILED || @state == :ABORTED
-      if active?
+      if active? && @grpc_transaction
         # We do a shoot-and-forget rollback here, as the error that caused the transaction to be rolled back could
         # also have invalidated the transaction (e.g. `Session not found`). If the rollback fails for any other
         # reason, we also do not need to retry it or propagate the error to the application, as the transaction will
@@ -113,11 +131,24 @@ module ActiveRecordSpannerAdapter
       @state = :ABORTED
     end
 
+    # Sets the underlying gRPC transaction to use for this Transaction.
+    # This is used for queries/DML statements that inlined the BeginTransaction option and returned
+    # a transaction in the metadata.
+    def grpc_transaction= grpc
+      @grpc_transaction = Google::Cloud::Spanner::Transaction.from_grpc grpc, @connection.session
+    end
+
     def transaction_selector
       return unless active?
 
-      Google::Spanner::V1::TransactionSelector.new \
-        id: @grpc_transaction.transaction_id
+      # Use the transaction that has been started by a BeginTransaction RPC or returned by a
+      # statement, if present.
+      return Google::Spanner::V1::TransactionSelector.new id: @grpc_transaction.transaction_id \
+          if @grpc_transaction
+
+      # Return a transaction selector that will instruct the statement to also start a transaction
+      # and return its id as a side effect.
+      @begin_transaction_selector
     end
   end
 end

data/lib/activerecord_spanner_adapter/version.rb

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

data/lib/arel/visitors/spanner.rb

@@ -14,11 +14,24 @@ module Arel # :nodoc: all
       end
     end
 
+    class StatementHint
+      attr_reader :value
+
+      def initialize value
+        @value = value
+      end
+    end
+
     class Spanner < Arel::Visitors::ToSql
       def compile node, collector = Arel::Collectors::SQLString.new
         collector.class.module_eval { attr_accessor :hints }
+        collector.class.module_eval { attr_accessor :table_hints }
+        collector.class.module_eval { attr_accessor :join_hints }
         collector.hints = {}
+        collector.table_hints = {}
+        collector.join_hints = {}
         sql, binds = accept(node, collector).value
+        sql = collector.hints[:statement_hint].value + sql if collector.hints[:statement_hint]
         binds << collector.hints[:staleness] if collector.hints[:staleness]
         [sql, binds]
       end
@@ -32,9 +45,25 @@ module Arel # :nodoc: all
         BIND_BLOCK
       end
 
+      def visit_table_hint v, collector
+        value = v.delete_prefix("table_hint:").strip
+        # TODO: This does not support FORCE_INDEX hints that reference an index that contains '@{' in the name.
+        start_of_hint_index = value.rindex "@{"
+        table_name = value[0, start_of_hint_index]
+        table_hint = value[start_of_hint_index, value.length]
+        collector.table_hints[table_name] = table_hint if table_name && table_hint
+      end
+
+      def visit_statement_hint v, collector
+        collector.hints[:statement_hint] = \
+          StatementHint.new v.delete_prefix("statement_hint:")
+      end
+
       # rubocop:disable Naming/MethodName
       def visit_Arel_Nodes_OptimizerHints o, collector
         o.expr.each do |v|
+          visit_table_hint v, collector if v.start_with? "table_hint:"
+          visit_statement_hint v, collector if v.start_with? "statement_hint:"
           if v.start_with? "max_staleness:"
             collector.hints[:staleness] = \
               StalenessHint.new max_staleness: v.delete_prefix("max_staleness:").to_f
@@ -59,6 +88,16 @@ module Arel # :nodoc: all
         collector
       end
 
+      def visit_Arel_Table o, collector
+        return super unless collector.table_hints[o.name]
+        if o.table_alias
+          collector << quote_table_name(o.name) << collector.table_hints[o.name] \
+                    << " " << quote_table_name(o.table_alias)
+        else
+          collector << quote_table_name(o.name) << collector.table_hints[o.name]
+        end
+      end
+
       def visit_Arel_Nodes_BindParam o, collector
         # Do not generate a query parameter if the value should be set to the PENDING_COMMIT_TIMESTAMP(), as that is
         # not supported as a parameter value by Cloud Spanner.

data/lib/spanner_client_ext.rb

@@ -95,6 +95,10 @@ module Google
         end
       end
 
+      class Results
+        attr_reader :metadata
+      end
+
       class Transaction
         attr_accessor :seqno, :commit
       end