activerecord-spanner-adapter 1.2.2 → 1.4.0

data/examples/snippets/optimistic-locking/db/schema.rb

@@ -2,8 +2,8 @@
 # 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 `rails
-# db:schema:load`. When creating a new database, `rails db:schema:load` tends to
+# 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.
@@ -11,19 +11,24 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
-  create_table "albums", force: :cascade do |t|
+  create_table "albums", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
     t.decimal "marketing_budget"
     t.integer "singer_id", limit: 8
     t.integer "lock_version", limit: 8
   end
 
-  create_table "singers", force: :cascade do |t|
+  create_table "singers", id: { limit: 8 }, force: :cascade do |t|
     t.string "first_name"
     t.string "last_name"
     t.integer "lock_version", limit: 8
   end
 
   add_foreign_key "albums", "singers"
+  connection.run_batch
+rescue
+  abort_batch
+  raise
 end

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

@@ -11,6 +11,7 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
   create_table "albums", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
@@ -23,4 +24,8 @@ ActiveRecord::Schema.define(version: 1) do
   end
 
   add_foreign_key "albums", "singers"
+  connection.run_batch
+rescue
+  abort_batch
+  raise
 end

data/examples/snippets/quickstart/db/schema.rb

@@ -2,8 +2,8 @@
 # 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 `rails
-# db:schema:load`. When creating a new database, `rails db:schema:load` tends to
+# 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.
@@ -11,16 +11,21 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
-  create_table "albums", force: :cascade do |t|
+  create_table "albums", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
     t.integer "singer_id", limit: 8
   end
 
-  create_table "singers", force: :cascade do |t|
+  create_table "singers", id: { limit: 8 }, 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/read-only-transactions/db/schema.rb

@@ -11,6 +11,7 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
   create_table "albums", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
@@ -23,4 +24,8 @@ ActiveRecord::Schema.define(version: 1) do
   end
 
   add_foreign_key "albums", "singers"
+  connection.run_batch
+rescue
+  abort_batch
+  raise
 end

data/examples/snippets/read-write-transactions/db/schema.rb

@@ -2,8 +2,8 @@
 # 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 `rails
-# db:schema:load`. When creating a new database, `rails db:schema:load` tends to
+# 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.
@@ -11,17 +11,22 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
-  create_table "albums", force: :cascade do |t|
+  create_table "albums", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
     t.decimal "marketing_budget"
     t.integer "singer_id", limit: 8
   end
 
-  create_table "singers", force: :cascade do |t|
+  create_table "singers", id: { limit: 8 }, 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/stale-reads/db/schema.rb

@@ -11,6 +11,7 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
   create_table "albums", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
@@ -23,4 +24,8 @@ ActiveRecord::Schema.define(version: 1) do
   end
 
   add_foreign_key "albums", "singers"
+  connection.run_batch
+rescue
+  abort_batch
+  raise
 end

data/examples/snippets/timestamp-data-type/db/schema.rb

@@ -2,8 +2,8 @@
 # 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 `rails
-# db:schema:load`. When creating a new database, `rails db:schema:load` tends to
+# 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.
@@ -11,11 +11,16 @@
 # It's strongly recommended that you check this file into your version control system.
 
 ActiveRecord::Schema.define(version: 1) do
+  connection.start_batch_ddl
 
-  create_table "meetings", force: :cascade do |t|
+  create_table "meetings", id: { limit: 8 }, force: :cascade do |t|
     t.string "title"
     t.time "meeting_time"
     t.string "meeting_timezone"
   end
 
+  connection.run_batch
+rescue
+  abort_batch
+  raise
 end

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

@@ -0,0 +1,23 @@
+# 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.
+#
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module Spanner
+      class Column < ConnectionAdapters::Column
+        def has_default? # rubocop:disable Naming/PredicateName
+          super && !virtual?
+        end
+
+        def virtual?
+          sql_type_metadata.generated
+        end
+      end
+    end
+  end
+end

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

@@ -35,7 +35,10 @@ module ActiveRecord
               binds.delete staleness_hint
             end
 
-            log sql, name do
+            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
@@ -219,9 +222,9 @@ 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
-            value = bind
-            value = type.serialize bind.value, :dml if type.respond_to?(:serialize) && type.method(:serialize).arity < 0
-            value = type.serialize bind.value if type.respond_to?(:serialize) && type.method(:serialize).arity >= 0
+            bind_value = bind.respond_to?(:value) ? bind.value : bind
+            value = ActiveRecord::Type::Spanner::SpannerActiveRecordConverter
+                    .serialize_with_transaction_isolation_level(type, bind_value, :dml)
 
             ["p#{i + 1}", value]
           end.to_h

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

@@ -54,6 +54,15 @@ module ActiveRecord
         def quoted_binary value
           "b'#{value}'"
         end
+
+        def _type_cast value
+          case value
+          when Array
+            ActiveSupport::JSON.encode value
+          else
+            super
+          end
+        end
       end
     end
   end

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

@@ -10,7 +10,8 @@ module ActiveRecord
       class SchemaCreation < SchemaCreation
         private
 
-        # rubocop:disable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity
+        # rubocop:disable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
+        # rubocop:disable Metrics/MethodLength
 
         def visit_TableDefinition o
           create_sql = +"CREATE TABLE #{quote_table_name o.name} "
@@ -26,6 +27,14 @@ module ActiveRecord
             end
           end
 
+          if ActiveRecord::VERSION::MAJOR >= 7
+            statements.concat(o.check_constraints.map { |chk| accept chk })
+          elsif ActiveRecord::VERSION::MAJOR == 6 && ActiveRecord::VERSION::MINOR >= 1
+            statements.concat(
+              o.check_constraints.map { |expression, options| check_constraint_in_create o.name, expression, options }
+            )
+          end
+
           create_sql << "(#{statements.join ', '}) " if statements.any?
 
           primary_keys = if o.primary_keys
@@ -63,7 +72,7 @@ module ActiveRecord
         def visit_ColumnDefinition o
           o.sql_type = type_to_sql o.type, **o.options
           column_sql = +"#{quote_column_name o.name} #{o.sql_type}"
-          add_column_options! column_sql, column_options(o)
+          add_column_options! o, column_sql, column_options(o)
           column_sql
         end
 
@@ -112,12 +121,16 @@ module ActiveRecord
           sql
         end
 
-        # rubocop:enable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity
+        # rubocop:enable Naming/MethodName, Metrics/AbcSize, Metrics/PerceivedComplexity, Metrics/CyclomaticComplexity
+        # rubocop:enable Metrics/MethodLength
 
-        def add_column_options! sql, options
+        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}"
+          end
 
           if !options[:allow_commit_timestamp].nil? &&
              options[:column].sql_type == "TIMESTAMP"

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

@@ -15,8 +15,8 @@ module ActiveRecord
           @on_delete = on_delete
         end
 
-        def parent_key name
-          column name, :parent_key, null: false
+        def parent_key name, type: nil
+          column name, :parent_key, null: false, passed_type: type
         end
 
         def interleave_in?
@@ -33,6 +33,15 @@ module ActiveRecord
           end
         end
         alias belongs_to references
+
+        def new_column_definition name, type, **options
+          case type
+          when :virtual
+            type = options[:type]
+          end
+
+          super
+        end
       end
 
       class Table < ActiveRecord::ConnectionAdapters::Table

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

@@ -13,6 +13,62 @@ module ActiveRecord
         def default_primary_key? column
           schema_type(column) == :integer
         end
+
+        def prepare_column_options column
+          spec = super
+
+          unless column.sql_type_metadata.allow_commit_timestamp.nil?
+            spec[:allow_commit_timestamp] = column.sql_type_metadata.allow_commit_timestamp
+          end
+
+          if column.virtual?
+            spec[:as] = extract_expression_for_virtual_column column
+            spec[:stored] = true
+            spec = { type: schema_type(column).inspect }.merge! spec
+          end
+
+          spec
+        end
+
+        def header stream
+          str = StringIO.new
+          super str
+          stream.puts <<~HEADER
+            #{str.string.rstrip}
+              connection.start_batch_ddl
+
+          HEADER
+        end
+
+        def trailer stream
+          stream.puts <<~TRAILER
+              connection.run_batch
+            rescue
+              abort_batch
+              raise
+          TRAILER
+          super
+        end
+
+        def index_parts index
+          index_parts = super
+          index_parts << "null_filtered: #{index.null_filtered.inspect}" if index.null_filtered
+          index_parts << "interleave_in: #{index.interleave_in.inspect}" if index.interleave_in
+          index_parts << "storing: #{format_index_parts index.storing}" if index.storing.present?
+          index_parts
+        end
+
+        private
+
+        def column_spec_for_primary_key column
+          spec = super
+          spec.except! :limit if default_primary_key? column
+          spec
+        end
+
+        def extract_expression_for_virtual_column column
+          column.default_function.inspect
+        end
       end
     end
   end

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

@@ -8,6 +8,7 @@
 
 require "active_record/connection_adapters/spanner/schema_creation"
 require "active_record/connection_adapters/spanner/schema_dumper"
+require "active_record/connection_adapters/spanner/column"
 
 module ActiveRecord
   module ConnectionAdapters
@@ -39,18 +40,19 @@ module ActiveRecord
         end
         alias data_source_exists? table_exists?
 
-        def create_table table_name, **options
+        def create_table table_name, id: :primary_key, **options
           td = create_table_definition table_name, options
 
-          if options[:id] != false
+          if id
             pk = options.fetch :primary_key do
               Base.get_primary_key table_name.to_s.singularize
             end
+            id = id.fetch :type, :primary_key if id.is_a? Hash
 
             if pk.is_a? Array
               td.primary_keys pk
             else
-              td.primary_key pk, options.fetch(:id, :primary_key), **{}
+              td.primary_key pk, id, **{}
             end
           end
 
@@ -108,16 +110,23 @@ module ActiveRecord
         end
 
         def new_column_from_field _table_name, field
-          ConnectionAdapters::Column.new \
+          Spanner::Column.new \
             field.name,
             field.default,
-            fetch_type_metadata(field.spanner_type, field.ordinal_position),
-            field.nullable
+            fetch_type_metadata(field.spanner_type,
+                                field.ordinal_position,
+                                field.allow_commit_timestamp,
+                                field.generated),
+            field.nullable,
+            field.default_function
         end
 
-        def fetch_type_metadata sql_type, ordinal_position = nil
+        def fetch_type_metadata sql_type, ordinal_position = nil, allow_commit_timestamp = nil, generated = nil
           Spanner::TypeMetadata.new \
-            super(sql_type), ordinal_position: ordinal_position
+            super(sql_type),
+            ordinal_position: ordinal_position,
+            allow_commit_timestamp: allow_commit_timestamp,
+            generated: generated
         end
 
         def add_column table_name, column_name, type, **options
@@ -376,6 +385,44 @@ module ActiveRecord
         end
         alias add_belongs_to add_reference
 
+        # Check Contraints
+
+        def check_constraints table_name
+          information_schema { |i| i.check_constraints table_name }
+        end
+
+        def assume_migrated_upto_version version
+          version = version.to_i
+          sm_table = quote_table_name schema_migration.table_name
+
+          migrated = migration_context.get_all_versions
+          versions = migration_context.migrations.map(&:version)
+
+          execute "INSERT INTO #{sm_table} (version) VALUES (#{quote version.to_s})" unless migrated.include? version
+
+          inserting = (versions - migrated).select { |v| v < version }
+          return unless inserting.any?
+
+          if (duplicate = inserting.detect { |v| inserting.count(v) > 1 })
+            raise "Duplicate migration #{duplicate}. Please renumber your migrations to resolve the conflict."
+          end
+
+          execute insert_versions_sql(inserting)
+        end
+
+        def insert_versions_sql versions
+          sm_table = quote_table_name schema_migration.table_name
+
+          if versions.is_a? Array
+            sql = +"INSERT INTO #{sm_table} (version) VALUES\n"
+            sql << versions.reverse.map { |v| "(#{quote v.to_s})" }.join(",\n")
+            sql << ";"
+            sql
+          else
+            "INSERT INTO #{sm_table} (version) VALUES (#{quote versions.to_s});"
+          end
+        end
+
         def quoted_scope name = nil, type: nil
           scope = { schema: quote("") }
           scope[:name] = quote name if name
@@ -389,7 +436,7 @@ module ActiveRecord
 
         # rubocop:disable Lint/UnusedMethodArgument
         def type_to_sql type, limit: nil, precision: nil, scale: nil, **opts
-          type = type.to_sym if type
+          type = opts[:passed_type] || type&.to_sym
           native = native_database_types[type]
 
           return type.to_s unless native

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

@@ -12,24 +12,39 @@ module ActiveRecord
       class TypeMetadata < DelegateClass(SqlTypeMetadata)
         undef to_yaml if method_defined? :to_yaml
 
-        attr_reader :ordinal_position
+        include Deduplicable if defined?(Deduplicable)
 
-        def initialize type_metadata, ordinal_position: nil
+        attr_reader :ordinal_position, :allow_commit_timestamp, :generated
+
+        def initialize type_metadata, ordinal_position: nil, allow_commit_timestamp: nil, generated: nil
           super type_metadata
           @ordinal_position = ordinal_position
+          @allow_commit_timestamp = allow_commit_timestamp
+          @generated = generated
         end
 
         def == other
           other.is_a?(TypeMetadata) &&
             __getobj__ == other.__getobj__ &&
-            ordinal_position == other.ordinal_position
+            ordinal_position == other.ordinal_position &&
+            allow_commit_timestamp == other.allow_commit_timestamp &&
+            generated == other.generated
         end
         alias eql? ==
 
         def hash
           TypeMetadata.hash ^
             __getobj__.hash ^
-            ordinal_position.hash
+            ordinal_position.hash ^
+            allow_commit_timestamp.hash ^
+            generated.hash
+        end
+
+        private
+
+        def deduplicated
+          __setobj__ __getobj__.deduplicate
+          super
         end
       end
     end

data/lib/active_record/connection_adapters/spanner_adapter.rb

@@ -74,6 +74,9 @@ module ActiveRecord
       include Spanner::DatabaseStatements
       include Spanner::SchemaStatements
 
+      # Determines whether or not to log query binds when executing statements
+      class_attribute :log_statement_binds, instance_writer: false, default: false
+
       def initialize connection, logger, connection_options, config
         super connection, logger, config
         @connection_options = connection_options
@@ -170,6 +173,14 @@ module ActiveRecord
         true
       end
 
+      def supports_check_constraints?
+        true
+      end
+
+      def supports_virtual_columns?
+        true
+      end
+
       # Generate next sequence number for primary key
       def next_sequence_value _sequence_name
         SecureRandom.uuid.gsub("-", "").hex & 0x7FFFFFFFFFFFFFFF

data/lib/active_record/tasks/spanner_database_tasks.rb

@@ -29,7 +29,12 @@ module ActiveRecord
       end
 
       def purge
-        drop
+        begin
+          drop
+        rescue ActiveRecord::NoDatabaseError # rubocop:disable Lint/HandleExceptions
+          # ignored; create the database
+        end
+
         create
       end
 
@@ -54,15 +59,24 @@ module ActiveRecord
           next if ignore_tables.any? &&
                   (table_regx =~ statement || index_regx =~ statement)
           file.write statement
-          file.write "\n"
+          file.write ";\n"
         end
       ensure
         file.close
       end
 
       def structure_load filename, _extra_flags
-        statements = File.read(filename).split(/(?=^CREATE)/)
-        @connection.execute_ddl statements
+        statements = File.read(filename).split(/;/).map(&:strip).reject(&:empty?)
+        ddls = statements.select { |s| s =~ /^(CREATE|ALTER|DROP|GRANT|REVOKE|ANALYZE)/ }
+        @connection.execute_ddl ddls
+
+        client = @connection.spanner.client @connection.instance_id,
+                                            @connection.database_id
+        dmls = statements.reject { |s| s =~ /^(CREATE|ALTER|DROP|GRANT|REVOKE|ANALYZE)/ }
+
+        client.transaction do |tx|
+          dmls.each { |dml| tx.execute_query dml }
+        end
       end
     end
 

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

@@ -10,6 +10,16 @@ module ActiveRecord
   module Type
     module Spanner
       class SpannerActiveRecordConverter
+        def self.serialize_with_transaction_isolation_level type, value, isolation_level
+          if type.respond_to? :serialize_with_isolation_level
+            type.serialize_with_isolation_level value, isolation_level
+          elsif type.respond_to? :serialize
+            type.serialize value
+          else
+            value
+          end
+        end
+
         ##
         # Converts an ActiveModel::Type to a Spanner type code.
         def self.convert_active_model_type_to_spanner type # rubocop:disable Metrics/CyclomaticComplexity

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

@@ -10,9 +10,16 @@ module ActiveRecord
   module Type
     module Spanner
       class Time < ActiveRecord::Type::Time
-        def serialize value, *options
-          return "PENDING_COMMIT_TIMESTAMP()" if value == :commit_timestamp && options.length && options[0] == :dml
-          return "spanner.commit_timestamp()" if value == :commit_timestamp && options.length && options[0] == :mutation
+        def serialize_with_isolation_level value, isolation_level
+          if value == :commit_timestamp
+            return "PENDING_COMMIT_TIMESTAMP()" if isolation_level == :dml
+            return "spanner.commit_timestamp()" if isolation_level == :mutation
+          end
+
+          serialize value
+        end
+
+        def serialize value
           val = super value
           val.acts_like?(:time) ? val.utc.rfc3339(9) : val
         end