activerecord-spanner-adapter 0.3.0 → 0.5.0

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

@@ -0,0 +1,26 @@
+# Copyright 2020 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 Type
+    module Spanner
+      class Bytes < ActiveRecord::Type::Binary
+        def serialize value
+          return super value if value.nil?
+
+          if value.respond_to?(:read) && value.respond_to?(:rewind)
+            value.rewind
+            value = value.read
+          end
+
+          Base64.strict_encode64 value.force_encoding("ASCII-8BIT")
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,32 @@
+# Copyright 2021 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 Type
+    module Spanner
+      class SpannerActiveRecordConverter
+        ##
+        # Converts an ActiveModel::Type to a Spanner type code.
+        def self.convert_active_model_type_to_spanner type # rubocop:disable Metrics/CyclomaticComplexity
+          case type
+          when NilClass then nil
+          when ActiveModel::Type::Integer, ActiveModel::Type::BigInteger then :INT64
+          when ActiveModel::Type::Boolean then :BOOL
+          when ActiveModel::Type::String, ActiveModel::Type::ImmutableString then :STRING
+          when ActiveModel::Type::Binary, ActiveRecord::Type::Spanner::Bytes then :BYTES
+          when ActiveModel::Type::Float then :FLOAT64
+          when ActiveModel::Type::Decimal then :NUMERIC
+          when ActiveModel::Type::DateTime, ActiveModel::Type::Time, ActiveRecord::Type::Spanner::Time then :TIMESTAMP
+          when ActiveModel::Type::Date then :DATE
+          when ActiveRecord::Type::Spanner::Array then [convert_active_model_type_to_spanner(type.element_type)]
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,37 @@
+# Copyright 2020 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 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
+          val = super value
+          val.acts_like?(:time) ? val.utc.rfc3339(9) : val
+        end
+
+        def user_input_in_time_zone value
+          return value.in_time_zone if value.is_a? ::Time
+          super value
+        end
+
+        private
+
+        def cast_value value
+          if value.is_a? ::String
+            value = value.empty? ? nil : ::Time.parse(value)
+          end
+
+          value
+        end
+      end
+    end
+  end
+end

data/lib/activerecord-spanner-adapter.rb

@@ -0,0 +1,23 @@
+# Copyright 2020 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 "activerecord_spanner_adapter/version"
+
+if defined?(Rails)
+  module ActiveRecord
+    module ConnectionAdapters
+      class SpannerRailtie < ::Rails::Railtie
+        rake_tasks do
+          require "active_record/tasks/spanner_database_tasks"
+        end
+
+        ActiveSupport.on_load :active_record do
+          require "active_record/connection_adapters/spanner_adapter"
+        end
+      end
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/base.rb

@@ -0,0 +1,217 @@
+# Copyright 2021 Google LLC
+#
+# Use of this source code is governed by an MIT-style
+# license that can be found in the LICENSE file or at
+# https://opensource.org/licenses/MIT.
+
+module ActiveRecord
+  class TableMetadata # :nodoc:
+    # This attr_reader is private in ActiveRecord 6.0.x and public in 6.1.x. This makes sure it is always available in
+    # the Spanner adapter.
+    attr_reader :arel_table
+  end
+
+  class Base
+    # 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.
+    def self.create attributes = nil, &block
+      return super if active_transaction?
+
+      transaction isolation: :buffered_mutations do
+        return super
+      end
+    end
+
+    def self._insert_record values
+      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+
+      primary_key = self.primary_key
+      primary_key_value = nil
+
+      if primary_key && values.is_a?(Hash)
+        primary_key_value = values[primary_key]
+
+        if !primary_key_value && prefetch_primary_key?
+          primary_key_value = next_sequence_value
+          values[primary_key] = primary_key_value
+        end
+      end
+
+      metadata = TableMetadata.new self, arel_table
+      columns, grpc_values = _create_grpc_values_for_insert metadata, values
+
+      mutation = Google::Cloud::Spanner::V1::Mutation.new(
+        insert: Google::Cloud::Spanner::V1::Mutation::Write.new(
+          table: arel_table.name,
+          columns: columns,
+          values: [grpc_values.list_value]
+        )
+      )
+      Base.connection.current_spanner_transaction.buffer mutation
+
+      primary_key_value
+    end
+
+    # Deletes all records of this class. This method will use mutations instead of DML if there is no active
+    # transaction, or if the active transaction has been created with the option isolation: :buffered_mutations.
+    def self.delete_all
+      return super if active_transaction?
+
+      transaction isolation: :buffered_mutations do
+        return super
+      end
+    end
+
+    def self.active_transaction?
+      current_transaction = connection.current_transaction
+      !(current_transaction.nil? || current_transaction.is_a?(ConnectionAdapters::NullTransaction))
+    end
+
+    # Updates the given attributes of the object in the database. This method will use mutations instead
+    # of DML if there is no active transaction, or if the active transaction has been created with the option
+    # isolation: :buffered_mutations.
+    def update attributes
+      return super if Base.active_transaction?
+
+      transaction isolation: :buffered_mutations do
+        return super
+      end
+    end
+
+    # Deletes the object in the database. This method will use mutations instead
+    # of DML if there is no active transaction, or if the active transaction has been created with the option
+    # isolation: :buffered_mutations.
+    def destroy
+      return super if Base.active_transaction?
+
+      transaction isolation: :buffered_mutations do
+        return super
+      end
+    end
+
+    private
+
+    def self._create_grpc_values_for_insert metadata, values
+      serialized_values = []
+      columns = []
+      values.each_pair do |k, v|
+        type = metadata.type k
+        serialized_values << (type.method(:serialize).arity < 0 ? type.serialize(v, :mutation) : type.serialize(v))
+        columns << metadata.arel_table[k].name
+      end
+      [columns, Google::Protobuf::Value.new(list_value:
+        Google::Protobuf::ListValue.new(
+          values: serialized_values.map do |value|
+            Google::Cloud::Spanner::Convert.object_to_grpc_value value
+          end
+        ))]
+    end
+    private_class_method :_create_grpc_values_for_insert
+
+    def _update_row attribute_names, attempted_action = "update"
+      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+
+      if locking_enabled?
+        _execute_version_check attempted_action
+        attribute_names << self.class.locking_column
+        self[self.class.locking_column] += 1
+      end
+
+      metadata = TableMetadata.new self.class, self.class.arel_table
+      values = attributes_with_values attribute_names
+      columns, grpc_values = _create_grpc_values_for_update metadata, values
+
+      mutation = Google::Cloud::Spanner::V1::Mutation.new(
+        update: Google::Cloud::Spanner::V1::Mutation::Write.new(
+          table: self.class.arel_table.name,
+          columns: columns,
+          values: [grpc_values.list_value]
+        )
+      )
+      Base.connection.current_spanner_transaction.buffer mutation
+      1 # Affected rows
+    end
+
+    def _create_grpc_values_for_update metadata, values
+      constraints = {}
+      keys = self.class.primary_and_parent_key
+      keys.each do |key|
+        constraints[key] = attribute_in_database key
+      end
+
+      # Use both the where values + the values that are actually set.
+      all_values = [constraints, values]
+      all_serialized_values = []
+      all_columns = []
+      all_values.each do |h|
+        h.each_pair do |k, v|
+          type = metadata.type k
+          has_serialize_options = type.method(:serialize).arity < 0
+          all_serialized_values << (has_serialize_options ? type.serialize(v, :mutation) : type.serialize(v))
+          all_columns << metadata.arel_table[k].name
+        end
+      end
+      [all_columns, Google::Protobuf::Value.new(list_value:
+        Google::Protobuf::ListValue.new(
+          values: all_serialized_values.map do |value|
+            Google::Cloud::Spanner::Convert.object_to_grpc_value value
+          end
+        ))]
+    end
+
+    def destroy_row
+      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+
+      _delete_row
+    end
+
+    def _delete_row
+      return super unless Base.connection&.current_spanner_transaction&.isolation == :buffered_mutations
+      if locking_enabled?
+        _execute_version_check "destroy"
+      end
+
+      metadata = TableMetadata.new self.class, self.class.arel_table
+      keys = self.class.primary_and_parent_key
+      serialized_values = serialize_keys metadata, keys
+      list_value = Google::Protobuf::ListValue.new(
+        values: serialized_values.map do |value|
+          Google::Cloud::Spanner::Convert.object_to_grpc_value value
+        end
+      )
+      mutation = Google::Cloud::Spanner::V1::Mutation.new(
+        delete: Google::Cloud::Spanner::V1::Mutation::Delete.new(
+          table: self.class.arel_table.name,
+          key_set: { keys: [list_value] }
+        )
+      )
+      Base.connection.current_spanner_transaction.buffer mutation
+      1 # Affected rows
+    end
+
+    def serialize_keys metadata, keys
+      serialized_values = []
+      keys.each do |key|
+        type = metadata.type key
+        has_serialize_options = type.method(:serialize).arity < 0
+        serialized_values << type.serialize(attribute_in_database(key), :mutation) if has_serialize_options
+        serialized_values << type.serialize(attribute_in_database(key)) unless has_serialize_options
+      end
+      serialized_values
+    end
+
+    def _execute_version_check attempted_action
+      locking_column = self.class.locking_column
+      previous_lock_value = read_attribute_before_type_cast locking_column
+
+      # We need to check the version using a SELECT query, as a mutation cannot include a WHERE clause.
+      sql = "SELECT 1 FROM `#{self.class.arel_table.name}` " \
+              "WHERE `#{self.class.primary_key}` = @id AND `#{locking_column}` = @lock_version"
+      params = { "id" => id_in_database, "lock_version" => previous_lock_value }
+      param_types = { "id" => :INT64, "lock_version" => :INT64 }
+      locked_row = Base.connection.raw_connection.execute_query sql, params: params, types: param_types
+      raise ActiveRecord::StaleObjectError.new(self, attempted_action) unless locked_row.rows.any?
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/connection.rb

@@ -0,0 +1,324 @@
+# Copyright 2020 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 "google/cloud/spanner"
+require "spanner_client_ext"
+require "activerecord_spanner_adapter/information_schema"
+
+module ActiveRecordSpannerAdapter
+  class Connection
+    attr_reader :instance_id, :database_id, :spanner
+    attr_accessor :current_transaction
+
+    def initialize config
+      @instance_id = config[:instance]
+      @database_id = config[:database]
+      @spanner = self.class.spanners config
+    end
+
+    def self.spanners config
+      config = config.symbolize_keys
+      @spanners ||= {}
+      @mutex ||= Mutex.new
+      @mutex.synchronize do
+        @spanners[database_path(config)] ||= Google::Cloud::Spanner.new(
+          project_id: config[:project],
+          credentials: config[:credentials],
+          emulator_host: config[:emulator_host],
+          scope: config[:scope],
+          timeout: config[:timeout],
+          lib_name: "spanner-activerecord-adapter",
+          lib_version: ActiveRecordSpannerAdapter::VERSION
+        )
+      end
+    end
+
+    def self.information_schema config
+      @information_schemas ||= {}
+      @information_schemas[database_path(config)] ||= \
+        ActiveRecordSpannerAdapter::InformationSchema.new new(config)
+    end
+
+    def session
+      @last_used = Time.current
+      @session ||= spanner.create_session instance_id, database_id
+    end
+    alias connect! session
+
+    def active?
+      # This method should not initialize a session.
+      unless @session
+        return false
+      end
+      # Assume that it is still active if it has been used in the past 50 minutes.
+      if ((Time.current - @last_used) / 60).round < 50
+        return true
+      end
+      session.execute_query "SELECT 1"
+      true
+    rescue StandardError
+      false
+    end
+
+    def disconnect!
+      session.release!
+      true
+    ensure
+      @session = nil
+    end
+
+    def reset!
+      disconnect!
+      session
+      true
+    end
+
+    # Database Operations
+
+    def create_database
+      job = spanner.create_database instance_id, database_id
+      job.wait_until_done!
+      raise Google::Cloud::Error.from_error job.error if job.error?
+      job.database
+    end
+
+    def database
+      @database ||= begin
+        database = spanner.database instance_id, database_id
+        unless database
+          raise ActiveRecord::NoDatabaseError(
+            "#{spanner.project}/#{instance_id}/#{database_id}"
+          )
+        end
+        database
+      end
+    end
+
+    # DDL Statements
+
+    # @params [Array<String>, String] sql Single or list of statements
+    def execute_ddl statements, operation_id: nil, wait_until_done: true
+      raise "DDL cannot be executed during a transaction" if current_transaction&.active?
+      self.current_transaction = nil
+
+      statements = Array statements
+      return unless statements.any?
+
+      # If a DDL batch is active we only buffer the statements on the connection until the batch is run.
+      if @ddl_batch
+        @ddl_batch.push(*statements)
+        return true
+      end
+
+      execute_ddl_statements statements, operation_id, wait_until_done
+    end
+
+    # DDL Batching
+
+    ##
+    # Executes a set of DDL statements as one batch. This method raises an error if no block is given.
+    #
+    # @example
+    #   connection.ddl_batch do
+    #     connection.execute_ddl "CREATE TABLE `Users` (Id INT64, Name STRING(MAX)) PRIMARY KEY (Id)"
+    #     connection.execute_ddl "CREATE INDEX Idx_Users_Name ON `Users` (Name)"
+    #   end
+    def ddl_batch
+      raise Google::Cloud::FailedPreconditionError, "No block given for the DDL batch" unless block_given?
+      begin
+        start_batch_ddl
+        yield
+        run_batch
+      rescue StandardError
+        abort_batch
+        raise
+      ensure
+        @ddl_batch = nil
+      end
+    end
+
+    ##
+    # Returns true if this connection is currently executing a DDL batch, and otherwise false.
+    def ddl_batch?
+      return true if @ddl_batch
+      false
+    end
+
+    ##
+    # Starts a manual DDL batch. The batch must be ended by calling either run_batch or abort_batch.
+    #
+    # @example
+    #   begin
+    #     connection.start_batch_ddl
+    #     connection.execute_ddl "CREATE TABLE `Users` (Id INT64, Name STRING(MAX)) PRIMARY KEY (Id)"
+    #     connection.execute_ddl "CREATE INDEX Idx_Users_Name ON `Users` (Name)"
+    #     connection.run_batch
+    #   rescue StandardError
+    #     connection.abort_batch
+    #     raise
+    #   end
+    def start_batch_ddl
+      if @ddl_batch
+        raise Google::Cloud::FailedPreconditionError, "A DDL batch is already active on this connection"
+      end
+      @ddl_batch = []
+    end
+
+    ##
+    # Aborts the current batch on this connection. This is a no-op if there is no batch on this connection.
+    #
+    # @see start_batch_ddl
+    def abort_batch
+      @ddl_batch = nil
+    end
+
+    ##
+    # Runs the current batch on this connection. This will raise a FailedPreconditionError if there is no
+    # active batch on this connection.
+    #
+    # @see start_batch_ddl
+    def run_batch
+      unless @ddl_batch
+        raise Google::Cloud::FailedPreconditionError, "There is no batch active on this connection"
+      end
+      # Just return if the batch is empty.
+      return true if @ddl_batch.empty?
+      begin
+        execute_ddl_statements @ddl_batch, nil, true
+      ensure
+        @ddl_batch = nil
+      end
+    end
+
+    # DQL, DML Statements
+
+    def execute_query sql, params: nil, types: nil
+      if params
+        converted_params, types = \
+          Google::Cloud::Spanner::Convert.to_input_params_and_types(
+            params, types
+          )
+      end
+
+      # Clear the transaction from the previous statement.
+      unless current_transaction&.active?
+        self.current_transaction = nil
+      end
+
+      begin
+        session.execute_query \
+          sql,
+          params: converted_params,
+          types: types,
+          transaction: transaction_selector,
+          seqno: (current_transaction&.next_sequence_number)
+      rescue Google::Cloud::AbortedError
+        # Mark the current transaction as aborted to prevent any unnecessary further requests on the transaction.
+        current_transaction&.mark_aborted
+        raise
+      rescue Google::Cloud::NotFoundError => e
+        if session_not_found?(e) || transaction_not_found?(e)
+          reset!
+          # Force a retry of the entire transaction if this statement was executed as part of a transaction.
+          # Otherwise, just retry the statement itself.
+          raise_aborted_err if current_transaction&.active?
+          retry
+        end
+        raise
+      end
+    end
+
+    # Transactions
+
+    def begin_transaction isolation = nil
+      raise "Nested transactions are not allowed" if current_transaction&.active?
+      self.current_transaction = Transaction.new self, isolation
+      current_transaction.begin
+      current_transaction
+    end
+
+    def commit_transaction
+      raise "This connection does not have a transaction" unless current_transaction
+      current_transaction.commit
+    end
+
+    def rollback_transaction
+      raise "This connection does not have a transaction" unless current_transaction
+      current_transaction.rollback
+    end
+
+    def transaction_selector
+      return current_transaction&.transaction_selector if current_transaction&.active?
+    end
+
+    def truncate table_name
+      session.delete table_name
+    end
+
+    def self.database_path config
+      "#{config[:emulator_host]}/#{config[:project]}/#{config[:instance]}/#{config[:database]}"
+    end
+
+    def session_not_found? err
+      if err.respond_to?(:metadata) && err.metadata["google.rpc.resourceinfo-bin"]
+        resource_info = Google::Rpc::ResourceInfo.decode err.metadata["google.rpc.resourceinfo-bin"]
+        type = resource_info["resource_type"]
+        return "type.googleapis.com/google.spanner.v1.Session".eql? type
+      end
+      false
+    end
+
+    def transaction_not_found? err
+      if err.respond_to?(:metadata) && err.metadata["google.rpc.resourceinfo-bin"]
+        resource_info = Google::Rpc::ResourceInfo.decode err.metadata["google.rpc.resourceinfo-bin"]
+        type = resource_info["resource_type"]
+        return "type.googleapis.com/google.spanner.v1.Transaction".eql? type
+      end
+      false
+    end
+
+    def raise_aborted_err
+      retry_info = Google::Rpc::RetryInfo.new retry_delay: Google::Protobuf::Duration.new(seconds: 0, nanos: 1)
+      begin
+        raise GRPC::BadStatus.new(
+          GRPC::Core::StatusCodes::ABORTED,
+          "Transaction aborted",
+          "google.rpc.retryinfo-bin": Google::Rpc::RetryInfo.encode(retry_info)
+        )
+      rescue GRPC::BadStatus
+        raise Google::Cloud::AbortedError
+      end
+    end
+
+    private
+
+    def execute_ddl_statements statements, operation_id, wait_until_done
+      job = database.update statements: statements, operation_id: operation_id
+      job.wait_until_done! if wait_until_done
+      raise Google::Cloud::Error.from_error job.error if job.error?
+      job.done?
+    end
+
+    ##
+    # Retrieves the delay value from Google::Cloud::AbortedError or
+    # GRPC::Aborted
+    def delay_from_aborted err
+      return nil if err.nil?
+      if err.respond_to?(:metadata) && err.metadata["google.rpc.retryinfo-bin"]
+        retry_info = Google::Rpc::RetryInfo.decode err.metadata["google.rpc.retryinfo-bin"]
+        seconds = retry_info["retry_delay"].seconds
+        nanos = retry_info["retry_delay"].nanos
+        return seconds if nanos.zero?
+        return seconds + (nanos / 1_000_000_000.0)
+      end
+      # No metadata? Try the inner error
+      delay_from_aborted err.cause
+    rescue StandardError
+      # Any error indicates the backoff should be handled elsewhere
+      nil
+    end
+  end
+end