activerecord-spanner-adapter 0.3.0 → 1.0.0

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

@@ -1,335 +1,313 @@
+# 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 ConnectionAdapters
     module Spanner
       module DatabaseStatements
-        include ConnectionAdapters::DatabaseStatements
-
-        class QueryVisitor < ::Arel::Visitors::ToSql
-          def visit_Arel_Nodes_BindParam(o, collector)
-            collector.add_bind(o) {|bind_idx| "@p#{bind_idx}" }
-          end
-        end
+        # DDL, DML and DQL Statements
 
-        # A mixin which provides ways to resolve rvalues.
-        module RightValueResolveable
-          def visit_Arel_Nodes_BindParam(o)
-            binds.shift
-          end
+        def execute sql, name = nil, binds = []
+          statement_type = sql_statement_type sql
 
-          def visit_Array(o)
-            o.map {|item| accept(item) }
+          if preventing_writes? && [:dml, :ddl].include?(statement_type)
+            raise ActiveRecord::ReadOnlyError(
+              "Write query attempted while in readonly mode: #{sql}"
+            )
           end
 
-          def visit_Arel_Nodes_Casted(o)
-            a = o.attribute
-            if a.able_to_type_cast?
-              a.type_cast_for_database(o.val)
-            else
-              o.val
-            end
-          end
-
-          private
-          # To be overridden
-          def binds
-            raise NotImplementedError
-          end
-        end
-
-        # Converts ASTs of INSERT, UPDATE or DELETE statements into forms
-        # convenient for DatabaseStatements#insert, #update and #delete.
-        class MutationVisitor < ::Arel::Visitors::Visitor
-          include RightValueResolveable
-
-          def initialize(schema_reader, binds)
-            super()
-            @schema_reader = schema_reader
-            @binds = binds
-          end
-
-          attr_reader :binds
-          private :binds
-
-          def visit_Arel_Nodes_InsertStatement(o)
-            raise NotImplementedError, 'INSERT INTO SELECT statement is not supported' if o.select
-            table = o.relation.name
-            columns = if o.columns.any?
-                        o.columns.map(&:name)
-                      else
-                        columns(table).map(&:name)
-                      end
-            values = o.values ? accept(o.values) : []
-
-            [table, columns, values]
-          end
-
-          def visit_Arel_Nodes_UpdateStatement(o)
-            table = o.relation.name
-            values = accept(o.values)
-            pk = @schema_reader.primary_key(table)
-
-            unless o.orders.empty? and o.limit.nil? and o.wheres.size == 1
-              raise NotImplementedError, 'UPDATE statements with ORDER, LIMIT or complex WHERE clause'
-            end
-
-            ids = WhereVisitor.new(o.relation, pk, binds).accept(o.wheres[0])
-            raise NotImplementedError 'UPDATE statements with complex WHERE clause' unless ids
-
-            [table, ids, values]
-          end
-
-          def visit_Arel_Nodes_DeleteStatement(o)
-            table = o.relation.name
-
-            # fallback_result lets the caller query the target id set at first and then
-            # delete the ids.
-            fallback_result = [table, nil, o.wheres]
-
-            case o.wheres.size
-            when 0
-              return [table, :all]
-            when 1
-              # it might be a simple "id = ?". Let's check later
-            else
-              return fallback_result
-            end
-
-            pk = @schema_reader.primary_key(table)
-            ids = WhereVisitor.new(o.relation, pk, binds).accept(o.wheres[0])
-
-            if ids
-              return [table, ids]
-            else
-              return fallback_result
+          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
-          end
 
-          def visit_Arel_Nodes_Values o
-            o.expressions.map.with_index do |value|
-              case value
-              when ::Arel::Nodes::SqlLiteral
-                raise NotImplementedError, "mutation with SQL literal is not supported"
-              else
-                accept(value)
+            log sql, name 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
+                  end
+                else
+                  @connection.execute_query sql, params: params, types: types, single_use_selector: selector
+                end
               end
             end
           end
+        end
 
-          def visit_Arel_Nodes_Assignment(o)
-            [accept(o.left), accept(o.right)]
-          end
-
-          def visit_Arel_Nodes_UnqualifiedColumn(o)
-            o.name
-          end
+        def query sql, name = nil
+          exec_query sql, name
         end
 
-        # Tries to convert where clause into a set of ids if it is simple enough.
-        # Returns nil if the clause is not simple.
-        class WhereVisitor < ::Arel::Visitors::Visitor
-          include RightValueResolveable
+        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
 
-          NOT_SIMPLE = nil
+        def exec_mutation mutation
+          @connection.current_transaction.buffer mutation
+        end
 
-          def initialize(relation, pk, binds)
-            super()
-            @relation = relation
-            @pk = pk
-            @binds = binds
+        def update arel, name = nil, binds = []
+          # Add a `WHERE TRUE` if it is an update_all or delete_all call that uses DML.
+          if !should_use_mutation(arel) && arel.respond_to?(:ast) && arel.ast.wheres.empty?
+            arel.ast.wheres << Arel::Nodes::SqlLiteral.new("TRUE")
           end
+          return super unless should_use_mutation arel
 
-          attr_reader :binds
-          private :binds
+          raise "Unsupported update for use with mutations: #{arel}" unless arel.is_a? Arel::DeleteManager
 
-          def visit_Arel_Nodes_And(o)
-            if o.children.size == 1
-              accept(o.left)
-            else
-              NOT_SIMPLE
-            end
-          end
+          exec_mutation create_delete_all_mutation arel if arel.is_a? Arel::DeleteManager
+          0 # Affected rows (unknown)
+        end
+        alias delete update
+
+        def exec_update sql, name = "SQL", binds = []
+          result = execute sql, name, binds
+          # Make sure that we consume the entire result stream before trying to get the stats.
+          # This is required because the ExecuteStreamingSql RPC is also used for (Partitioned) DML,
+          # and this RPC can return multiple partial result sets for DML as well. Only the last partial
+          # result set will contain the statistics. Although there will never be any rows, this makes
+          # sure that the stream is fully consumed.
+          result.rows.each { |_| }
+          return result.row_count if result.row_count
+
+          raise ActiveRecord::StatementInvalid.new(
+            "DML statement is invalid.", sql: sql
+          )
+        end
+        alias exec_delete exec_update
 
-          def visit_Arel_Nodes_Equality(o)
-            if pk_cond?(o)
-              [accept(o.right)]
-            else
-              NOT_SIMPLE
+        def truncate table_name, name = nil
+          Array(table_name).each do |t|
+            log "TRUNCATE #{t}", name do
+              @connection.truncate t
             end
           end
+        end
 
-          def visit_Arel_Nodes_In(o)
-            return nil unless pk_cond?(o)
+        def write_query? sql
+          sql_statement_type(sql) == :dml
+        end
 
-            if o.kind_of?(Array) and o.empty?
-              []
-            else
-              accept(o.right)
+        def execute_ddl statements
+          log "MIGRATION", "SCHEMA" do
+            ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
+              @connection.execute_ddl statements
             end
           end
+        rescue Google::Cloud::Error => error
+          raise ActiveRecord::StatementInvalid, error
+        end
+
+        # Transaction
 
-          def unsupported(o)
-            return NOT_SIMPLE
+        def transaction requires_new: nil, isolation: nil, joinable: true
+          if !requires_new && current_transaction.joinable?
+            return super
           end
 
-          alias visit_Arel_Nodes_Grouping unsupported
-          alias visit_Arel_Nodes_NotIn unsupported
-          alias visit_Arel_Nodes_Or unsupported
-          alias visit_Arel_Nodes_NotEqual unsupported
-          alias visit_Arel_Nodes_Case unsupported
-          alias visit_Arel_Nodes_Between unsupported
-          alias visit_Arel_Nodes_GreaterThanOrEqual unsupported
-          alias visit_Arel_Nodes_GreaterThan unsupported
-          alias visit_Arel_Nodes_LessThanOrEqual unsupported
-          alias visit_Arel_Nodes_LessThan unsupported
-          alias visit_Arel_Nodes_Matches unsupported
-          alias visit_Arel_Nodes_DoesNotMatch unsupported
-
-          private
-          def pk_cond?(o)
-            o.left.kind_of?(Arel::Attributes::Attribute) &&
-              o.left.relation == @relation &&
-              o.left.name == @pk
+          backoff = 0.2
+          begin
+            super
+          rescue ActiveRecord::StatementInvalid => err
+            if err.cause.is_a? Google::Cloud::AbortedError
+              sleep(delay_from_aborted(err) || backoff *= 1.3)
+              retry
+            end
+            raise
           end
         end
 
-        def insert(arel, name = nil, pk = nil, id_value = nil, sequence_name = nil, binds = [])
-          raise NotImplementedError, "INSERT in raw SQL is not supported" unless arel.respond_to?(:ast)
-
-          type_casted_binds = binds.map {|attr| type_cast(attr.value_for_database) }
-          table, columns, values = MutationVisitor.new(self, type_casted_binds).accept(arel.ast)
-          fake_sql = <<~"SQL"
-            INSERT INTO #{table}(#{columns.join(", ")}) VALUES (#{values.join(", ")})
-          SQL
-
-          row = columns.zip(values).inject({}) {|out, (col, value)|
-            out[col] = value
-            out
+        def transaction_isolation_levels
+          {
+            read_uncommitted:   "READ UNCOMMITTED",
+            read_committed:     "READ COMMITTED",
+            repeatable_read:    "REPEATABLE READ",
+            serializable:       "SERIALIZABLE",
+
+            # These are not really isolation levels, but it is the only (best) way to pass in additional
+            # transaction options to the connection.
+            read_only:          "READ_ONLY",
+            buffered_mutations: "BUFFERED_MUTATIONS"
           }
+        end
 
-          log(fake_sql, name) do
-            with_phase_transition {|client| client.insert(table, row) }
+        def begin_db_transaction
+          log "BEGIN" do
+            @connection.begin_transaction
           end
-
-          id_value
         end
 
-        def update(arel, name = nil, binds = [])
-          raise NotImplementedError, "UPDATE statement in raw SQL is not supported" unless arel.respond_to?(:ast)
-
-          type_casted_binds = binds.map {|attr| type_cast(attr.value_for_database) }
-          table, target, values = MutationVisitor.new(self, type_casted_binds.dup).accept(arel.ast)
-
-          fake_set = values.map {|col, val|
-            "#{quote_column_name(col)} = #{quote(val)}"
-          }.join(', ')
-
-          pk = primary_key(table)
-          if target.size > 1
-            fake_sql = "UPDATE #{quote_column_name(table)} SET #{fake_set} WHERE #{pk} IN ?"
+        # Begins a transaction on the database with the specified isolation level. Cloud Spanner only supports
+        # isolation level :serializable, but also defines three additional 'isolation levels' that can be used
+        # to start specific types of Spanner transactions:
+        # * :read_only: Starts a read-only snapshot transaction using a strong timestamp bound.
+        # * :buffered_mutations: Starts a read/write transaction that will use mutations instead of DML for single-row
+        #                        inserts/updates/deletes. Mutations are buffered locally until the transaction is
+        #                        committed, and any changes during a transaction cannot be read by the application.
+        # * :pdml: Starts a Partitioned DML transaction. Executing multiple DML statements in one PDML transaction
+        #          block is NOT supported A PDML transaction is not guaranteed to be atomic.
+        #          See https://cloud.google.com/spanner/docs/dml-partitioned for more information.
+        #
+        # In addition to the above, a Hash containing read-only snapshot options may be used to start a specific
+        # read-only snapshot:
+        # * { timestamp: Time } Starts a read-only snapshot at the given timestamp.
+        # * { staleness: Integer } Starts a read-only snapshot with the given staleness in seconds.
+        # * { strong: <any value>} Starts a read-only snapshot with strong timestamp bound
+        #                          (this is the same as :read_only)
+        #
+        def begin_isolated_db_transaction isolation
+          if isolation.is_a? Hash
+            raise "Unsupported isolation level: #{isolation}" unless \
+              isolation[:timestamp] || isolation[:staleness] || isolation[:strong]
+            raise "Only one option is supported. It must be one of `timestamp`, `staleness` or `strong`." \
+              if isolation.count != 1
           else
-            fake_sql = "UPDATE #{quote_column_name(table)} SET #{fake_set} WHERE #{pk} = ?"
+            raise "Unsupported isolation level: #{isolation}" unless \
+              [:serializable, :read_only, :buffered_mutations, :pdml].include? isolation
           end
 
-          row = values.inject({}) {|r, (col, val)| r[col] = val; r }
-          rows = target.map {|id| row.merge(pk => id) }
-
-          log(fake_sql, name, binds) do
-            with_phase_transition {|client| client.update(table, rows) }
+          log "BEGIN #{isolation}" do
+            @connection.begin_transaction isolation
           end
-          target.size
         end
 
-        def delete(arel, name, binds)
-          raise NotImplementedError, "DELETE statement in raw SQL is not supported" unless arel.respond_to?(:ast)
-
-          type_casted_binds = binds.map {|attr| type_cast(attr.value_for_database) }
-          table, target, wheres = MutationVisitor.new(self, type_casted_binds.dup).accept(arel.ast)
-
-          # TODO(yugui) Support composite primary key?
-          pk = primary_key(table)
-          if target.nil?
-            where_clause = visitor.accept(wheres, collector).compile(binds.dup, self)
-            # TODO(yugui) keep consistency with transaction
-            target = select_values(<<~"SQL", name, binds)
-              SELECT #{quote_column_name(pk)} FROM #{quote_table_name(table)} WHERE #{where_clause}
-            SQL
+        def commit_db_transaction
+          log "COMMIT" do
+            @connection.commit_transaction
           end
+        end
 
-          if target == :all
-            keyset = []
-            fake_sql = "DELETE FROM #{quote_column_name(table)}"
-          elsif target.size > 1
-            fake_sql = "DELETE FROM #{quote_column_name(table)} WHERE (primary-key) IN ?"
-            keyset = target
-          else
-            fake_sql = "DELETE FROM #{quote_column_name(table)} WHERE (primary-key) = ?"
-            keyset = target
+        def rollback_db_transaction
+          log "ROLLBACK" do
+            @connection.rollback_transaction
           end
+        end
 
-          log(fake_sql, name, binds) do
-            with_phase_transition {|client| client.delete(table, keyset) }
-          end
+        private
 
-          keyset.size
+        # Translates binds to Spanner types and params.
+        def to_types_and_params binds
+          types = binds.enum_for(:each_with_index).map do |bind, i|
+            type = :INT64
+            if bind.respond_to? :type
+              type = ActiveRecord::Type::Spanner::SpannerActiveRecordConverter
+                     .convert_active_model_type_to_spanner(bind.type)
+            end
+            [
+              # Generates binds for named parameters in the format `@p1, @p2, ...`
+              "p#{i + 1}", type
+            ]
+          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
+
+            ["p#{i + 1}", value]
+          end.to_h
+          [types, params]
         end
 
-        def execute(stmt)
-          case stmt
-          when Spanner::DDL
-            execute_ddl(stmt)
-          else
-            super(stmt)
-          end
+        # An insert/update/delete statement could use mutations in some specific circumstances.
+        # This method returns an indication whether a specific operation should use mutations instead of DML
+        # based on the operation itself, and the current transaction.
+        def should_use_mutation arel
+          !@connection.current_transaction.nil? \
+            && @connection.current_transaction.isolation == :buffered_mutations \
+            && can_use_mutation(arel) \
         end
 
-        def exec_query(sql, name = 'SQL', binds = [], prepare: :ignored)
-          case
-          when binds.kind_of?(Hash)
-            spanner_binds = binds
-            binds = binds.values
-          when binds.respond_to?(:to_hash)
-            spanner_binds = binds.to_hash
-            binds = spanner_binds.values
-          else
-            spanner_binds = binds.each_with_index.inject({}) {|b, (attr, i)|
-              b["p#{i+1}"] = type_cast(attr.value_for_database)
-              b
-            }
-          end
+        def can_use_mutation arel
+          return true if arel.is_a?(Arel::DeleteManager) && arel.respond_to?(:ast) && arel.ast.wheres.empty?
+          false
+        end
 
-          log(sql, name, binds) do
-            executor = client
-            executor = raw_client if name == 'SCHEMA'
-            results = executor.execute(sql, params: spanner_binds) 
-            columns = results.fields.keys
-            rows = results.rows.map(&:to_a)
-            ActiveRecord::Result.new(columns.map(&:to_s), rows)
+        def create_delete_all_mutation arel
+          unless arel.is_a? Arel::DeleteManager
+            raise "A delete mutation can only be created from a DeleteManager"
+          end
+          # Check if it is a delete_all operation.
+          unless arel.ast.wheres.empty?
+            raise "A delete mutation can only be created without a WHERE clause"
+          end
+          table_name = arel.ast.relation.name if arel.ast.relation.is_a? Arel::Table
+          table_name = arel.ast.relation.left.name if arel.ast.relation.is_a? Arel::Nodes::JoinSource
+          unless table_name
+            raise "Could not find table for delete mutation"
           end
-        end
 
-        def begin_db_transaction
-          with_phase_transition {|client| client.begin_transaction }
+          Google::Cloud::Spanner::V1::Mutation.new(
+            delete: Google::Cloud::Spanner::V1::Mutation::Delete.new(
+              table: table_name,
+              key_set: { all: true }
+            )
+          )
         end
 
-        def begin_isolated_db_transaction(isolation)
-          with_phase_transition {|client| client.begin_snapshot(**isolation) }
-        end
+        COMMENT_REGEX = %r{(?:--.*\n)*|/\*(?:[^*]|\*[^/])*\*/}m.freeze \
+            unless defined? ActiveRecord::ConnectionAdapters::AbstractAdapter::COMMENT_REGEX
+        COMMENT_REGEX = ActiveRecord::ConnectionAdapters::AbstractAdapter::COMMENT_REGEX \
+            if defined? ActiveRecord::ConnectionAdapters::AbstractAdapter::COMMENT_REGEX
 
-        def commit_db_transaction
-          with_phase_transition {|client| client.commit }
+        private_class_method def self.build_sql_statement_regexp *parts # :nodoc:
+          parts = parts.map { |part| /#{part}/i }
+          /\A(?:[\(\s]|#{COMMENT_REGEX})*#{Regexp.union(*parts)}/
         end
 
-        def exec_rollback_db_transaction
-          with_phase_transition {|client| client.rollback }
+        DDL_REGX = build_sql_statement_regexp(:create, :alter, :drop).freeze
+
+        DML_REGX = build_sql_statement_regexp(:insert, :delete, :update).freeze
+
+        def sql_statement_type sql
+          case sql
+          when DDL_REGX
+            :ddl
+          when DML_REGX
+            :dml
+          else
+            :dql
+          end
         end
 
-        def with_phase_transition
-          result = yield client
-          @client = client.next
-          result
+        ##
+        # 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
   end
 end
-