activerecord-spanner-adapter 1.5.0 → 1.6.0

data/examples/snippets/query-logs/db/schema.rb

@@ -0,0 +1,31 @@
+# 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[7.1].define(version: 1) do
+  connection.start_batch_ddl
+
+  create_table "albums", force: :cascade do |t|
+    t.string "title"
+    t.integer "singer_id", limit: 8
+  end
+
+  create_table "singers", force: :cascade do |t|
+    t.string "first_name"
+    t.string "last_name"
+  end
+
+  add_foreign_key "albums", "singers"
+  connection.run_batch
+rescue
+  abort_batch
+  raise
+end

data/examples/snippets/query-logs/db/seeds.rb

@@ -0,0 +1,24 @@
+# Copyright 2023 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 = ["Pete", "Alice", "John", "Ethel", "Trudy", "Naomi", "Wendy", "Ruben", "Thomas", "Elly"]
+last_names = ["Wendelson", "Allison", "Peterson", "Johnson", "Henderson", "Ericsson", "Aronson", "Tennet", "Courtou"]
+
+adjectives = ["daily", "happy", "blue", "generous", "cooked", "bad", "open"]
+nouns = ["windows", "potatoes", "bank", "street", "tree", "glass", "bottle"]
+
+5.times do
+  Singer.create first_name: first_names.sample, last_name: last_names.sample
+end
+
+20.times do
+  singer_id = Singer.all.sample.id
+  Album.create title: "#{adjectives.sample} #{nouns.sample}", singer_id: singer_id
+end

data/examples/snippets/query-logs/models/album.rb

@@ -0,0 +1,9 @@
+# Copyright 2023 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/query-logs/models/singer.rb

@@ -0,0 +1,9 @@
+# Copyright 2023 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/column.rb

@@ -10,6 +10,15 @@ module ActiveRecord
   module ConnectionAdapters
     module Spanner
       class Column < ConnectionAdapters::Column
+        # rubocop:disable Style/MethodDefParentheses
+        def initialize(name, default, sql_type_metadata = nil, null = true,
+                       default_function = nil, collation: nil, comment: nil,
+                       primary_key: false, **)
+          # rubocop:enable Style/MethodDefParentheses
+          super
+          @primary_key = primary_key
+        end
+
         def has_default? # rubocop:disable Naming/PredicateName
           super && !virtual?
         end
@@ -17,6 +26,10 @@ module ActiveRecord
         def virtual?
           sql_type_metadata.generated
         end
+
+        def primary_key?
+          @primary_key
+        end
       end
     end
   end

data/lib/active_record/connection_adapters/spanner/database_statements.rb

@@ -6,14 +6,34 @@
 
 # frozen_string_literal: true
 
+require "active_record/gem_version"
+
 module ActiveRecord
   module ConnectionAdapters
     module Spanner
       module DatabaseStatements
+        VERSION_7_1_0 = Gem::Version.create "7.1.0"
+        RequestOptions = Google::Cloud::Spanner::V1::RequestOptions
+
         # DDL, DML and DQL Statements
 
         def execute sql, name = nil, binds = []
+          internal_execute sql, name, binds
+        end
+
+        def internal_exec_query sql, name = "SQL", binds = [], prepare: false, async: false
+          result = internal_execute sql, name, binds, prepare: prepare, async: async
+          ActiveRecord::Result.new(
+            result.fields.keys.map(&:to_s), result.rows.map(&:values)
+          )
+        end
+
+        def internal_execute sql, name = "SQL", binds = [],
+                             prepare: false, async: false # rubocop:disable Lint/UnusedMethodArgument
           statement_type = sql_statement_type sql
+          # Call `transform` to invoke any query transformers that might have been registered.
+          sql = transform sql
+          append_request_tag_from_query_logs sql, binds
 
           if preventing_writes? && [:dml, :ddl].include?(statement_type)
             raise ActiveRecord::ReadOnlyError(
@@ -24,49 +44,127 @@ module ActiveRecord
           if statement_type == :ddl
             execute_ddl sql
           else
-            transaction_required = statement_type == :dml
-            materialize_transactions
-
-            # First process and remove any hints in the binds that indicate that
-            # a different read staleness should be used than the default.
-            staleness_hint = binds.find { |b| b.is_a? Arel::Visitors::StalenessHint }
-            if staleness_hint
-              selector = Google::Cloud::Spanner::Session.single_use_transaction staleness_hint.value
-              binds.delete staleness_hint
-            end
-            request_options = binds.find { |b| b.is_a? Google::Cloud::Spanner::V1::RequestOptions }
-            if request_options
-              binds.delete request_options
-            end
+            execute_query_or_dml statement_type, sql, name, binds
+          end
+        end
+
+        def execute_query_or_dml statement_type, sql, name, binds
+          transaction_required = statement_type == :dml
+          materialize_transactions
+
+          # First process and remove any hints in the binds that indicate that
+          # a different read staleness should be used than the default.
+          staleness_hint = binds.find { |b| b.is_a? Arel::Visitors::StalenessHint }
+          if staleness_hint
+            selector = Google::Cloud::Spanner::Session.single_use_transaction staleness_hint.value
+            binds.delete staleness_hint
+          end
+          request_options = binds.find { |b| b.is_a? RequestOptions }
+          if request_options
+            binds.delete request_options
+          end
+
+          log_args = [sql, name]
+          log_args.concat [binds, type_casted_binds(binds)] if log_statement_binds
 
-            log_args = [sql, name]
-            log_args.concat [binds, type_casted_binds(binds)] if log_statement_binds
-
-            log(*log_args) do
-              types, params = to_types_and_params binds
-              ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
-                if transaction_required
-                  transaction do
-                    @connection.execute_query sql, params: params, types: types, request_options: request_options
-                  end
-                else
-                  @connection.execute_query sql, params: params, types: types, single_use_selector: selector,
-                                            request_options: request_options
+          log(*log_args) do
+            types, params = to_types_and_params binds
+            ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
+              if transaction_required
+                transaction do
+                  @connection.execute_query sql, params: params, types: types, request_options: request_options
                 end
+              else
+                @connection.execute_query sql, params: params, types: types, single_use_selector: selector,
+                                          request_options: request_options
               end
             end
           end
         end
 
-        def query sql, name = nil
-          exec_query sql, name
+        def append_request_tag_from_query_logs sql, binds
+          legacy_formatter_prefix = "/*request_tag:true,"
+          sql_commenter_prefix = "/*request_tag='true',"
+          if sql.start_with? legacy_formatter_prefix
+            append_request_tag_from_query_logs_with_format sql, binds, legacy_formatter_prefix
+          elsif sql.start_with? sql_commenter_prefix
+            append_request_tag_from_query_logs_with_format sql, binds, sql_commenter_prefix
+          end
         end
 
-        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)
-          )
+        def append_request_tag_from_query_logs_with_format sql, binds, prefix
+          end_of_comment = sql.index "*/", prefix.length
+          return unless end_of_comment
+
+          request_tag = sql[prefix.length, end_of_comment - prefix.length]
+          options = binds.find { |bind| bind.is_a? RequestOptions } || RequestOptions.new
+          if options.request_tag == ""
+            options.request_tag = request_tag
+          else
+            options.request_tag += "," + request_tag
+          end
+
+          binds.append options
+        end
+
+        # The method signatures for executing queries and DML statements changed between Rails 7.0 and 7.1.
+
+        if ActiveRecord.gem_version >= VERSION_7_1_0
+          def sql_for_insert sql, pk, binds, returning
+            if pk && !_has_pk_binding(pk, binds)
+              # Add the primary key to the columns that should be returned if there is no value specified for it.
+              returning ||= []
+              returning |= if pk.respond_to? :each
+                             pk
+                           else
+                             [pk]
+                           end
+            end
+            if returning&.any?
+              returning_columns_statement = returning.map { |c| quote_column_name c }.join(", ")
+              sql = "#{sql} THEN RETURN #{returning_columns_statement}"
+            end
+
+            [sql, binds]
+          end
+
+          def query sql, name = nil
+            exec_query sql, name
+          end
+        else # ActiveRecord.gem_version < VERSION_7_1_0
+          def query sql, name = nil
+            exec_query sql, name
+          end
+
+          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)
+            )
+          end
+
+          def sql_for_insert sql, pk, binds
+            if pk && !_has_pk_binding(pk, binds)
+              # Add the primary key to the columns that should be returned if there is no value specified for it.
+              returning_columns_statement = if pk.respond_to? :each
+                                              pk.map { |c| quote_column_name c }.join(", ")
+                                            else
+                                              quote_column_name pk
+                                            end
+              sql = "#{sql} THEN RETURN #{returning_columns_statement}" if returning_columns_statement
+            end
+            super
+          end
+        end
+
+        def _has_pk_binding pk, binds
+          if pk.respond_to? :each
+            has_value = true
+            pk.each { |col| has_value &&= binds.any? { |bind| bind.name == col } }
+            has_value
+          else
+            binds.any? { |bind| bind.name == pk }
+          end
         end
 
         def exec_mutation mutation
@@ -219,6 +317,9 @@ module ActiveRecord
             if bind.respond_to? :type
               type = ActiveRecord::Type::Spanner::SpannerActiveRecordConverter
                      .convert_active_model_type_to_spanner(bind.type)
+            elsif bind.class == Symbol
+              # This ensures that for example :environment is sent as the string 'environment' to Cloud Spanner.
+              type = :STRING
             end
             [
               # Generates binds for named parameters in the format `@p1, @p2, ...`
@@ -226,7 +327,15 @@ module ActiveRecord
             ]
           end.to_h
           params = binds.enum_for(:each_with_index).map do |bind, i|
-            type = bind.respond_to?(:type) ? bind.type : ActiveModel::Type::Integer
+            type = if bind.respond_to? :type
+                     bind.type
+                   elsif bind.class == Symbol
+                     # This ensures that for example :environment is sent as the string 'environment' to Cloud Spanner.
+                     :STRING
+                   else
+                     # The Cloud Spanner default type is INT64 if no other type is known.
+                     ActiveModel::Type::Integer
+                   end
             bind_value = bind.respond_to?(:value) ? bind.value : bind
             value = ActiveRecord::Type::Spanner::SpannerActiveRecordConverter
                     .serialize_with_transaction_isolation_level(type, bind_value, :dml)

data/lib/active_record/connection_adapters/spanner/schema_cache.rb

@@ -6,37 +6,19 @@
 
 module ActiveRecord
   module ConnectionAdapters
-    class SpannerSchemaCache < SchemaCache
+    class SpannerSchemaCache
       def initialize conn
+        @connection = conn
         @primary_and_parent_keys = {}
-        super
-      end
-
-      def initialize_dup other
-        @primary_and_parent_keys = @primary_and_parent_keys.dup
-        super
-      end
-
-      def encode_with coder
-        coder["primary_and_parent_keys"] = @primary_and_parent_keys
-        super
-      end
-
-      def init_with coder
-        @primary_and_parent_keys = coder["primary_and_parent_keys"]
-        super
       end
 
       def primary_and_parent_keys table_name
         @primary_and_parent_keys[table_name] ||=
-          if data_source_exists? table_name
-            connection.primary_and_parent_keys table_name
-          end
+          @connection.primary_and_parent_keys table_name
       end
 
       def clear!
         @primary_and_parent_keys.clear
-        super
       end
     end
   end

data/lib/active_record/connection_adapters/spanner/schema_creation.rb

@@ -39,6 +39,17 @@ module ActiveRecord
 
           primary_keys = if o.primary_keys
                            o.primary_keys
+                         elsif o.options[:primary_key]
+                           columns = if o.options[:primary_key].is_a? Array
+                                       o.options[:primary_key]
+                                     else
+                                       [o.options[:primary_key]]
+                                     end
+                           pk_names = []
+                           columns.each do |c|
+                             pk_names.append c.to_s
+                           end
+                           PrimaryKeyDefinition.new pk_names
                          else
                            pk_names = o.columns.each_with_object [] do |c, r|
                              if c.type == :primary_key || c.primary_key?

data/lib/active_record/connection_adapters/spanner/schema_definitions.rb

@@ -42,6 +42,10 @@ module ActiveRecord
 
           super
         end
+
+        def valid_column_definition_options
+          super + [:type, :array, :allow_commit_timestamp, :as, :stored, :parent_key, :passed_type, :index]
+        end
       end
 
       class Table < ActiveRecord::ConnectionAdapters::Table

data/lib/active_record/connection_adapters/spanner/schema_statements.rb

@@ -109,7 +109,7 @@ module ActiveRecord
           information_schema { |i| i.table_columns table_name }
         end
 
-        def new_column_from_field _table_name, field
+        def new_column_from_field _table_name, field, _definitions = nil
           Spanner::Column.new \
             field.name,
             field.default,
@@ -118,7 +118,8 @@ module ActiveRecord
                                 field.allow_commit_timestamp,
                                 field.generated),
             field.nullable,
-            field.default_function
+            field.default_function,
+            primary_key: field.primary_key
         end
 
         def fetch_type_metadata sql_type, ordinal_position = nil, allow_commit_timestamp = nil, generated = nil

data/lib/active_record/connection_adapters/spanner_adapter.rb

@@ -42,14 +42,6 @@ module ActiveRecord
   end
 
   module ConnectionAdapters
-    module AbstractPool
-      def get_schema_cache connection
-        self.schema_cache ||= SpannerSchemaCache.new connection
-        schema_cache.connection = connection
-        schema_cache
-      end
-    end
-
     class SpannerAdapter < AbstractAdapter
       ADAPTER_NAME = "spanner".freeze
       NATIVE_DATABASE_TYPES = {
@@ -78,8 +70,10 @@ module ActiveRecord
       class_attribute :log_statement_binds, instance_writer: false, default: false
 
       def initialize connection, logger, connection_options, config
-        super connection, logger, config
+        @connection = connection
         @connection_options = connection_options
+        super connection, logger, config
+        @raw_connection ||= connection
       end
 
       def max_identifier_length
@@ -117,6 +111,10 @@ module ActiveRecord
       end
       alias reconnect! reset!
 
+      def spanner_schema_cache
+        @spanner_schema_cache ||= SpannerSchemaCache.new self
+      end
+
       # Spanner Connection API
       delegate :ddl_batch, :ddl_batch?, :start_batch_ddl, :abort_batch, :run_batch, to: :@connection
 
@@ -186,6 +184,10 @@ module ActiveRecord
         SecureRandom.uuid.gsub("-", "").hex & 0x7FFFFFFFFFFFFFFF
       end
 
+      def return_value_after_insert? column
+        column.auto_incremented_by_db? || column.primary_key?
+      end
+
       def arel_visitor
         Arel::Visitors::Spanner.new self
       end
@@ -259,6 +261,23 @@ module ActiveRecord
         include TypeMapBuilder
       end
 
+      def transform sql
+        if ActiveRecord::VERSION::MAJOR >= 7
+          transform_query sql
+        else
+          sql
+        end
+      end
+
+      # Overwrite the standard log method to be able to translate exceptions.
+      def log sql, name = "SQL", binds = [], type_casted_binds = [], statement_name = nil, *args
+        super
+      rescue ActiveRecord::StatementInvalid
+        raise
+      rescue StandardError => e
+        raise translate_exception_class(e, sql, binds)
+      end
+
       def translate_exception exception, message:, sql:, binds:
         if exception.is_a? Google::Cloud::FailedPreconditionError
           case exception.message

data/lib/activerecord_spanner_adapter/base.rb

@@ -4,6 +4,7 @@
 # license that can be found in the LICENSE file or at
 # https://opensource.org/licenses/MIT.
 
+require "active_record/gem_version"
 require "activerecord_spanner_adapter/relation"
 
 module ActiveRecord
@@ -14,6 +15,8 @@ module ActiveRecord
   end
 
   class Base
+    VERSION_7_1 = Gem::Version.create "7.1.0"
+
     # 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.
@@ -30,7 +33,9 @@ module ActiveRecord
       return super unless spanner_adapter?
       return super if active_transaction?
 
-      transaction isolation: :buffered_mutations do
+      # Only use mutations to create new records if the primary key is generated client-side.
+      isolation = sequence_name ? nil : :buffered_mutations
+      transaction isolation: isolation do
         return super
       end
     end
@@ -43,35 +48,67 @@ module ActiveRecord
       spanner_adapter? && connection&.current_spanner_transaction&.isolation == :buffered_mutations
     end
 
-    def self._insert_record values
+    def self._insert_record values, returning = []
       return super unless buffered_mutations? || (primary_key && values.is_a?(Hash))
 
-      return _buffer_record values, :insert if buffered_mutations?
+      # Mutations cannot be used in combination with a sequence, as mutations do not support a THEN RETURN clause.
+      if buffered_mutations? && sequence_name
+        raise StatementInvalid, "Mutations cannot be used to create records that use a sequence " \
+                                     "to generate the primary key. #{self} uses #{sequence_name}."
+      end
 
-      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
+      return _buffer_record values, :insert, returning if buffered_mutations?
+
+      primary_key_value = _set_primary_key_value values
       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)
+      result = connection.insert(im, "#{self} Create", primary_key || false, primary_key_value)
+
+      _convert_primary_key result, returning
     end
 
-    def self._upsert_record values
-      _buffer_record values, :insert_or_update
+    def self._set_primary_key_value values
+      if primary_key.is_a? Array
+        _set_composite_primary_key_values primary_key, values
+      else
+        _set_single_primary_key_value primary_key, values
+      end
+    end
+
+    def self._convert_primary_key primary_key_value, returning
+      # Rails 7.1 and higher supports composite primary keys, and therefore require the provider to return an array
+      # instead of a single value in all cases. The order of the values should be equal to the order of the returning
+      # columns (or the primary key if no returning columns were specified).
+      return primary_key_value if ActiveRecord.gem_version < VERSION_7_1
+      return primary_key_value if primary_key_value.is_a?(Array) && primary_key_value.length == 1
+      return [primary_key_value] unless primary_key_value.is_a? Array
+
+      # Re-order the returned values according to the returning columns if it is not equal to the primary key of the
+      # table.
+      keys = returning || primary_key
+      return primary_key_value if keys == primary_key
+
+      primary_key_values_hash = Hash[primary_key.zip(primary_key_value)]
+      values = []
+      keys.each do |column|
+        values.append primary_key_values_hash[column]
+      end
+      values
+    end
+
+    def self._upsert_record values, returning
+      _buffer_record values, :insert_or_update, returning
     end
 
-    def self.insert_all _attributes, _returning: nil, _unique_by: nil
-      raise NotImplementedError, "Cloud Spanner does not support skip_duplicates."
+    def self.insert_all _attributes, **_kwargs
+      raise NotImplementedError, "Cloud Spanner does not support skip_duplicates. Use insert! or upsert instead."
     end
 
-    def self.insert_all! attributes, returning: nil
+    def self.insert_all! attributes, returning: nil, **_kwargs
       return super unless spanner_adapter?
       return super if active_transaction? && !buffered_mutations?
 
@@ -90,7 +127,7 @@ module ActiveRecord
       end
     end
 
-    def self.upsert_all attributes, returning: nil, unique_by: nil
+    def self.upsert_all attributes, returning: nil, unique_by: nil, **_kwargs
       return super unless spanner_adapter?
       if active_transaction? && !buffered_mutations?
         raise NotImplementedError, "Cloud Spanner does not support upsert using DML. " \
@@ -102,18 +139,18 @@ module ActiveRecord
       # The mutations will be sent as one batch when the transaction is committed.
       if active_transaction?
         attributes.each do |record|
-          _upsert_record record
+          _upsert_record record, returning
         end
       else
         transaction isolation: :buffered_mutations do
           attributes.each do |record|
-            _upsert_record record
+            _upsert_record record, returning
           end
         end
       end
     end
 
-    def self._buffer_record values, method
+    def self._buffer_record values, method, returning
       primary_key_value =
         if primary_key.is_a? Array
           _set_composite_primary_key_values primary_key, values
@@ -132,10 +169,9 @@ module ActiveRecord
       mutation = Google::Cloud::Spanner::V1::Mutation.new(
         "#{method}": write
       )
-
       connection.current_spanner_transaction.buffer mutation
 
-      primary_key_value
+      _convert_primary_key primary_key_value, returning
     end
 
     def self._set_composite_primary_key_values primary_key, values
@@ -174,6 +210,7 @@ module ActiveRecord
     def self._set_single_primary_key_value primary_key, values
       primary_key_value = values[primary_key] || values[primary_key.to_sym]
 
+      return primary_key_value if sequence_name
       return primary_key_value unless prefetch_primary_key?
 
       if primary_key_value.nil?