activerecord-spanner-adapter 1.8.0 → 2.1.0

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

@@ -8,6 +8,11 @@ module ActiveRecord
   module ConnectionAdapters
     module Spanner
       class SchemaCreation < SchemaCreation
+        def initialize connection
+          super
+          @connection = connection
+        end
+
         private
 
         # rubocop:disable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
@@ -44,10 +49,7 @@ module ActiveRecord
                                      else
                                        [o.options[:primary_key]]
                                      end
-                           pk_names = []
-                           columns.each do |c|
-                             pk_names.append c.to_s
-                           end
+                           pk_names = columns.map(&:to_s)
                            PrimaryKeyDefinition.new pk_names
                          else
                            pk_names = o.columns.each_with_object [] do |c, r|
@@ -93,8 +95,8 @@ module ActiveRecord
         end
 
         def visit_DropColumnDefinition o
-          "ALTER TABLE #{quote_table_name o.table_name} DROP" \
-           " COLUMN #{quote_column_name o.name}"
+          "ALTER TABLE #{quote_table_name o.table_name} DROP " \
+            "COLUMN #{quote_column_name o.name}"
         end
 
         def visit_ChangeColumnDefinition o
@@ -131,19 +133,23 @@ module ActiveRecord
           sql
         end
 
-        # rubocop:enable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
-
         def add_column_options! column, sql, options
           if options[:null] == false || options[:primary_key] == true
             sql << " NOT NULL"
           end
           if options.key? :default
             sql << " DEFAULT (#{quote_default_expression options[:default], column})"
+          elsif column.type == :primary_key
+            if @connection.use_auto_increment?
+              sql << " AUTO_INCREMENT"
+            elsif @connection.use_identity?
+              sql << " GENERATED BY DEFAULT AS IDENTITY (#{@connection.default_sequence_kind})"
+            end
           end
 
           if !options[:allow_commit_timestamp].nil? &&
              options[:column].sql_type == "TIMESTAMP"
-            sql << " OPTIONS (allow_commit_timestamp = "\
+            sql << " OPTIONS (allow_commit_timestamp = " \
                    "#{options[:allow_commit_timestamp]})"
           end
 
@@ -153,13 +159,16 @@ module ActiveRecord
             sql << " STORED" if options[:stored]
             unless options[:stored]
               raise ArgumentError, "" \
-                "Cloud Spanner currently does not support generated columns without the STORED option." \
-                "Specify 'stored: true' option for `#{options[:column].name}`"
+                                   "Cloud Spanner currently does not support generated columns" \
+                                   "without the STORED option." \
+                                   "Specify 'stored: true' option for `#{options[:column].name}`"
             end
           end
 
           sql
         end
+
+        # rubocop:enable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
       end
     end
   end

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

@@ -5,7 +5,7 @@
 # https://opensource.org/licenses/MIT.
 
 module ActiveRecord
-  module ConnectionAdapters #:nodoc:
+  module ConnectionAdapters # :nodoc:
     module Spanner
       class TableDefinition < ActiveRecord::ConnectionAdapters::TableDefinition
         attr_reader :interleave_in_parent
@@ -62,6 +62,9 @@ module ActiveRecord
       DropIndexDefinition = Struct.new :name
 
       class ReferenceDefinition < ActiveRecord::ConnectionAdapters::ReferenceDefinition
+        # This constructor intentionally does not call super to prevent ActiveRecord
+        # from creating an additional secondary index for the foreign key.
+        # rubocop:disable Lint/MissingSuper
         def initialize \
             name,
             polymorphic: false,
@@ -81,6 +84,7 @@ module ActiveRecord
           return unless polymorphic && foreign_key
           raise ArgumentError, "Cannot add a foreign key to a polymorphic relation"
         end
+        # rubocop:enable Lint/MissingSuper
 
         private
 
@@ -96,8 +100,13 @@ module ActiveRecord
       end
 
       class IndexDefinition < ActiveRecord::ConnectionAdapters::IndexDefinition
-        attr_reader :null_filtered, :interleave_in, :storing, :orders
+        attr_reader :null_filtered
+        attr_reader :interleave_in
+        attr_reader :storing
+        attr_reader :orders
 
+        # This constructor intentionally does not call super.
+        # rubocop:disable Lint/MissingSuper
         def initialize \
             table_name,
             name,
@@ -123,6 +132,7 @@ module ActiveRecord
 
           @orders = @orders.symbolize_keys
         end
+        # rubocop:enable Lint/MissingSuper
 
         def columns_with_order
           columns.each_with_object({}) do |c, result|

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

@@ -21,8 +21,6 @@ module ActiveRecord
       # [Schema Doc](https://cloud.google.com/spanner/docs/information-schema)
       #
       module SchemaStatements
-        VERSION_6_1_0 = Gem::Version.create "6.1.0"
-        VERSION_6_0_3 = Gem::Version.create "6.0.3"
         VERSION_7_2 = Gem::Version.create "7.2.0"
 
         def current_database
@@ -109,8 +107,7 @@ module ActiveRecord
         end
 
         def rename_table _table_name, _new_name
-          raise ActiveRecordSpannerAdapter::NotSupportedError, \
-                "rename_table is not implemented"
+          raise ActiveRecordSpannerAdapter::NotSupportedError, "rename_table is not implemented"
         end
 
         # Column
@@ -126,18 +123,20 @@ module ActiveRecord
             fetch_type_metadata(field.spanner_type,
                                 field.ordinal_position,
                                 field.allow_commit_timestamp,
-                                field.generated),
+                                field.generated,
+                                is_identity: field.is_identity),
             field.nullable,
             field.default_function,
             primary_key: field.primary_key
         end
 
-        def fetch_type_metadata sql_type, ordinal_position = nil, allow_commit_timestamp = nil, generated = nil
+        def fetch_type_metadata sql_type, ordinal_position = nil, allow_commit_timestamp = nil, generated = nil,
+                                is_identity: false
           Spanner::TypeMetadata.new \
             super(sql_type),
             ordinal_position: ordinal_position,
             allow_commit_timestamp: allow_commit_timestamp,
-            generated: generated
+            generated: generated, is_identity: is_identity
         end
 
         def add_column table_name, column_name, type, **options
@@ -168,20 +167,14 @@ module ActiveRecord
           execute_schema_statements statements
         end
 
-        if ActiveRecord.gem_version < VERSION_6_1_0
-          def remove_columns table_name, *column_names
-            _remove_columns table_name, *column_names
-          end
-        else
-          def remove_columns table_name, *column_names, _type: nil, **_options
-            _remove_columns table_name, *column_names
-          end
+        def remove_columns table_name, *column_names, _type: nil, **_options
+          _remove_columns table_name, *column_names
         end
 
         def _remove_columns table_name, *column_names
           if column_names.empty?
-            raise ArgumentError, "You must specify at least one column name. "\
-              "Example: remove_columns(:people, :first_name)"
+            raise ArgumentError, "You must specify at least one column name. " \
+                                 "Example: remove_columns(:people, :first_name)"
           end
 
           statements = []
@@ -193,14 +186,8 @@ module ActiveRecord
           execute_schema_statements statements
         end
 
-        if ActiveRecord.gem_version < VERSION_6_1_0
-          def change_column table_name, column_name, type, options = {}
-            _change_column table_name, column_name, type, **options
-          end
-        else
-          def change_column table_name, column_name, type, **options
-            _change_column table_name, column_name, type, **options
-          end
+        def change_column table_name, column_name, type, **options
+          _change_column table_name, column_name, type, **options
         end
 
         def change_column_null table_name, column_name, null, _default = nil
@@ -208,14 +195,12 @@ module ActiveRecord
         end
 
         def change_column_default _table_name, _column_name, _default_or_changes
-          raise ActiveRecordSpannerAdapter::NotSupportedError, \
-                "change column with default value not supported."
+          raise ActiveRecordSpannerAdapter::NotSupportedError, "change column with default value not supported."
         end
 
         def rename_column table_name, column_name, new_column_name
           if ActiveRecord::Base.connection.ddl_batch?
-            raise ActiveRecordSpannerAdapter::NotSupportedError, \
-                  "rename_column in a DDL Batch is not supported."
+            raise ActiveRecordSpannerAdapter::NotSupportedError, "rename_column in a DDL Batch is not supported."
           end
           column = information_schema do |i|
             i.table_column table_name, column_name
@@ -280,16 +265,9 @@ module ActiveRecord
           execute_schema_statements schema_creation.accept(id)
         end
 
-        if ActiveRecord.gem_version < VERSION_6_1_0
-          def remove_index table_name, options = {}
-            index_name = index_name_for_remove table_name, options
-            execute "DROP INDEX #{quote_table_name index_name}"
-          end
-        else
-          def remove_index table_name, column_name = nil, **options
-            index_name = index_name_for_remove table_name, column_name, options
-            execute "DROP INDEX #{quote_table_name index_name}"
-          end
+        def remove_index table_name, column_name = nil, **options
+          index_name = index_name_for_remove table_name, column_name, options
+          execute "DROP INDEX #{quote_table_name index_name}"
         end
 
         def rename_index table_name, old_name, new_name
@@ -358,14 +336,8 @@ module ActiveRecord
           end
         end
 
-        if ActiveRecord.gem_version < VERSION_6_0_3
-          def add_foreign_key from_table, to_table, options = {}
-            _add_foreign_key from_table, to_table, **options
-          end
-        else
-          def add_foreign_key from_table, to_table, **options
-            _add_foreign_key from_table, to_table, **options
-          end
+        def add_foreign_key from_table, to_table, **options
+          _add_foreign_key from_table, to_table, **options
         end
 
         def _add_foreign_key from_table, to_table, **options
@@ -508,8 +480,8 @@ module ActiveRecord
           type ||= column.type
           options[:null] = column.null unless options.key? :null
 
-          if ["STRING", "BYTES"].include? type
-            options[:limit] = column.limit unless options.key? :limit
+          if ["STRING", "BYTES"].include?(type) && !options.key?(:limit)
+            options[:limit] = column.limit
           end
 
           # Only timestamp type can set commit timestamp
@@ -642,8 +614,7 @@ module ActiveRecord
         end
 
         def information_schema
-          info_schema = \
-            ActiveRecordSpannerAdapter::Connection.information_schema @config
+          info_schema = ActiveRecordSpannerAdapter::Connection.information_schema @config
 
           return info_schema unless block_given?
 

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

@@ -14,13 +14,18 @@ module ActiveRecord
 
         include Deduplicable if defined?(Deduplicable)
 
-        attr_reader :ordinal_position, :allow_commit_timestamp, :generated
+        attr_reader :ordinal_position
+        attr_reader :allow_commit_timestamp
+        attr_reader :generated
+        attr_reader :is_identity
 
-        def initialize type_metadata, ordinal_position: nil, allow_commit_timestamp: nil, generated: nil
+        def initialize type_metadata, ordinal_position: nil, allow_commit_timestamp: nil, generated: nil,
+                       is_identity: false
           super type_metadata
           @ordinal_position = ordinal_position
           @allow_commit_timestamp = allow_commit_timestamp
           @generated = generated
+          @is_identity = is_identity
         end
 
         def == other
@@ -28,16 +33,13 @@ module ActiveRecord
             __getobj__ == other.__getobj__ &&
             ordinal_position == other.ordinal_position &&
             allow_commit_timestamp == other.allow_commit_timestamp &&
-            generated == other.generated
+            generated == other.generated &&
+            is_identity == other.is_identity
         end
         alias eql? ==
 
         def hash
-          TypeMetadata.hash ^
-            __getobj__.hash ^
-            ordinal_position.hash ^
-            allow_commit_timestamp.hash ^
-            generated.hash
+          [TypeMetadata.name, __getobj__, ordinal_position, allow_commit_timestamp, generated, is_identity].hash
         end
 
         private

data/lib/active_record/connection_adapters/spanner_adapter.rb

@@ -71,6 +71,9 @@ module ActiveRecord
       # Determines whether or not to log query binds when executing statements
       class_attribute :log_statement_binds, instance_writer: false, default: false
 
+      attr_accessor :default_sequence_kind
+      attr_accessor :use_client_side_id_for_mutations
+
       def initialize config_or_deprecated_connection, deprecated_logger = nil,
                      deprecated_connection_options = nil, deprecated_config = nil
         if config_or_deprecated_connection.is_a? Hash
@@ -86,6 +89,25 @@ module ActiveRecord
         end
         # Spanner does not support unprepared statements
         @prepared_statements = true
+        # The default for default_sequence_kind will be changed to BIT_REVERSED_POSITIVE
+        # in the next major version. The default value is currently DISABLED to prevent
+        # breaking changes to existing code.
+        @default_sequence_kind = @config.fetch :default_sequence_kind, "DISABLED"
+        @use_auto_increment = @default_sequence_kind&.casecmp? "AUTO_INCREMENT"
+        @auto_increment_disabled = @default_sequence_kind&.casecmp? "DISABLED"
+        @use_client_side_id_for_mutations = self.class.type_cast_config_to_boolean(
+          @config.fetch(:use_client_side_id_for_mutations, false)
+        )
+      end
+
+      def use_auto_increment?
+        "AUTO_INCREMENT".casecmp?(@default_sequence_kind || "")
+      end
+
+      def use_identity?
+        !use_auto_increment? \
+        && @default_sequence_kind \
+        && !@default_sequence_kind.casecmp?("DISABLED")
       end
 
       def max_identifier_length
@@ -276,7 +298,9 @@ module ActiveRecord
       end
 
       def transform sql
-        if ActiveRecord::VERSION::MAJOR >= 7
+        if ActiveRecord::VERSION::MAJOR >= 8
+          preprocess_query sql
+        elsif ActiveRecord::VERSION::MAJOR == 7
           transform_query sql
         else
           sql
@@ -284,12 +308,23 @@ module ActiveRecord
       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)
+      # sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &block
+      if ActiveRecord::VERSION::MAJOR >= 8
+        def log sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &block
+          super
+        rescue ActiveRecord::StatementInvalid
+          raise
+        rescue StandardError => e
+          raise translate_exception_class(e, sql, binds)
+        end
+      else
+        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
       end
 
       def translate_exception exception, message:, sql:, binds:

data/lib/active_record/tasks/spanner_database_tasks.rb

@@ -31,7 +31,7 @@ module ActiveRecord
       def purge
         begin
           drop
-        rescue ActiveRecord::NoDatabaseError # rubocop:disable Lint/HandleExceptions
+        rescue ActiveRecord::NoDatabaseError
           # ignored; create the database
         end
 
@@ -51,8 +51,8 @@ module ActiveRecord
         ignore_tables = ActiveRecord::SchemaDumper.ignore_tables
 
         if ignore_tables.any?
-          index_regx = /^CREATE(.*)INDEX(.*)ON (#{ignore_tables.join "|"})\(/
-          table_regx = /^CREATE TABLE (#{ignore_tables.join "|"})/
+          index_regx = /^CREATE(.*)INDEX(.*)ON (#{ignore_tables.join '|'})\(/
+          table_regx = /^CREATE TABLE (#{ignore_tables.join '|'})/
         end
 
         @connection.database.ddl(force: true).each do |statement|
@@ -66,7 +66,7 @@ module ActiveRecord
       end
 
       def structure_load filename, _extra_flags
-        statements = File.read(filename).split(/;/).map(&:strip).reject(&:empty?)
+        statements = File.read(filename).split(";").map(&:strip).reject(&:empty?)
         ddls = statements.select { |s| s =~ /^(CREATE|ALTER|DROP|GRANT|REVOKE|ANALYZE)/ }
         @connection.execute_ddl ddls
 

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

@@ -9,11 +9,15 @@ module ActiveRecord
     module Spanner
       class Array < Type::Value
         attr_reader :element_type
+
         delegate :type, :user_input_in_time_zone, :limit, :precision, :scale, to: :element_type
 
+        # This constructor intentionally does not call super.
+        # rubocop:disable Lint/MissingSuper
         def initialize element_type
           @element_type = element_type
         end
+        # rubocop:enable Lint/MissingSuper
 
         def cast value
           return super if value.nil?

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

@@ -10,6 +10,16 @@ module ActiveRecord
   module Type
     module Spanner
       class Bytes < ActiveRecord::Type::Binary
+        def deserialize value
+          # Set this environment variable to disable de-serializing BYTES
+          # to a StringIO instance.
+          return super if ENV["SPANNER_BYTES_DESERIALIZE_DISABLED"]
+
+          return super value if value.nil?
+          return StringIO.new Base64.strict_decode64(value) if value.is_a? ::String
+          value
+        end
+
         def serialize value
           return super value if value.nil?
 

data/lib/activerecord_spanner_adapter/base.rb

@@ -35,7 +35,7 @@ module ActiveRecord
       return super if active_transaction?
 
       # Only use mutations to create new records if the primary key is generated client-side.
-      isolation = sequence_name ? nil : :buffered_mutations
+      isolation = has_auto_generated_primary_key? ? nil : :buffered_mutations
       transaction isolation: isolation do
         return super
       end
@@ -53,6 +53,21 @@ module ActiveRecord
       !(buffered_mutations? || (primary_key && values.is_a?(Hash))) || !spanner_adapter?
     end
 
+    def self.has_auto_generated_primary_key?
+      return true if sequence_name
+      pk = primary_key
+      if pk.is_a? Array
+        return pk.any? do |col|
+          columns_hash[col].auto_incremented_by_db?
+        end
+      end
+      columns_hash[pk].auto_incremented_by_db?
+    end
+
+    def self.is_auto_generated? col
+      columns_hash[col]&.auto_incremented_by_db?
+    end
+
     def self._internal_insert_record values
       if ActiveRecord.gem_version < VERSION_7_2
         _insert_record values
@@ -73,10 +88,13 @@ module ActiveRecord
         return super
       end
 
-      # 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}."
+      # Mutations cannot be used in combination with an auto-generated primary key,
+      # as mutations do not support a THEN RETURN clause.
+      if buffered_mutations? \
+        && has_auto_generated_primary_key? \
+        && !_has_all_primary_key_values?(primary_key, values) \
+        && !connection.use_client_side_id_for_mutations
+        raise StatementInvalid, "Mutations cannot be used to create records that use an auto-generated primary key."
       end
 
       return _buffer_record values, :insert, returning if buffered_mutations?
@@ -85,7 +103,7 @@ module ActiveRecord
     end
 
     def self._insert_record_dml values, returning
-      primary_key_value = _set_primary_key_value values
+      primary_key_value = _set_primary_key_value values, false
       if ActiveRecord::VERSION::MAJOR >= 7
         im = Arel::InsertManager.new arel_table
         im.insert(values.transform_keys { |name| arel_table[name] })
@@ -97,11 +115,11 @@ module ActiveRecord
       _convert_primary_key result, returning
     end
 
-    def self._set_primary_key_value values
+    def self._set_primary_key_value values, is_mutation
       if primary_key.is_a? Array
-        _set_composite_primary_key_values primary_key, values
+        _set_composite_primary_key_values primary_key, values, is_mutation
       else
-        _set_single_primary_key_value primary_key, values
+        _set_single_primary_key_value primary_key, values, is_mutation
       end
     end
 
@@ -118,12 +136,10 @@ module ActiveRecord
       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]
+      primary_key_values_hash = primary_key.zip(primary_key_value).to_h
+      keys.map do |column|
+        primary_key_values_hash[column]
       end
-      values
     end
 
     def self._upsert_record values, returning
@@ -194,9 +210,9 @@ module ActiveRecord
     def self._buffer_record values, method, returning
       primary_key_value =
         if primary_key.is_a? Array
-          _set_composite_primary_key_values primary_key, values
+          _set_composite_primary_key_values primary_key, values, true
         else
-          _set_single_primary_key_value primary_key, values
+          _set_single_primary_key_value primary_key, values, true
         end
 
       metadata = TableMetadata.new self, arel_table
@@ -215,15 +231,25 @@ module ActiveRecord
       _convert_primary_key primary_key_value, returning
     end
 
-    def self._set_composite_primary_key_values primary_key, values
-      primary_key_value = []
-      primary_key.each do |col|
-        primary_key_value.append _set_composite_primary_key_value col, values
+    def self._has_all_primary_key_values? primary_key, values
+      if primary_key.is_a? Array
+        all = TrueClass
+        primary_key.each do |key|
+          all &&= values.key? key
+        end
+        all
+      else
+        values.key? primary_key
+      end
+    end
+
+    def self._set_composite_primary_key_values primary_key, values, is_mutation
+      primary_key.map do |col|
+        _set_composite_primary_key_value col, values, is_mutation
       end
-      primary_key_value
     end
 
-    def self._set_composite_primary_key_value primary_key, values
+    def self._set_composite_primary_key_value primary_key, values, is_mutation
       value = values[primary_key]
       type = ActiveModel::Type::BigInteger.new
 
@@ -232,6 +258,8 @@ module ActiveRecord
         value = value.value
       end
 
+      return value if is_auto_generated?(primary_key) \
+        && !(is_mutation && connection.use_client_side_id_for_mutations)
       return value unless prefetch_primary_key?
 
       if value.nil?
@@ -248,10 +276,11 @@ module ActiveRecord
       value
     end
 
-    def self._set_single_primary_key_value primary_key, values
+    def self._set_single_primary_key_value primary_key, values, is_mutation
       primary_key_value = values[primary_key] || values[primary_key.to_sym]
 
-      return primary_key_value if sequence_name
+      return primary_key_value if has_auto_generated_primary_key? \
+        && !(is_mutation && connection.use_client_side_id_for_mutations)
       return primary_key_value unless prefetch_primary_key?
 
       if primary_key_value.nil?
@@ -419,14 +448,12 @@ module ActiveRecord
     end
 
     def serialize_keys metadata, keys
-      serialized_values = []
-      keys.each do |key|
-        serialized_values << ActiveRecord::Type::Spanner::SpannerActiveRecordConverter
-                             .serialize_with_transaction_isolation_level(metadata.type(key),
-                                                                         attribute_in_database(key),
-                                                                         :mutation)
+      keys.map do |key|
+        ActiveRecord::Type::Spanner::SpannerActiveRecordConverter
+          .serialize_with_transaction_isolation_level(metadata.type(key),
+                                                      attribute_in_database(key),
+                                                      :mutation)
       end
-      serialized_values
     end
 
     def _execute_version_check attempted_action # rubocop:disable Metrics/AbcSize
@@ -456,7 +483,7 @@ module ActiveRecord
 
       # 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 #{pk_sql} AND `#{locking_column}` = @lock_version"
+            "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