activerecord-spanner-adapter 1.0.1 → 1.2.1

data/lib/active_record/connection_adapters/spanner_adapter.rb

@@ -8,6 +8,7 @@ require "securerandom"
 require "google/cloud/spanner"
 require "spanner_client_ext"
 require "active_record/connection_adapters/abstract_adapter"
+require "active_record/connection_adapters/abstract/connection_pool"
 require "active_record/connection_adapters/spanner/database_statements"
 require "active_record/connection_adapters/spanner/schema_statements"
 require "active_record/connection_adapters/spanner/schema_cache"
@@ -43,9 +44,9 @@ module ActiveRecord
   module ConnectionAdapters
     module AbstractPool
       def get_schema_cache connection
-        @schema_cache ||= SpannerSchemaCache.new connection
-        @schema_cache.connection = connection
-        @schema_cache
+        self.schema_cache ||= SpannerSchemaCache.new connection
+        schema_cache.connection = connection
+        schema_cache
       end
     end
 
@@ -178,41 +179,73 @@ module ActiveRecord
         Arel::Visitors::Spanner.new self
       end
 
-      private
+      def build_insert_sql insert
+        if current_spanner_transaction&.isolation == :buffered_mutations
+          raise "ActiveRecordSpannerAdapter does not support insert_sql with buffered_mutations transaction."
+        end
 
-      def initialize_type_map m = type_map
-        m.register_type "BOOL", Type::Boolean.new
-        register_class_with_limit(
-          m, %r{^BYTES}i, ActiveRecord::Type::Spanner::Bytes
-        )
-        m.register_type "DATE", Type::Date.new
-        m.register_type "FLOAT64", Type::Float.new
-        m.register_type "NUMERIC", Type::Decimal.new
-        m.register_type "INT64", Type::Integer.new(limit: 8)
-        register_class_with_limit m, %r{^STRING}i, Type::String
-        m.register_type "TIMESTAMP", ActiveRecord::Type::Spanner::Time.new
-        m.register_type "JSON", ActiveRecord::Type::Json.new
+        if insert.skip_duplicates? || insert.update_duplicates?
+          raise NotImplementedError, "CloudSpanner does not support skip_duplicates and update_duplicates."
+        end
 
-        register_array_types m
+        values_list, = insert.values_list
+        "INSERT #{insert.into} #{values_list}"
       end
 
-      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(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)
-        m.register_type %r{^ARRAY<INT64>}i, Type::Spanner::Array.new(Type::Integer.new(limit: 8))
-        m.register_type %r{^ARRAY<STRING\((MAX|d+)\)>}i, Type::Spanner::Array.new(Type::String.new)
-        m.register_type %r{^ARRAY<TIMESTAMP>}i, Type::Spanner::Array.new(ActiveRecord::Type::Spanner::Time.new)
-        m.register_type %r{^ARRAY<JSON>}i, Type::Spanner::Array.new(ActiveRecord::Type::Json.new)
+      module TypeMapBuilder
+        private
+
+        def initialize_type_map m = type_map
+          m.register_type "BOOL", Type::Boolean.new
+          register_class_with_limit(
+            m, %r{^BYTES}i, ActiveRecord::Type::Spanner::Bytes
+          )
+          m.register_type "DATE", Type::Date.new
+          m.register_type "FLOAT64", Type::Float.new
+          m.register_type "NUMERIC", Type::Decimal.new
+          m.register_type "INT64", Type::Integer.new(limit: 8)
+          register_class_with_limit m, %r{^STRING}i, Type::String
+          m.register_type "TIMESTAMP", ActiveRecord::Type::Spanner::Time.new
+          m.register_type "JSON", ActiveRecord::Type::Json.new
+
+          register_array_types m
+        end
+
+        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(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)
+          m.register_type %r{^ARRAY<INT64>}i, Type::Spanner::Array.new(Type::Integer.new(limit: 8))
+          m.register_type %r{^ARRAY<STRING\((MAX|d+)\)>}i, Type::Spanner::Array.new(Type::String.new)
+          m.register_type %r{^ARRAY<TIMESTAMP>}i, Type::Spanner::Array.new(ActiveRecord::Type::Spanner::Time.new)
+          m.register_type %r{^ARRAY<JSON>}i, Type::Spanner::Array.new(ActiveRecord::Type::Json.new)
+        end
+
+        def extract_limit sql_type
+          value = /\((.*)\)/.match sql_type
+          return unless value
+
+          value[1] == "MAX" ? "MAX" : value[1].to_i
+        end
       end
 
-      def extract_limit sql_type
-        value = /\((.*)\)/.match sql_type
-        return unless value
+      if ActiveRecord::VERSION::MAJOR >= 7
+        class << self
+          include TypeMapBuilder
+        end
+
+        TYPE_MAP = Type::TypeMap.new.tap { |m| initialize_type_map m }
+
+        private
 
-        value[1] == "MAX" ? "MAX" : value[1].to_i
+        def type_map
+          TYPE_MAP
+        end
+      else
+        include TypeMapBuilder
       end
 
       def translate_exception exception, message:, sql:, binds:

data/lib/activerecord_spanner_adapter/base.rb

@@ -4,6 +4,8 @@
 # license that can be found in the LICENSE file or at
 # https://opensource.org/licenses/MIT.
 
+require "activerecord_spanner_adapter/relation"
+
 module ActiveRecord
   class TableMetadata # :nodoc:
     # This attr_reader is private in ActiveRecord 6.0.x and public in 6.1.x. This makes sure it is always available in
@@ -42,35 +44,137 @@ module ActiveRecord
     end
 
     def self._insert_record values
-      return super unless buffered_mutations?
+      return super unless buffered_mutations? || (primary_key && values.is_a?(Hash))
 
-      primary_key = self.primary_key
-      primary_key_value = nil
+      return _buffer_record values, :insert if buffered_mutations?
 
-      if primary_key && values.is_a?(Hash)
-        primary_key_value = values[primary_key]
+      primary_key_value =
+        if primary_key.is_a? Array
+          _set_composite_primary_key_values primary_key, values
+        else
+          _set_single_primary_key_value primary_key, values
+        end
+      if ActiveRecord::VERSION::MAJOR >= 7
+        im = Arel::InsertManager.new arel_table
+        im.insert(values.transform_keys { |name| arel_table[name] })
+      else
+        im = arel_table.compile_insert _substitute_values(values)
+      end
+      connection.insert(im, "#{self} Create", primary_key || false, primary_key_value)
+    end
 
-        if !primary_key_value && prefetch_primary_key?
-          primary_key_value = next_sequence_value
-          values[primary_key] = primary_key_value
+    def self._upsert_record values
+      _buffer_record values, :insert_or_update
+    end
+
+    def self.insert_all _attributes, _returning: nil, _unique_by: nil
+      raise NotImplementedError, "Cloud Spanner does not support skip_duplicates."
+    end
+
+    def self.insert_all! attributes, returning: nil
+      return super unless spanner_adapter?
+      return super if active_transaction? && !buffered_mutations?
+
+      # This might seem inefficient, but is actually not, as it is only buffering a mutation locally.
+      # The mutations will be sent as one batch when the transaction is committed.
+      if active_transaction?
+        attributes.each do |record|
+          _insert_record record
+        end
+      else
+        transaction isolation: :buffered_mutations do
+          attributes.each do |record|
+            _insert_record record
+          end
         end
       end
+    end
+
+    def self.upsert_all attributes, returning: nil, unique_by: nil
+      return super unless spanner_adapter?
+      if active_transaction? && !buffered_mutations?
+        raise NotImplementedError, "Cloud Spanner does not support upsert using DML. " \
+                                   "Use upsert outside a transaction block or in a transaction " \
+                                   "block with isolation: :buffered_mutations"
+      end
+
+      # This might seem inefficient, but is actually not, as it is only buffering a mutation locally.
+      # The mutations will be sent as one batch when the transaction is committed.
+      if active_transaction?
+        attributes.each do |record|
+          _upsert_record record
+        end
+      else
+        transaction isolation: :buffered_mutations do
+          attributes.each do |record|
+            _upsert_record record
+          end
+        end
+      end
+    end
+
+    def self._buffer_record values, method
+      primary_key_value =
+        if primary_key.is_a? Array
+          _set_composite_primary_key_values primary_key, values
+        else
+          _set_single_primary_key_value primary_key, values
+        end
 
       metadata = TableMetadata.new self, arel_table
       columns, grpc_values = _create_grpc_values_for_insert metadata, values
 
+      write = Google::Cloud::Spanner::V1::Mutation::Write.new(
+        table: arel_table.name,
+        columns: columns,
+        values: [grpc_values.list_value]
+      )
       mutation = Google::Cloud::Spanner::V1::Mutation.new(
-        insert: Google::Cloud::Spanner::V1::Mutation::Write.new(
-          table: arel_table.name,
-          columns: columns,
-          values: [grpc_values.list_value]
-        )
+        "#{method}": write
       )
+
       connection.current_spanner_transaction.buffer mutation
 
       primary_key_value
     end
 
+    def self._set_composite_primary_key_values primary_key, values
+      primary_key_value = []
+      primary_key.each do |col|
+        value = values[col]
+
+        if !value && prefetch_primary_key?
+          value =
+            if ActiveRecord::VERSION::MAJOR >= 7
+              ActiveModel::Attribute.from_database col, next_sequence_value, ActiveModel::Type::BigInteger.new
+            else
+              next_sequence_value
+            end
+          values[col] = value
+        end
+        if value.is_a? ActiveModel::Attribute
+          value = value.value
+        end
+        primary_key_value.append value
+      end
+      primary_key_value
+    end
+
+    def self._set_single_primary_key_value primary_key, values
+      primary_key_value = values[primary_key] || values[primary_key.to_sym]
+
+      if !primary_key_value && prefetch_primary_key?
+        primary_key_value = next_sequence_value
+        if ActiveRecord::VERSION::MAJOR >= 7
+          values[primary_key] = ActiveModel::Attribute.from_database primary_key, primary_key_value,
+                                                                     ActiveModel::Type::BigInteger.new
+        else
+          values[primary_key] = primary_key_value
+        end
+      end
+      primary_key_value
+    end
+
     # 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
@@ -87,6 +191,14 @@ module ActiveRecord
       !(current_transaction.nil? || current_transaction.is_a?(ConnectionAdapters::NullTransaction))
     end
 
+    def self.unwrap_attribute attr_or_value
+      if attr_or_value.is_a? ActiveModel::Attribute
+        attr_or_value.value
+      else
+        attr_or_value
+      end
+    end
+
     # Updates the given attributes of the object in 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.
@@ -117,6 +229,7 @@ module ActiveRecord
       serialized_values = []
       columns = []
       values.each_pair do |k, v|
+        v = unwrap_attribute v
         type = metadata.type k
         serialized_values << (type.method(:serialize).arity < 0 ? type.serialize(v, :mutation) : type.serialize(v))
         columns << metadata.arel_table[k].name
@@ -168,6 +281,7 @@ module ActiveRecord
       all_values.each do |h|
         h.each_pair do |k, v|
           type = metadata.type k
+          v = self.class.unwrap_attribute v
           has_serialize_options = type.method(:serialize).arity < 0
           all_serialized_values << (has_serialize_options ? type.serialize(v, :mutation) : type.serialize(v))
           all_columns << metadata.arel_table[k].name
@@ -222,15 +336,34 @@ module ActiveRecord
       serialized_values
     end
 
-    def _execute_version_check attempted_action
+    def _execute_version_check attempted_action # rubocop:disable Metrics/AbcSize
       locking_column = self.class.locking_column
       previous_lock_value = read_attribute_before_type_cast locking_column
 
+      primary_key = self.class.primary_key
+      if primary_key.is_a? Array
+        pk_sql = ""
+        params = {}
+        param_types = {}
+        id = id_in_database
+        primary_key.each_with_index do |col, idx|
+          pk_sql.concat "`#{col}`=@id#{idx}"
+          pk_sql.concat " AND " if idx < primary_key.length - 1
+
+          params["id#{idx}"] = id[idx]
+          param_types["id#{idx}"] = :INT64
+        end
+        params["lock_version"] = previous_lock_value
+        param_types["lock_version"] = :INT64
+      else
+        pk_sql = "`#{self.class.primary_key}` = @id"
+        params = { "id" => id_in_database, "lock_version" => previous_lock_value }
+        param_types = { "id" => :INT64, "lock_version" => :INT64 }
+      end
+
       # We need to check the version using a SELECT query, as a mutation cannot include a WHERE clause.
       sql = "SELECT 1 FROM `#{self.class.arel_table.name}` " \
-              "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 }
+              "WHERE #{pk_sql} AND `#{locking_column}` = @lock_version"
       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

data/lib/activerecord_spanner_adapter/connection.rb

@@ -251,7 +251,7 @@ module ActiveRecordSpannerAdapter
     # 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
+      Google::Cloud::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

data/lib/activerecord_spanner_adapter/relation.rb

@@ -0,0 +1,21 @@
+# Copyright 2022 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.
+
+module ActiveRecord
+  module CpkExtension
+    def cpk_subquery stmt
+      return super unless spanner_adapter?
+      # The composite_primary_key gem will by default generate WHERE clauses using an IN clause with a multi-column
+      # sub select, e.g.: SELECT * FROM my_table WHERE (id1, id2) IN (SELECT id1, id2 FROM my_table WHERE ...).
+      # This is not supported in Cloud Spanner. Instead, composite_primary_key should generate an EXISTS clause.
+      cpk_exists_subquery stmt
+    end
+  end
+
+  class Relation
+    prepend CpkExtension
+  end
+end

data/lib/activerecord_spanner_adapter/transaction.rb

@@ -55,9 +55,9 @@ module ActiveRecordSpannerAdapter
         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
+          @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
             )
 
         end
@@ -143,7 +143,7 @@ module ActiveRecordSpannerAdapter
 
       # 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 \
+      return Google::Cloud::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

data/lib/activerecord_spanner_adapter/version.rb

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

data/lib/arel/visitors/spanner.rb

@@ -98,6 +98,16 @@ module Arel # :nodoc: all
         end
       end
 
+      # For ActiveRecord 7.0
+      def visit_ActiveModel_Attribute 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.
+        return collector << "PENDING_COMMIT_TIMESTAMP()" \
+          if o.type.is_a?(ActiveRecord::Type::Spanner::Time) && o.value == :commit_timestamp
+        collector.add_bind(o, &bind_block)
+      end
+
+      # For ActiveRecord 6.x
       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.

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-spanner-adapter
 version: !ruby/object:Gem::Version
-  version: 1.0.1
+  version: 1.2.1
 platform: ruby
 authors:
 - Google LLC
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2022-04-21 00:00:00.000000000 Z
+date: 2022-08-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: google-cloud-spanner
@@ -28,16 +28,22 @@ dependencies:
   name: activerecord
   requirement: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 6.0.0
+    - - "<"
       - !ruby/object:Gem::Version
-        version: 6.1.4
+        version: '7.1'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
-    - - "~>"
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 6.0.0
+    - - "<"
       - !ruby/object:Gem::Version
-        version: 6.1.4
+        version: '7.1'
 - !ruby/object:Gem::Dependency
   name: autotest-suffix
   requirement: !ruby/object:Gem::Requirement
@@ -264,6 +270,7 @@ files:
 - acceptance/cases/migration/rename_column_test.rb
 - acceptance/cases/models/calculation_query_test.rb
 - acceptance/cases/models/generated_column_test.rb
+- acceptance/cases/models/insert_all_test.rb
 - acceptance/cases/models/mutation_test.rb
 - acceptance/cases/models/query_test.rb
 - acceptance/cases/sessions/session_not_found_test.rb
@@ -380,6 +387,16 @@ files:
 - examples/snippets/hints/db/seeds.rb
 - examples/snippets/hints/models/album.rb
 - examples/snippets/hints/models/singer.rb
+- examples/snippets/interleaved-tables/README.md
+- examples/snippets/interleaved-tables/Rakefile
+- examples/snippets/interleaved-tables/application.rb
+- examples/snippets/interleaved-tables/config/database.yml
+- examples/snippets/interleaved-tables/db/migrate/01_create_tables.rb
+- examples/snippets/interleaved-tables/db/schema.rb
+- examples/snippets/interleaved-tables/db/seeds.rb
+- examples/snippets/interleaved-tables/models/album.rb
+- examples/snippets/interleaved-tables/models/singer.rb
+- examples/snippets/interleaved-tables/models/track.rb
 - examples/snippets/migrations/README.md
 - examples/snippets/migrations/Rakefile
 - examples/snippets/migrations/application.rb
@@ -485,6 +502,7 @@ files:
 - lib/activerecord_spanner_adapter/index/column.rb
 - lib/activerecord_spanner_adapter/information_schema.rb
 - lib/activerecord_spanner_adapter/primary_key.rb
+- lib/activerecord_spanner_adapter/relation.rb
 - lib/activerecord_spanner_adapter/table.rb
 - lib/activerecord_spanner_adapter/table/column.rb
 - lib/activerecord_spanner_adapter/transaction.rb
@@ -512,7 +530,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.5
+rubygems_version: 3.3.14
 signing_key: 
 specification_version: 4
 summary: Rails ActiveRecord connector for Google Spanner Database