activerecord-spanner-adapter 0.3.0 → 1.0.0

data/lib/activerecord_spanner_adapter/index.rb

@@ -0,0 +1,80 @@
+# 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/index/column"
+
+module ActiveRecordSpannerAdapter
+  class Index
+    attr_accessor :table, :name, :columns, :type, :unique, :null_filtered,
+                  :interleave_in, :storing, :state
+
+    def initialize \
+        table,
+        name,
+        columns,
+        type: nil,
+        unique: false,
+        null_filtered: false,
+        interleave_in: nil,
+        storing: nil,
+        state: nil
+      @table = table.to_s
+      @name = name.to_s
+      @columns = Array(columns)
+      @type = type
+      @unique = unique
+      @null_filtered = null_filtered
+      @interleave_in = interleave_in unless interleave_in.to_s.empty?
+      @storing = storing || []
+      @state = state
+    end
+
+    def primary?
+      @type == "PRIMARY_KEY"
+    end
+
+    def columns_by_position
+      @columns.select(&:ordinal_position).sort do |c1, c2|
+        c1.ordinal_position <=> c2.ordinal_position
+      end
+    end
+
+    def column_names
+      columns_by_position.map(&:name)
+    end
+
+    def orders
+      columns_by_position.each_with_object({}) do |c, r|
+        r[c.name] = c.desc? ? :desc : :asc
+      end
+    end
+
+    def options
+      {
+        name: name,
+        order: orders,
+        unique: unique,
+        interleave_in: interleave_in,
+        null_filtered: null_filtered,
+        storing: storing
+      }.delete_if { |_, v| v.nil? }
+    end
+
+    def rename_column_options old_column, new_column
+      opts = options
+
+      opts[:order].transform_keys do |key|
+        key.to_s == new_column.to_s
+      end
+
+      columns = column_names.map do |c|
+        c.to_s == old_column.to_s ? new_column : c
+      end
+
+      { options: opts, columns: columns }
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/information_schema.rb

@@ -0,0 +1,262 @@
+# 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 "active_record/connection_adapters/abstract/quoting"
+require "activerecord_spanner_adapter/information_schema"
+require "activerecord_spanner_adapter/table"
+require "activerecord_spanner_adapter/index"
+require "activerecord_spanner_adapter/foreign_key"
+
+
+module ActiveRecordSpannerAdapter
+  class InformationSchema
+    include ActiveRecord::ConnectionAdapters::Quoting
+
+    attr_reader :connection
+
+    def initialize connection
+      @connection = connection
+      @mutex = Mutex.new
+    end
+
+    def tables table_name: nil, schema_name: nil, view: nil
+      sql = +"SELECT TABLE_CATALOG, TABLE_SCHEMA, TABLE_NAME, PARENT_TABLE_NAME, ON_DELETE_ACTION"
+      sql << " FROM INFORMATION_SCHEMA.TABLES"
+      sql << " WHERE TABLE_SCHEMA=%<schema_name>s"
+      sql << " AND TABLE_NAME=%<table_name>s" if table_name
+
+      rows = execute_query(
+        sql,
+        schema_name: (schema_name || ""), table_name: table_name
+      )
+
+      rows.map do |row|
+        table = Table.new(
+          row["TABLE_NAME"],
+          parent_table: row["PARENT_TABLE_NAME"],
+          on_delete: row["ON_DELETE_ACTION"],
+          schema_name: row["TABLE_SCHEMA"],
+          catalog: row["TABLE_CATALOG"]
+        )
+
+        if [:full, :columns].include? view
+          table.columns = table_columns table.name
+        end
+
+        if [:full, :indexes].include? view
+          table.indexes = indexes table.name
+        end
+        table
+      end
+    end
+
+    def table table_name, schema_name: nil, view: nil
+      tables(
+        table_name: table_name,
+        schema_name: schema_name,
+        view: view
+      ).first
+    end
+
+    def table_columns table_name, column_name: nil
+      sql = +"SELECT COLUMN_NAME, SPANNER_TYPE, IS_NULLABLE,"
+      sql << " CAST(COLUMN_DEFAULT AS STRING) AS COLUMN_DEFAULT, ORDINAL_POSITION"
+      sql << " FROM INFORMATION_SCHEMA.COLUMNS"
+      sql << " WHERE TABLE_NAME=%<table_name>s"
+      sql << " AND COLUMN_NAME=%<column_name>s" if column_name
+      sql << " ORDER BY ORDINAL_POSITION ASC"
+
+      execute_query(
+        sql,
+        table_name: table_name,
+        column_name: column_name
+      ).map do |row|
+        type, limit = parse_type_and_limit row["SPANNER_TYPE"]
+        Table::Column.new \
+          table_name,
+          row["COLUMN_NAME"],
+          type,
+          limit: limit,
+          ordinal_position: row["ORDINAL_POSITION"],
+          nullable: row["IS_NULLABLE"] == "YES",
+          default: row["COLUMN_DEFAULT"]
+      end
+    end
+
+    def table_column table_name, column_name
+      table_columns(table_name, column_name: column_name).first
+    end
+
+    # Returns the primary key columns of the given table. By default it will only return the columns that are not part
+    # of the primary key of the parent table (if any). These are the columns that are considered the primary key by
+    # ActiveRecord. The parent primary key columns are filtered out by default to allow interleaved tables to be
+    # considered as tables with a single-column primary key by ActiveRecord. The actual primary key of the table will
+    # include both the parent primary key columns and the 'own' primary key columns of a table.
+    def table_primary_keys table_name, include_parent_keys = false
+      sql = +"WITH TABLE_PK_COLS AS ( "
+      sql << "SELECT C.TABLE_NAME, C.COLUMN_NAME, C.INDEX_NAME, C.COLUMN_ORDERING, C.ORDINAL_POSITION "
+      sql << "FROM INFORMATION_SCHEMA.INDEX_COLUMNS C "
+      sql << "WHERE C.INDEX_TYPE = 'PRIMARY_KEY' "
+      sql << "AND TABLE_CATALOG = '' "
+      sql << "AND TABLE_SCHEMA = '') "
+      sql << "SELECT INDEX_NAME, COLUMN_NAME, COLUMN_ORDERING, ORDINAL_POSITION "
+      sql << "FROM TABLE_PK_COLS "
+      sql << "INNER JOIN INFORMATION_SCHEMA.TABLES T USING (TABLE_NAME) "
+      sql << "WHERE TABLE_NAME = %<table_name>s "
+      sql << "AND TABLE_CATALOG = '' "
+      sql << "AND TABLE_SCHEMA = '' "
+      unless include_parent_keys
+        sql << "AND (T.PARENT_TABLE_NAME IS NULL OR COLUMN_NAME NOT IN ( "
+        sql << "  SELECT COLUMN_NAME "
+        sql << "  FROM TABLE_PK_COLS "
+        sql << "  WHERE TABLE_NAME = T.PARENT_TABLE_NAME "
+        sql << ")) "
+      end
+      sql << "ORDER BY ORDINAL_POSITION"
+      execute_query(
+        sql,
+        table_name: table_name
+      ).map do |row|
+        Index::Column.new \
+          table_name,
+          row["INDEX_NAME"],
+          row["COLUMN_NAME"],
+          order: row["COLUMN_ORDERING"],
+          ordinal_position: row["ORDINAL_POSITION"]
+      end
+    end
+
+    def indexes table_name, index_name: nil, index_type: nil
+      table_indexes_columns = index_columns(
+        table_name,
+        index_name: index_name
+      )
+
+      sql = +"SELECT INDEX_NAME, INDEX_TYPE, IS_UNIQUE, IS_NULL_FILTERED, PARENT_TABLE_NAME, INDEX_STATE"
+      sql << " FROM INFORMATION_SCHEMA.INDEXES"
+      sql << " WHERE TABLE_NAME=%<table_name>s"
+      sql << " AND TABLE_CATALOG = ''"
+      sql << " AND TABLE_SCHEMA = ''"
+      sql << " AND INDEX_NAME=%<index_name>s" if index_name
+      sql << " AND INDEX_TYPE=%<index_type>s" if index_type
+      sql << " AND SPANNER_IS_MANAGED=FALSE"
+
+      execute_query(
+        sql,
+        table_name: table_name,
+        index_name: index_name,
+        index_type: index_type
+      ).map do |row|
+        columns = []
+        storing = []
+        table_indexes_columns.each do |c|
+          next unless c.index_name == row["INDEX_NAME"]
+          if c.ordinal_position
+            columns << c
+          else
+            storing << c.name
+          end
+        end
+
+        Index.new \
+          table_name,
+          row["INDEX_NAME"],
+          columns,
+          type: row["INDEX_TYPE"],
+          unique: row["IS_UNIQUE"],
+          null_filtered: row["IS_NULL_FILTERED"],
+          interleave_in: row["PARENT_TABLE_NAME"],
+          storing: storing,
+          state: row["INDEX_STATE"]
+      end
+    end
+
+    def index table_name, index_name
+      indexes(table_name, index_name: index_name).first
+    end
+
+    def index_columns table_name, index_name: nil
+      sql = +"SELECT INDEX_NAME, COLUMN_NAME, COLUMN_ORDERING, ORDINAL_POSITION"
+      sql << " FROM INFORMATION_SCHEMA.INDEX_COLUMNS"
+      sql << " WHERE TABLE_NAME=%<table_name>s"
+      sql << " AND TABLE_CATALOG = ''"
+      sql << " AND TABLE_SCHEMA = ''"
+      sql << " AND INDEX_NAME=%<index_name>s" if index_name
+      sql << " ORDER BY ORDINAL_POSITION ASC"
+
+      execute_query(
+        sql,
+        table_name: table_name, index_name: index_name
+      ).map do |row|
+        Index::Column.new \
+          table_name,
+          row["INDEX_NAME"],
+          row["COLUMN_NAME"],
+          order: row["COLUMN_ORDERING"],
+          ordinal_position: row["ORDINAL_POSITION"]
+      end
+    end
+
+    def indexes_by_columns table_name, column_names
+      column_names = Array(column_names).map(&:to_s)
+
+      indexes(table_name).select do |index|
+        index.columns.any? { |c| column_names.include? c.name }
+      end
+    end
+
+    def foreign_keys table_name
+      sql = <<~SQL
+        SELECT cc.table_name AS to_table,
+               cc.column_name AS primary_key,
+               fk.column_name as column,
+               fk.constraint_name AS name,
+               rc.update_rule AS on_update,
+               rc.delete_rule AS on_delete
+        FROM information_schema.referential_constraints rc
+        INNER JOIN information_schema.key_column_usage fk ON rc.constraint_name = fk.constraint_name
+        INNER JOIN information_schema.constraint_column_usage cc ON rc.constraint_name = cc.constraint_name
+        WHERE fk.table_name = %<table_name>s
+          AND fk.constraint_schema = %<constraint_schema>s
+      SQL
+
+      rows = execute_query(
+        sql, table_name: table_name, constraint_schema: ""
+      )
+
+      rows.map do |row|
+        ForeignKey.new(
+          table_name,
+          row["name"],
+          row["column"],
+          row["to_table"],
+          row["primary_key"],
+          on_delete: row["on_delete"],
+          on_update: row["on_update"]
+        )
+      end
+    end
+
+    def parse_type_and_limit value
+      matched = /^([A-Z]*)\((.*)\)/.match value
+      return [value] unless matched
+
+      limit = matched[2]
+      limit = limit.to_i unless limit == "MAX"
+
+      [matched[1], limit]
+    end
+
+    private
+
+    def execute_query sql, params = {}
+      params = params.transform_values { |v| quote v }
+      sql = format sql, params
+
+      @connection.execute_query(sql).rows
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/primary_key.rb

@@ -0,0 +1,31 @@
+# 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.
+
+module ActiveRecord
+  module AttributeMethods
+    module PrimaryKey
+      module ClassMethods
+        def primary_and_parent_key
+          reset_primary_and_parent_key unless defined? @primary_and_parent_key
+          @primary_and_parent_key
+        end
+
+        def reset_primary_and_parent_key
+          self.primary_and_parent_key = base_class? ? fetch_primary_and_parent_key : base_class.primary_and_parent_key
+        end
+
+        def fetch_primary_and_parent_key
+          return connection.schema_cache.primary_and_parent_keys table_name \
+            if ActiveRecord::Base != self && table_exists?
+        end
+
+        def primary_and_parent_key= value
+          @primary_and_parent_key = value
+        end
+      end
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/table/column.rb

@@ -0,0 +1,59 @@
+# 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.
+
+module ActiveRecordSpannerAdapter
+  class Table
+    class Column
+      attr_accessor :table_name, :name, :type, :limit, :ordinal_position,
+                    :allow_commit_timestamp, :default, :primary_key
+      attr_writer :nullable
+
+      def initialize \
+          table_name,
+          name,
+          type,
+          limit: nil,
+          ordinal_position: nil,
+          nullable: true,
+          allow_commit_timestamp: nil,
+          default: nil
+        @table_name = table_name.to_s
+        @name = name.to_s
+        @type = type
+        @limit = limit
+        @nullable = nullable != false
+        @ordinal_position = ordinal_position
+        @allow_commit_timestamp = allow_commit_timestamp
+        @default = default
+        @primary_key = false
+      end
+
+      def nullable
+        return false if primary_key
+        @nullable
+      end
+
+      def spanner_type
+        return "#{type}(#{limit || 'MAX'})" if limit_allowed?
+        type
+      end
+
+      def options
+        {
+          limit: limit,
+          null: nullable,
+          allow_commit_timestamp: allow_commit_timestamp
+        }.delete_if { |_, v| v.nil? }
+      end
+
+      private
+
+      def limit_allowed?
+        ["BYTES", "STRING"].include? type
+      end
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/table.rb

@@ -0,0 +1,61 @@
+# The MIT License (MIT)
+#
+# Copyright (c) 2020 Google LLC.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# ITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+# 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/table/column"
+
+module ActiveRecordSpannerAdapter
+  class Table
+    attr_accessor :name, :on_delete, :parent_table, :schema_name, :catalog,
+                  :indexes, :columns, :foreign_keys
+
+    # parent_table == interleave_in
+    def initialize \
+        name,
+        parent_table: nil,
+        on_delete: nil,
+        schema_name: nil,
+        catalog: nil
+      @name = name.to_s
+      @parent_table = parent_table.to_s if parent_table
+      @on_delete = on_delete
+      @schema_name = schema_name
+      @catalog = catalog
+      @columns = []
+      @indexes = []
+      @foreign_keys = []
+    end
+
+    def primary_keys
+      columns.select(&:primary_key).map(&:name)
+    end
+
+    def cascade?
+      @on_delete == "CASCADE"
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/transaction.rb

@@ -0,0 +1,154 @@
+# 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 ActiveRecordSpannerAdapter
+  class Transaction
+    attr_reader :state
+
+    def initialize connection, isolation
+      @connection = connection
+      @isolation = isolation
+      @committable = ![:read_only, :pdml].include?(isolation) && !isolation.is_a?(Hash)
+      @state = :INITIALIZED
+      @sequence_number = 0
+      @mutations = []
+    end
+
+    def active?
+      @state == :STARTED
+    end
+
+    def isolation
+      return nil unless active?
+      @isolation
+    end
+
+    def buffer mutation
+      @mutations << mutation
+    end
+
+    # Begins the transaction.
+    #
+    # Read-only and PDML transactions are started by executing a BeginTransaction RPC.
+    # Read/write transactions are not really started by this method, and instead a
+    # transaction selector is prepared that will be included with the first statement
+    # on the transaction.
+    def begin
+      raise "Nested transactions are not allowed" if @state != :INITIALIZED
+      begin
+        case @isolation
+        when Hash
+          if @isolation[:timestamp]
+            @grpc_transaction = @connection.session.create_snapshot timestamp: @isolation[:timestamp]
+          elsif @isolation[:staleness]
+            @grpc_transaction = @connection.session.create_snapshot staleness: @isolation[:staleness]
+          elsif @isolation[:strong]
+            @grpc_transaction = @connection.session.create_snapshot strong: true
+          else
+            raise "Invalid snapshot argument: #{@isolation}"
+          end
+        when :read_only
+          @grpc_transaction = @connection.session.create_snapshot strong: true
+        when :pdml
+          @grpc_transaction = @connection.session.create_pdml
+        else
+          @begin_transaction_selector = Google::Spanner::V1::TransactionSelector.new \
+            begin: Google::Spanner::V1::TransactionOptions.new(
+              read_write: Google::Spanner::V1::TransactionOptions::ReadWrite.new
+            )
+
+        end
+        @state = :STARTED
+      rescue Google::Cloud::NotFoundError => e
+        if @connection.session_not_found? e
+          @connection.reset!
+          retry
+        end
+        @state = :FAILED
+        raise
+      rescue StandardError
+        @state = :FAILED
+        raise
+      end
+    end
+
+    # Forces a BeginTransaction RPC for a read/write transaction. This is used by a
+    # connection if the first statement of a transaction failed.
+    def force_begin_read_write
+      @grpc_transaction = @connection.session.create_transaction
+    end
+
+    def next_sequence_number
+      @sequence_number += 1 if @committable
+    end
+
+    def commit
+      raise "This transaction is not active" unless active?
+
+      begin
+        # Start a transaction with an explicit BeginTransaction RPC if the transaction only contains mutations.
+        force_begin_read_write if @committable && !@mutations.empty? && !@grpc_transaction
+
+        @connection.session.commit_transaction @grpc_transaction, @mutations if @committable && @grpc_transaction
+        @state = :COMMITTED
+      rescue Google::Cloud::NotFoundError => e
+        if @connection.session_not_found? e
+          shoot_and_forget_rollback
+          @connection.reset!
+          @connection.raise_aborted_err
+        end
+        @state = :FAILED
+        raise
+      rescue StandardError
+        @state = :FAILED
+        raise
+      end
+    end
+
+    def rollback
+      # Allow rollback after abort and/or a failed commit.
+      raise "This transaction is not active" unless active? || @state == :FAILED || @state == :ABORTED
+      if active? && @grpc_transaction
+        # We do a shoot-and-forget rollback here, as the error that caused the transaction to be rolled back could
+        # also have invalidated the transaction (e.g. `Session not found`). If the rollback fails for any other
+        # reason, we also do not need to retry it or propagate the error to the application, as the transaction will
+        # automatically be aborted by Cloud Spanner after 10 seconds anyways.
+        shoot_and_forget_rollback
+      end
+      @state = :ROLLED_BACK
+    end
+
+    def shoot_and_forget_rollback
+      @connection.session.rollback @grpc_transaction.transaction_id if @committable
+    rescue StandardError # rubocop:disable Lint/HandleExceptions
+      # Ignored
+    end
+
+    def mark_aborted
+      @state = :ABORTED
+    end
+
+    # Sets the underlying gRPC transaction to use for this Transaction.
+    # This is used for queries/DML statements that inlined the BeginTransaction option and returned
+    # a transaction in the metadata.
+    def grpc_transaction= grpc
+      @grpc_transaction = Google::Cloud::Spanner::Transaction.from_grpc grpc, @connection.session
+    end
+
+    def transaction_selector
+      return unless active?
+
+      # Use the transaction that has been started by a BeginTransaction RPC or returned by a
+      # statement, if present.
+      return Google::Spanner::V1::TransactionSelector.new id: @grpc_transaction.transaction_id \
+          if @grpc_transaction
+
+      # Return a transaction selector that will instruct the statement to also start a transaction
+      # and return its id as a side effect.
+      @begin_transaction_selector
+    end
+  end
+end

data/lib/activerecord_spanner_adapter/version.rb

@@ -0,0 +1,9 @@
+# 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.
+
+module ActiveRecordSpannerAdapter
+  VERSION = "1.0.0".freeze
+end