activerecord-spanner-adapter 1.5.1 → 1.6.0

data/examples/snippets/interleaved-tables-before-7.1/models/track.rb

@@ -0,0 +1,28 @@
+# 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.
+
+class Track < ActiveRecord::Base
+  # Use the `composite_primary_key` gem to create a composite primary key definition for the model.
+  self.primary_keys = :singerid, :albumid, :trackid
+
+  # `tracks` is defined as INTERLEAVE IN PARENT `albums`. The primary key of `albums` is ()`singerid`, `albumid`).
+  belongs_to :album, foreign_key: [:singerid, :albumid]
+
+  # `tracks` also has a `singerid` column should be used to associate a Track with a Singer.
+  belongs_to :singer, foreign_key: :singerid
+
+  # Override the default initialize method to automatically set the singer attribute when an album is given.
+  def initialize attributes = nil
+    super
+    self.singer ||= album&.singer
+  end
+
+  def album=value
+    super
+    # Ensure the singer of this track is equal to the singer of the album that is set.
+    self.singer = value&.singer
+  end
+end

data/examples/snippets/query-logs/README.md

@@ -0,0 +1,43 @@
+# Sample - Query Logs
+
+__NOTE__: Query logs require additional configuration for Cloud Spanner. Please read the entire file.
+
+Rails 7.0 and higher supports [Query Logs](https://api.rubyonrails.org/classes/ActiveRecord/QueryLogs.html). Query Logs
+can be used to automatically annotate all queries that are executed based on the current execution context.
+
+The Cloud Spanner ActiveRecord provider can be used in combination with Query Logs. The query logs are automatically
+translated to request tags for the queries.
+See https://cloud.google.com/spanner/docs/introspection/troubleshooting-with-tags for more 
+information about request and transaction tags in Cloud Spanner.
+
+## Configuration
+Using Query Logs with Cloud Spanner requires some specific configuration:
+1. You must set `ActiveRecord::QueryLogs.prepend_comment = true`
+2. You must include `{ request_tag:  "true" }` as the first tag in your configuration.
+
+```ruby
+ActiveRecord::QueryLogs.prepend_comment = true
+config.active_record.query_log_tags = [
+  {
+    request_tag:  "true",
+  },
+  :namespaced_controller,
+  :action,
+  :job,
+  {
+    request_id: ->(context) { context[:controller]&.request&.request_id },
+    job_id: ->(context) { context[:job]&.job_id },
+    tenant_id: -> { Current.tenant&.id },
+    static: "value",
+  },
+]
+```
+
+The sample will automatically start a Spanner Emulator in a docker container and execute the sample
+against that emulator. The emulator will automatically be stopped when the application finishes.
+
+Run the application with the command
+
+```bash
+bundle exec rake run
+```

data/examples/snippets/query-logs/Rakefile

@@ -0,0 +1,13 @@
+# 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"
+require "sinatra/activerecord/rake"
+
+desc "Sample showing how to use automatic query log tagging on Cloud Spanner with ActiveRecord."
+task :run do
+  Dir.chdir("..") { sh "bundle exec rake run[query-logs]" }
+end

data/examples/snippets/query-logs/application.rb

@@ -0,0 +1,63 @@
+# 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 "io/console"
+require_relative "../config/environment"
+require_relative "models/singer"
+require_relative "models/album"
+
+class Application
+  def self.run
+    enable_query_logs
+
+    puts ""
+    puts "Query all Albums and include an automatically generated request tag"
+    albums = Album.all
+    puts "Queried #{albums.length} albums using an automatically generated request tag"
+
+    puts ""
+    puts "Press any key to end the application"
+    STDIN.getch
+  end
+
+  def self.enable_query_logs
+    # Enables Query Logs in a non-Rails application. Normally, this should be done
+    # as described here: https://api.rubyonrails.org/classes/ActiveRecord/QueryLogs.html
+    ActiveRecord.query_transformers << ActiveRecord::QueryLogs
+
+    # Query log comments *MUST* be prepended to be included as a request tag.
+    ActiveRecord::QueryLogs.prepend_comment = true
+
+    # This block manually enables Query Logs without a full Rails application.
+    # This should normally not be needed in your application.
+    ActiveRecord::QueryLogs.taggings.merge!(
+      application:  "example-app",
+      action:       "run-test-application",
+      pid:          -> { Process.pid.to_s },
+      socket:       ->(context) { context[:connection].pool.db_config.socket },
+      db_host:      ->(context) { context[:connection].pool.db_config.host },
+      database:     ->(context) { context[:connection].pool.db_config.database }
+    )
+
+    ActiveRecord::QueryLogs.tags = [
+      # The first tag *MUST* be the fixed value 'request_tag:true'.
+      {
+        request_tag:  "true"
+      },
+      :controller,
+      :action,
+      :job,
+      {
+        request_id: ->(context) { context[:controller]&.request&.request_id },
+        job_id: ->(context) { context[:job]&.job_id }
+      },
+      :db_host,
+      :database
+    ]
+  end
+end
+
+Application.run

data/examples/snippets/query-logs/config/database.yml

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

data/examples/snippets/query-logs/db/migrate/01_create_tables.rb

@@ -0,0 +1,21 @@
+# 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 CreateTables < ActiveRecord::Migration[6.0]
+  def change
+    connection.ddl_batch do
+      create_table :singers do |t|
+        t.string :first_name
+        t.string :last_name
+      end
+
+      create_table :albums do |t|
+        t.string :title
+        t.references :singer, index: false, foreign_key: true
+      end
+    end
+  end
+end

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