activerecord-cockroachdb-adapter 6.0.0beta2 → 6.1.0.pre.beta.3

data/lib/active_record/connection_adapters/cockroachdb/setup.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module ActiveRecord # :nodoc:
+  module ConnectionAdapters # :nodoc:
+    module CockroachDB  # :nodoc:
+      def self.initial_setup
+        ::ActiveRecord::SchemaDumper.ignore_tables |= %w[
+          geography_columns
+          geometry_columns
+          layer
+          raster_columns
+          raster_overviews
+          spatial_ref_sys
+          topology
+        ]
+      end
+    end
+  end
+end

data/lib/active_record/connection_adapters/cockroachdb/spatial_column_info.rb

@@ -0,0 +1,44 @@
+module ActiveRecord
+  module ConnectionAdapters
+    module CockroachDB
+      class SpatialColumnInfo
+        def initialize(adapter, table_name)
+          @adapter = adapter
+          @table_name = table_name
+        end
+
+        def all
+          info = @adapter.query(
+            "SELECT f_geometry_column,coord_dimension,srid,type FROM geometry_columns WHERE f_table_name='#{@table_name}'"
+          )
+          result = {}
+          info.each do |row|
+            name = row[0]
+            type = row[3]
+            dimension = row[1].to_i
+            has_m = !!(type =~ /m$/i)
+            type.sub!(/m$/, '')
+            has_z = dimension > 3 || dimension == 3 && !has_m
+            result[name] = {
+              dimension: dimension,
+              has_m: has_m,
+              has_z: has_z,
+              name: name,
+              srid: row[2].to_i,
+              type: type
+            }
+          end
+          result
+        end
+
+        # do not query the database for non-spatial columns/tables
+        def get(column_name, type)
+          return unless CockroachDBAdapter.spatial_column_options(type.to_sym)
+
+          @spatial_column_info ||= all
+          @spatial_column_info[column_name]
+        end
+      end
+    end
+  end
+end

data/lib/active_record/connection_adapters/cockroachdb/table_definition.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+module ActiveRecord # :nodoc:
+  module ConnectionAdapters # :nodoc:
+    module CockroachDB # :nodoc:
+      class TableDefinition < PostgreSQL::TableDefinition # :nodoc:
+        include ColumnMethods
+
+        # Support for spatial columns in tables
+        # super: https://github.com/rails/rails/blob/master/activerecord/lib/active_record/connection_adapters/abstract/schema_definitions.rb
+        def new_column_definition(name, type, **options)
+          if (info = CockroachDBAdapter.spatial_column_options(type.to_sym))
+            if (limit = options.delete(:limit)) && limit.is_a?(::Hash)
+              options.merge!(limit)
+            end
+
+            geo_type = ColumnDefinitionUtils.geo_type(options[:type] || type || info[:type])
+            base_type = info[:type] || (options[:geographic] ? :geography : :geometry)
+
+            options[:limit] = ColumnDefinitionUtils.limit_from_options(geo_type, options)
+            options[:spatial_type] = geo_type
+            column = super(name, base_type, **options)
+          else
+            column = super(name, type, **options)
+          end
+
+          column
+        end
+      end
+
+      module ColumnDefinitionUtils
+        class << self
+          def geo_type(type = 'GEOMETRY')
+            g_type = type.to_s.delete('_').upcase
+            return 'POINT' if g_type == 'STPOINT'
+            return 'POLYGON' if g_type == 'STPOLYGON'
+
+            g_type
+          end
+
+          def limit_from_options(type, options = {})
+            spatial_type = geo_type(type)
+            spatial_type << 'Z' if options[:has_z]
+            spatial_type << 'M' if options[:has_m]
+            spatial_type << ",#{options[:srid] || default_srid(options)}"
+            spatial_type
+          end
+
+          def default_srid(options)
+            options[:geographic] ? 4326 : CockroachDBAdapter::DEFAULT_SRID
+          end
+        end
+      end
+    end
+  end
+end

data/lib/active_record/connection_adapters/cockroachdb/transaction_manager.rb

@@ -9,7 +9,7 @@ module ActiveRecord
         # transactions will be retried until they pass or the max retry limit is
         # exceeded.
         def within_new_transaction(isolation: nil, joinable: true, attempts: 0)
-          super
+          super(isolation: isolation, joinable: joinable)
         rescue ActiveRecord::StatementInvalid => error
           raise unless retryable? error
           raise if attempts >= @connection.max_transaction_retries

data/lib/active_record/connection_adapters/cockroachdb/type.rb

@@ -1,12 +1,10 @@
 module ActiveRecord
   module Type
     class << self
-      private
-
       # Return :postgresql instead of :cockroachdb for current_adapter_name so
       # we can continue using the ActiveRecord::Types defined in
       # PostgreSQLAdapter.
-      def current_adapter_name
+      def adapter_name_from(_model)
         :postgresql
       end
     end

data/lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -1,36 +1,54 @@
-require 'active_record/connection_adapters/postgresql_adapter'
+require "rgeo/active_record"
+
+require "active_record/connection_adapters/postgresql_adapter"
+require "active_record/connection_adapters/cockroachdb/column_methods"
 require "active_record/connection_adapters/cockroachdb/schema_statements"
 require "active_record/connection_adapters/cockroachdb/referential_integrity"
 require "active_record/connection_adapters/cockroachdb/transaction_manager"
-require "active_record/connection_adapters/cockroachdb/column"
 require "active_record/connection_adapters/cockroachdb/database_statements"
+require "active_record/connection_adapters/cockroachdb/table_definition"
 require "active_record/connection_adapters/cockroachdb/quoting"
 require "active_record/connection_adapters/cockroachdb/type"
 require "active_record/connection_adapters/cockroachdb/attribute_methods"
+require "active_record/connection_adapters/cockroachdb/column"
+require "active_record/connection_adapters/cockroachdb/spatial_column_info"
+require "active_record/connection_adapters/cockroachdb/setup"
+require "active_record/connection_adapters/cockroachdb/oid/type_map_initializer"
+require "active_record/connection_adapters/cockroachdb/oid/spatial"
+require "active_record/connection_adapters/cockroachdb/oid/interval"
+require "active_record/connection_adapters/cockroachdb/arel_tosql"
+
+# Run to ignore spatial tables that will break schemna dumper.
+# Defined in ./setup.rb
+ActiveRecord::ConnectionAdapters::CockroachDB.initial_setup
 
 module ActiveRecord
   module ConnectionHandling
     def cockroachdb_connection(config)
       # This is copied from the PostgreSQL adapter.
-      conn_params = config.symbolize_keys
-
-      conn_params.delete_if { |_, v| v.nil? }
+      conn_params = config.symbolize_keys.compact
 
       # Map ActiveRecords param names to PGs.
       conn_params[:user] = conn_params.delete(:username) if conn_params[:username]
       conn_params[:dbname] = conn_params.delete(:database) if conn_params[:database]
 
       # Forward only valid config params to PG::Connection.connect.
-      valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:sslmode, :application_name]
+      valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:requiressl]
       conn_params.slice!(*valid_conn_param_keys)
 
-      conn = PG.connect(conn_params)
-      ConnectionAdapters::CockroachDBAdapter.new(conn, logger, conn_params, config)
-    rescue ::PG::Error, ActiveRecord::ActiveRecordError  => error
-      if error.message.include?("does not exist")
+      ConnectionAdapters::CockroachDBAdapter.new(
+        ConnectionAdapters::CockroachDBAdapter.new_client(conn_params),
+        logger,
+        conn_params,
+        config
+      )
+    # This rescue flow appears in new_client, but it is needed here as well
+    # since Cockroach will sometimes not raise until a query is made.
+    rescue ActiveRecord::StatementInvalid => error
+      if conn_params && conn_params[:dbname] && error.cause.message.include?(conn_params[:dbname])
         raise ActiveRecord::NoDatabaseError
       else
-        raise
+        raise ActiveRecord::ConnectionNotEstablished, error.message
       end
     end
   end
@@ -42,11 +60,49 @@ module ActiveRecord
       ADAPTER_NAME = "CockroachDB".freeze
       DEFAULT_PRIMARY_KEY = "rowid"
 
+      SPATIAL_COLUMN_OPTIONS =
+        {
+          geography:           { geographic: true },
+          geometry:            {},
+          geometry_collection: {},
+          line_string:         {},
+          multi_line_string:   {},
+          multi_point:         {},
+          multi_polygon:       {},
+          spatial:             {},
+          st_point:            {},
+          st_polygon:          {},
+        }
+
+      # http://postgis.17.x6.nabble.com/Default-SRID-td5001115.html      
+      DEFAULT_SRID = 0
+
       include CockroachDB::SchemaStatements
       include CockroachDB::ReferentialIntegrity
       include CockroachDB::DatabaseStatements
       include CockroachDB::Quoting
 
+      def self.spatial_column_options(key)
+        SPATIAL_COLUMN_OPTIONS[key]
+      end
+
+      def postgis_lib_version
+        @postgis_lib_version ||= select_value("SELECT PostGIS_Lib_Version()")
+      end
+
+      def default_srid
+        DEFAULT_SRID
+      end
+
+      def srs_database_columns
+        {
+          auth_name_column: "auth_name",
+          auth_srid_column: "auth_srid",
+          proj4text_column: "proj4text",
+          srtext_column:    "srtext",
+        }
+      end
+
       def debugging?
         !!ENV["DEBUG_COCKROACHDB_ADAPTER"]
       end
@@ -117,6 +173,10 @@ module ActiveRecord
         @crdb_version >= 202
       end
 
+      def supports_partitioned_indexes?
+        false
+      end
+
       # This is hardcoded to 63 (as previously was in ActiveRecord 5.0) to aid in
       # migration from PostgreSQL to CockroachDB. In practice, this limitation
       # is arbitrary since CockroachDB supports index name lengths and table alias
@@ -160,17 +220,46 @@ module ActiveRecord
       private
 
         def initialize_type_map(m = type_map)
+          %w(
+            geography
+            geometry
+            geometry_collection
+            line_string
+            multi_line_string
+            multi_point
+            multi_polygon
+            st_point
+            st_polygon
+          ).each do |geo_type|
+            m.register_type(geo_type) do |oid, _, sql_type|
+              CockroachDB::OID::Spatial.new(oid, sql_type)
+            end
+          end
+
+          # Belongs after other types are defined because of issues described
+          # in this https://github.com/rails/rails/pull/38571
+          # Once that PR is merged, we can call super at the top.
           super(m)
-          # NOTE(joey): PostgreSQL intervals have a precision.
-          # CockroachDB intervals do not, so overide the type
-          # definition. Returning a ArgumentError may not be correct.
-          # This needs to be tested.
-          m.register_type "interval" do |_, _, sql_type|
+
+          # Override numeric type. This is almost identical to the default,
+          # except that the conditional based on the fmod is changed.
+          m.register_type "numeric" do |_, fmod, sql_type|
             precision = extract_precision(sql_type)
-            if precision
-              raise(ArgumentError, "CockroachDB does not support precision on intervals, but got precision: #{precision}")
+            scale = extract_scale(sql_type)
+
+            # TODO(#178) this should never use DecimalWithoutScale since scale
+            # is assumed to be 0 if it is not explicitly defined.
+            #
+            # If fmod is -1, that means that precision is defined but not
+            # scale, or neither is defined.
+            if fmod && fmod == -1
+              # Below comment is from ActiveRecord
+              # FIXME: Remove this class, and the second argument to
+              # lookups on PG
+              Type::DecimalWithoutScale.new(precision: precision)
+            else
+              OID::Decimal.new(precision: precision, scale: scale)
             end
-            OID::SpecializedString.new(:interval, precision: precision)
           end
         end
 
@@ -281,6 +370,175 @@ module ActiveRecord
           return "{}"
         end
 
+        # override
+        # This method makes a query to gather information about columns
+        # in a table. It returns an array of arrays (one for each col) and
+        # passes each to the SchemaStatements#new_column_from_field method
+        # as the field parameter. This data is then used to format the column
+        # objects for the model and sent to the OID for data casting.
+        #
+        # Sometimes there are differences between how data is formatted
+        # in Postgres and CockroachDB, so additional queries for certain types
+        # may be necessary to properly form the column definition.
+        #
+        # @see: https://github.com/rails/rails/blob/8695b028261bdd244e254993255c6641bdbc17a5/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb#L829
+        def column_definitions(table_name)
+          fields = super
+
+          # 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)
+          fields.map do |field|
+            dtype = field[1]
+            if re.match(dtype)
+              crdb_column_definition(field, table_name)
+            else
+              field
+            end
+          end
+        end
+
+        # Use the crdb_sql_type instead of the sql_type returned by
+        # column_definitions. This will include limit,
+        # precision, and scale information in the type.
+        # Ex. geometry -> geometry(point, 4326)
+        def crdb_column_definition(field, table_name)
+          col_name = field[0]
+          data_type = \
+          query(<<~SQL, "SCHEMA")
+            SELECT c.crdb_sql_type
+              FROM information_schema.columns c
+            WHERE c.table_name = #{quote(table_name)}
+              AND c.column_name = #{quote(col_name)}
+          SQL
+          field[1] = data_type[0][0].downcase
+          field
+        end
+
+        # override
+        # This method is used to determine if a
+        # FEATURE_NOT_SUPPORTED error from the PG gem should
+        # be an ActiveRecord::PreparedStatementCacheExpired
+        # error.
+        #
+        # ActiveRecord handles this by checking that the sql state matches the
+        # FEATURE_NOT_SUPPORTED code and that the source function
+        # is "RevalidateCachedQuery" since that is the only function
+        # in postgres that will create this error.
+        #
+        # That method will not work for CockroachDB because the error
+        # originates from the "runExecBuilder" function, so we need
+        # to modify the original to match the CockroachDB behavior.
+        def is_cached_plan_failure?(e)
+          pgerror = e.cause
+
+          pgerror.result.result_error_field(PG::PG_DIAG_SQLSTATE) == FEATURE_NOT_SUPPORTED &&
+            pgerror.result.result_error_field(PG::PG_DIAG_SOURCE_FUNCTION) == "runExecBuilder"
+        rescue
+          false
+        end
+
+        # override
+        # This method loads info about data types from the database to
+        # populate the TypeMap.
+        #
+        # Currently, querying from the pg_type catalog can be slow due to geo-partitioning
+        # so this modified query uses AS OF SYSTEM TIME '-10s' to read historical data.
+        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)
+            end
+          else
+            super
+          end
+        rescue ActiveRecord::StatementInvalid => e
+          raise e unless e.cause.is_a? PG::InvalidCatalogName
+          # use original if database is younger than 10s
+          super
+        end
+
+        # override
+        # This method maps data types to their proper decoder.
+        #
+        # Currently, querying from the pg_type catalog can be slow due to geo-partitioning
+        # so this modified query uses AS OF SYSTEM TIME '-10s' to read historical data.
+        def add_pg_decoders
+          if @config[:use_follower_reads_for_type_introspection]
+            @default_timezone = nil
+            @timestamp_decoder = nil
+
+            coders_by_name = {
+              "int2" => PG::TextDecoder::Integer,
+              "int4" => PG::TextDecoder::Integer,
+              "int8" => PG::TextDecoder::Integer,
+              "oid" => PG::TextDecoder::Integer,
+              "float4" => PG::TextDecoder::Float,
+              "float8" => PG::TextDecoder::Float,
+              "numeric" => PG::TextDecoder::Numeric,
+              "bool" => PG::TextDecoder::Boolean,
+              "timestamp" => PG::TextDecoder::TimestampUtc,
+              "timestamptz" => PG::TextDecoder::TimestampWithTimeZone,
+            }
+
+            known_coder_types = coders_by_name.keys.map { |n| quote(n) }
+            query = <<~SQL % known_coder_types.join(", ")
+              SELECT t.oid, t.typname
+              FROM pg_type as t AS OF SYSTEM TIME '-10s'
+              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
+            end
+
+            map = PG::TypeMapByOid.new
+            coders.each { |coder| map.add_coder(coder) }
+            @connection.type_map_for_results = map
+
+            @type_map_for_results = PG::TypeMapByOid.new
+            @type_map_for_results.default_type_map = map
+            @type_map_for_results.add_coder(PG::TextDecoder::Bytea.new(oid: 17, name: "bytea"))
+            @type_map_for_results.add_coder(MoneyDecoder.new(oid: 790, name: "money"))
+
+            # extract timestamp decoder for use in update_typemap_for_default_timezone
+            @timestamp_decoder = coders.find { |coder| coder.name == "timestamp" }
+            update_typemap_for_default_timezone
+          else
+            super
+          end
+        rescue ActiveRecord::StatementInvalid => e
+          raise e unless e.cause.is_a? PG::InvalidCatalogName
+          # use original if database is younger than 10s
+          super
+        end
+
+        def arel_visitor
+          Arel::Visitors::CockroachDB.new(self)
+        end
+
       # end private
     end
   end