activerecord-cockroachdb-adapter 7.0.2 → 7.1.0

data/lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -1,5 +1,6 @@
 require "rgeo/active_record"
 
+require_relative "../../arel/nodes/join_source_ext"
 require "active_record/connection_adapters/postgresql_adapter"
 require "active_record/connection_adapters/cockroachdb/attribute_methods"
 require "active_record/connection_adapters/cockroachdb/column_methods"
@@ -23,6 +24,8 @@ require "active_record/connection_adapters/cockroachdb/arel_tosql"
 require_relative "../migration/cockroachdb/compatibility"
 require_relative "../../version"
 
+require_relative "../relation/query_methods_ext"
+
 # Run to ignore spatial tables that will break schemna dumper.
 # Defined in ./setup.rb
 ActiveRecord::ConnectionAdapters::CockroachDB.initial_setup
@@ -158,7 +161,7 @@ module ActiveRecord
       end
 
       def supports_bulk_alter?
-        false
+        true
       end
 
       def supports_json?
@@ -179,7 +182,15 @@ module ActiveRecord
       end
 
       def supports_partial_index?
-        @crdb_version >= 2020
+        true
+      end
+
+      def supports_index_include?
+        false
+      end
+
+      def supports_exclusion_constraints?
+        false
       end
 
       def supports_expression_index?
@@ -194,7 +205,7 @@ module ActiveRecord
       end
 
       def supports_comments?
-        @crdb_version >= 2010
+        true
       end
 
       def supports_comments_in_create?
@@ -206,11 +217,11 @@ module ActiveRecord
       end
 
       def supports_virtual_columns?
-        @crdb_version >= 2110
+        true
       end
 
       def supports_string_to_array_coercion?
-        @crdb_version >= 2020
+        true
       end
 
       def supports_partitioned_indexes?
@@ -236,62 +247,30 @@ module ActiveRecord
       alias index_name_length max_identifier_length
       alias table_alias_length max_identifier_length
 
-      def initialize(connection, logger, conn_params, config)
-        super(connection, logger, conn_params, config)
-
-        # crdb_version is the version of the binary running on the node. We
-        # really want to use `SHOW CLUSTER SETTING version` to get the cluster
-        # version, but that is only available to admins. Instead, we can use
-        # crdb_internal.is_at_least_version, but that's only available in 22.1.
-        crdb_version_string = query_value("SHOW crdb_version")
-        if crdb_version_string.include? "v22.1"
-          version_num = query_value(<<~SQL, "VERSION")
-            SELECT
-              CASE
-              WHEN crdb_internal.is_at_least_version('22.2') THEN 2220
-              WHEN crdb_internal.is_at_least_version('22.1') THEN 2210
-              ELSE 2120
-              END;
-          SQL
-        else
-          # This branch can be removed once the dialect stops supporting v21.2
-          # and earlier.
-          if crdb_version_string.include? "v1."
-            version_num = 1
-          elsif crdb_version_string.include? "v2."
-            version_num 2
-          elsif crdb_version_string.include? "v19.1."
-            version_num = 1910
-          elsif crdb_version_string.include? "v19.2."
-            version_num = 1920
-          elsif crdb_version_string.include? "v20.1."
-            version_num = 2010
-          elsif crdb_version_string.include? "v20.2."
-            version_num = 2020
-          elsif crdb_version_string.include? "v21.1."
-            version_num = 2110
-          else
-            version_num = 2120
-          end
-        end
-        @crdb_version = version_num.to_i
-
-        # NOTE: this is normally in configure_connection, but that is run
-        # before crdb_version is determined. Once all supported versions
-        # of CockroachDB support SET intervalstyle it can safely be moved
-        # back.
-        # Set interval output format to ISO 8601 for ease of parsing by ActiveSupport::Duration.parse
-        if @crdb_version >= 2120
-          begin
-            execute("SET intervalstyle_enabled = true", "SCHEMA")
-            execute("SET intervalstyle = iso_8601", "SCHEMA")
-          rescue
-            # Ignore any error. This can happen with a cluster that has
-            # not yet finalized the v21.2 upgrade. v21.2 does not have
-            # a way to tell if the upgrade was finalized (see comment above).
-          end
-        end
-      end
+      # NOTE: This commented bit of code allows to have access to crdb version,
+      # which can be useful for feature detection. However, we currently don't
+      # need, hence we avoid the extra queries.
+      #
+      #   def initialize(connection, logger, conn_params, config)
+      #     super(connection, logger, conn_params, config)
+
+      #     # crdb_version is the version of the binary running on the node. We
+      #     # really want to use `SHOW CLUSTER SETTING version` to get the cluster
+      #     # version, but that is only available to admins. Instead, we can use
+      #     # crdb_internal.is_at_least_version, but that's only available in 22.1.
+      #     crdb_version_string = query_value("SHOW crdb_version")
+      #     if crdb_version_string.include? "v22.1"
+      #       version_num = query_value(<<~SQL, "VERSION")
+      #         SELECT
+      #           CASE
+      #           WHEN crdb_internal.is_at_least_version('22.2') THEN 2220
+      #           WHEN crdb_internal.is_at_least_version('22.1') THEN 2210
+      #           ELSE 2120
+      #           END;
+      #       SQL
+      #     end
+      #     @crdb_version = version_num.to_i
+      #   end
 
       def self.database_exists?(config)
         !!ActiveRecord::Base.cockroachdb_connection(config)
@@ -304,12 +283,12 @@ module ActiveRecord
       # (DO $$) that CockroachDB does not support.
       #
       # Given a name and an array of values, creates an enum type.
-      def create_enum(name, values)
-        sql_values = values.map { |s| "'#{s}'" }.join(", ")
+      def create_enum(name, values, **options)
+        sql_values = values.map { |s| quote(s) }.join(", ")
         query = <<~SQL
-          CREATE TYPE IF NOT EXISTS \"#{name}\" AS ENUM (#{sql_values});
+          CREATE TYPE IF NOT EXISTS #{quote_table_name(name)} AS ENUM (#{sql_values});
         SQL
-        exec_query(query)
+        internal_exec_query(query).tap { reload_type_map }
       end
 
       class << self
@@ -367,56 +346,6 @@ module ActiveRecord
 
       private
 
-        # Configures the encoding, verbosity, schema search path, and time zone of the connection.
-        # This is called by #connect and should not be called manually.
-        #
-        # NOTE(joey): This was cradled from postgresql_adapter.rb. This
-        # was due to needing to override configuration statements.
-        def configure_connection
-          if @config[:encoding]
-            @connection.set_client_encoding(@config[:encoding])
-          end
-          self.client_min_messages = @config[:min_messages] || "warning"
-          self.schema_search_path = @config[:schema_search_path] || @config[:schema_order]
-
-          # Use standard-conforming strings so we don't have to do the E'...' dance.
-          set_standard_conforming_strings
-
-          variables = @config.fetch(:variables, {}).stringify_keys
-
-          # If using Active Record's time zone support configure the connection to return
-          # TIMESTAMP WITH ZONE types in UTC.
-          unless variables["timezone"]
-            if ActiveRecord.default_timezone == :utc
-              variables["timezone"] = "UTC"
-            elsif @local_tz
-              variables["timezone"] = @local_tz
-            end
-          end
-
-          # NOTE(joey): This is a workaround as CockroachDB 1.1.x
-          # supports SET TIME ZONE <...> and SET "time zone" = <...> but
-          # not SET timezone = <...>.
-          if variables.key?("timezone")
-            tz = variables.delete("timezone")
-            execute("SET TIME ZONE #{quote(tz)}", "SCHEMA")
-          end
-
-          # SET statements from :variables config hash
-          # https://www.postgresql.org/docs/current/static/sql-set.html
-          variables.map do |k, v|
-            if v == ":default" || v == :default
-              # Sets the value to the global or compile default
-
-              # NOTE(joey): I am not sure if simply commenting this out
-              # is technically correct.
-              # execute("SET #{k} = DEFAULT", "SCHEMA")
-            elsif !v.nil?
-              execute("SET SESSION #{k} = #{quote(v)}", "SCHEMA")
-            end
-          end
-        end
-
         # Override extract_value_from_default because the upstream definition
         # doesn't handle the variations in CockroachDB's behavior.
         def extract_value_from_default(default)
@@ -500,7 +429,8 @@ module ActiveRecord
               SELECT a.attname, format_type(a.atttypid, a.atttypmod),
                      pg_get_expr(d.adbin, d.adrelid), a.attnotnull, a.atttypid, a.atttypmod,
                      c.collname, NULL AS comment,
-                     #{supports_virtual_columns? ? 'attgenerated' : quote('')} as attgenerated,
+                     attidentity,
+                     attgenerated,
                      NULL as is_hidden
                 FROM pg_attribute a
                 LEFT JOIN pg_attrdef d ON a.attrelid = d.adrelid AND a.attnum = d.adnum
@@ -515,26 +445,31 @@ module ActiveRecord
 
           # Use regex comparison because if a type is an array it will
           # have [] appended to the end of it.
-          target_types = [
-            /geometry/,
-            /geography/,
-            /interval/,
-            /numeric/
-          ]
-
-          re = Regexp.union(target_types)
+          re = /\A(?:geometry|geography|interval|numeric)/
+
+          # 0: attname
+          # 1: type
+          # 2: default
+          # 3: attnotnull
+          # 4: atttypid
+          # 5: atttypmod
+          # 6: collname
+          # 7: comment
+          # 8: attidentity
+          # 9: attgenerated
+          # 10: is_hidden
           fields.map do |field|
             dtype = field[1]
             field[1] = crdb_fields[field[0]][2].downcase if re.match(dtype)
             field[7] = crdb_fields[field[0]][1]&.gsub!(/^\'|\'?$/, '')
-            field[9] = true if crdb_fields[field[0]][3]
+            field[10] = true if crdb_fields[field[0]][3]
             field
           end
           fields.delete_if do |field|
             # Don't include rowid column if it is hidden and the primary key
             # is not defined (meaning CRDB implicitly created it).
             if field[0] == CockroachDBAdapter::DEFAULT_PRIMARY_KEY
-              field[9] && !primary_key(table_name)
+              field[10] && !primary_key(table_name)
             else
               false # Keep this entry.
             end
@@ -547,11 +482,14 @@ module ActiveRecord
         # precision, and scale information in the type.
         # Ex. geometry -> geometry(point, 4326)
         def crdb_column_definitions(table_name)
+          table_name = PostgreSQL::Utils.extract_schema_qualified_name(table_name)
+          table = table_name.identifier
+          with_schema = " AND c.table_schema = #{quote(table_name.schema)}" if table_name.schema
           fields = \
           query(<<~SQL, "SCHEMA")
             SELECT c.column_name, c.column_comment, c.crdb_sql_type, c.is_hidden::BOOLEAN
-              FROM information_schema.columns c
-            WHERE c.table_name = #{quote(table_name)}
+            FROM information_schema.columns c
+            WHERE c.table_name = #{quote(table)}#{with_schema}
           SQL
 
           fields.reduce({}) do |a, e|
@@ -592,21 +530,10 @@ module ActiveRecord
         def load_additional_types(oids = nil)
           if @config[:use_follower_reads_for_type_introspection]
             initializer = OID::TypeMapInitializer.new(type_map)
-
-            query = <<~SQL
-              SELECT t.oid, t.typname, t.typelem, t.typdelim, t.typinput, r.rngsubtype, t.typtype, t.typbasetype
-              FROM pg_type as t
-              LEFT JOIN pg_range as r ON oid = rngtypid AS OF SYSTEM TIME '-10s'
-            SQL
-
-            if oids
-              query += "WHERE t.oid IN (%s)" % oids.join(", ")
-            else
-              query += initializer.query_conditions_for_initial_load
-            end
-
-            execute_and_clear(query, "SCHEMA", []) do |records|
-              initializer.run(records)
+            load_types_queries_with_aost(initializer, oids) do |query|
+              execute_and_clear(query, "SCHEMA", [], allow_retry: true, materialize_transactions: false) do |records|
+                initializer.run(records)
+              end
             end
           else
             super
@@ -617,6 +544,21 @@ module ActiveRecord
           super
         end
 
+        def load_types_queries_with_aost(initializer, oids)
+          query = <<~SQL
+            SELECT t.oid, t.typname, t.typelem, t.typdelim, t.typinput, r.rngsubtype, t.typtype, t.typbasetype
+            FROM pg_type as t
+            LEFT JOIN pg_range as r ON oid = rngtypid AS OF SYSTEM TIME '-10s'
+          SQL
+          if oids
+            yield query + "WHERE t.oid IN (%s)" % oids.join(", ")
+          else
+            yield query + initializer.query_conditions_for_known_type_names
+            yield query + initializer.query_conditions_for_known_type_types
+            yield query + initializer.query_conditions_for_array_types
+          end
+        end
+
         # override
         # This method maps data types to their proper decoder.
         #
@@ -647,15 +589,13 @@ module ActiveRecord
               WHERE t.typname IN (%s)
             SQL
 
-            coders = execute_and_clear(query, "SCHEMA", []) do |result|
-              result
-                .map { |row| construct_coder(row, coders_by_name[row["typname"]]) }
-                .compact
+            coders = execute_and_clear(query, "SCHEMA", [], allow_retry: true, materialize_transactions: false) do |result|
+              result.filter_map { |row| construct_coder(row, coders_by_name[row["typname"]]) }
             end
 
             map = PG::TypeMapByOid.new
             coders.each { |coder| map.add_coder(coder) }
-            @connection.type_map_for_results = map
+            @raw_connection.type_map_for_results = map
 
             @type_map_for_results = PG::TypeMapByOid.new
             @type_map_for_results.default_type_map = map

data/lib/active_record/relation/query_methods_ext.rb

@@ -0,0 +1,127 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  class Relation
+    module QueryMethodsExt
+      def aost!(time) # :nodoc:
+        unless time.nil? || time.is_a?(Time)
+          raise ArgumentError, "Unsupported argument type: #{time} (#{time.class})"
+        end
+
+        @aost = time
+        self
+      end
+
+      # Set system time for the current query. Using
+      # `.aost(nil)` resets.
+      #
+      # See cockroachlabs.com/docs/stable/as-of-system-time
+      def aost(time)
+        spawn.aost!(time)
+      end
+
+      def from!(...) # :nodoc:
+        @force_index = nil
+        @index_hint = nil
+        super
+      end
+
+      # Set table index hint for the query to the
+      # given `index_name`, and `direction` (either
+      # `ASC` or `DESC`).
+      #
+      # Any call to `ActiveRecord::QueryMethods#from`
+      # will reset the index hint. Index hints are
+      # not set if the `from` clause is not a table
+      # name.
+      #
+      # @see https://www.cockroachlabs.com/docs/v22.2/table-expressions#force-index-selection
+      def force_index(index_name, direction: nil)
+        spawn.force_index!(index_name, direction: direction)
+      end
+
+      def force_index!(index_name, direction: nil)
+        return self unless from_clause_is_a_table_name?
+
+        index_name = sanitize_sql(index_name.to_s)
+        direction = direction.to_s.upcase
+        direction = %w[ASC DESC].include?(direction) ? ",#{direction}" : ""
+
+        @force_index = "FORCE_INDEX=#{index_name}#{direction}"
+        self.from_clause = build_from_clause_with_hints
+        self
+      end
+
+      # Set table index hint for the query with the
+      # given `hint`. This allows more control over
+      # the hint than `ActiveRecord::Relation#force_index`.
+      # For instance, you could set it to `NO_FULL_SCAN`.
+      #
+      # Any call to `ActiveRecord::QueryMethods#from`
+      # will reset the index hint. Index hints are
+      # not set if the `from` clause is not a table
+      # name.
+      #
+      # @see https://www.cockroachlabs.com/docs/v22.2/table-expressions#force-index-selection
+      def index_hint(hint)
+        spawn.index_hint!(hint)
+      end
+
+      def index_hint!(hint)
+        return self unless from_clause_is_a_table_name?
+
+        hint = sanitize_sql(hint.to_s)
+        @index_hint = hint.to_s
+        self.from_clause = build_from_clause_with_hints
+        self
+      end
+
+      # TODO: reset or no reset?
+
+      def show_create
+        connection.execute("show create table #{connection.quote_table_name self.table_name}").first["create_statement"]
+      end
+
+      private
+
+      def build_arel(...)
+        arel = super
+        arel.aost(@aost) if @aost.present?
+        arel
+      end
+
+      def from_clause_is_a_table_name?
+        # if empty, we are just dealing with the current table.
+        return true if from_clause.empty?
+        # `from_clause` can be a subquery.
+        return false unless from_clause.value.is_a?(String)
+        # `from_clause` can be a list of tables or a function.
+        # A simple way to check is to see if the string
+        # contains special characters. But we have to
+        # not check against an existing table hint.
+        return !from_clause.value.gsub(/\@{.*?\}/, "").match?(/[,\(]/)
+      end
+
+      def build_from_clause_with_hints
+        table_hints = [@index_hint, @force_index].compact.join(",")
+
+        table_name =
+          if from_clause.empty?
+            quoted_table_name
+          else
+            # Remove previous table hints if any. And spaces.
+            from_clause.value.partition("@").first.strip
+          end
+        Relation::FromClause.new("#{table_name}@{#{table_hints}}", nil)
+      end
+    end
+
+    QueryMethods.prepend(QueryMethodsExt)
+  end
+  # `ActiveRecord::Base` ancestors do not include `QueryMethods`.
+  # But the `#all` method returns a relation, which has `QueryMethods`
+  # as ancestor. That is how active_record is doing is as well.
+  #
+  # @see https://github.com/rails/rails/blob/914130a9f/activerecord/lib/active_record/querying.rb#L23
+  Querying.delegate(:force_index, :index_hint, :aost, :show_create, to: :all)
+end

data/lib/activerecord-cockroachdb-adapter.rb

@@ -1,4 +1,4 @@
-if defined?(Rails)
+if defined?(Rails::Railtie)
   module ActiveRecord
     module ConnectionAdapters
       class CockroachDBRailtie < ::Rails::Railtie

data/lib/arel/nodes/join_source_ext.rb

@@ -0,0 +1,28 @@
+module Arel
+  module Nodes
+    module JoinSourceExt
+      def initialize(...)
+        super
+        @aost = nil
+      end
+
+      def hash
+        [*super, aost].hash
+      end
+
+      def eql?(other)
+        super && aost == other.aost
+      end
+      alias_method :==, :eql?
+    end
+    JoinSource.attr_accessor :aost
+    JoinSource.prepend JoinSourceExt
+  end
+  module SelectManagerExt
+    def aost(time)
+      @ctx.source.aost = time
+      nil
+    end
+  end
+  SelectManager.prepend SelectManagerExt
+end

data/lib/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ActiveRecord
-  COCKROACH_DB_ADAPTER_VERSION = "7.0.2"
+  COCKROACH_DB_ADAPTER_VERSION = "7.1.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-cockroachdb-adapter
 version: !ruby/object:Gem::Version
-  version: 7.0.2
+  version: 7.1.0
 platform: ruby
 authors:
 - Cockroach Labs
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-05-30 00:00:00.000000000 Z
+date: 2024-01-30 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -16,14 +16,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 7.0.3
+        version: 7.1.0
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 7.0.3
+        version: 7.1.0
 - !ruby/object:Gem::Dependency
   name: pg
   requirement: !ruby/object:Gem::Requirement
@@ -74,6 +74,8 @@ executables: []
 extensions: []
 extra_rdoc_files: []
 files:
+- ".editorconfig"
+- ".github/workflows/ci.yml"
 - ".github/workflows/docker.yml"
 - ".gitignore"
 - ".gitmodules"
@@ -86,7 +88,10 @@ files:
 - Rakefile
 - activerecord-cockroachdb-adapter.gemspec
 - bin/console
+- bin/console_schemas/default.rb
+- bin/console_schemas/schemas.rb
 - bin/setup
+- bin/start-cockroachdb
 - build/Dockerfile
 - build/config.teamcity.yml
 - build/local-test.sh
@@ -114,7 +119,9 @@ files:
 - lib/active_record/connection_adapters/cockroachdb/type.rb
 - lib/active_record/connection_adapters/cockroachdb_adapter.rb
 - lib/active_record/migration/cockroachdb/compatibility.rb
+- lib/active_record/relation/query_methods_ext.rb
 - lib/activerecord-cockroachdb-adapter.rb
+- lib/arel/nodes/join_source_ext.rb
 - lib/version.rb
 homepage: https://github.com/cockroachdb/activerecord-cockroachdb-adapter
 licenses:
@@ -136,7 +143,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.4.9
+rubygems_version: 3.4.10
 signing_key:
 specification_version: 4
 summary: CockroachDB adapter for ActiveRecord.