activerecord-cockroachdb-adapter 6.0.0beta1 → 6.0.0

data/lib/active_record/connection_adapters/cockroachdb_adapter.rb

@@ -1,12 +1,25 @@
+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/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
@@ -24,12 +37,9 @@ module ActiveRecord
       valid_conn_param_keys = PG::Connection.conndefaults_hash.keys + [:sslmode, :application_name]
       conn_params.slice!(*valid_conn_param_keys)
 
-      # The postgres drivers don't allow the creation of an unconnected
-      # PG::Connection object, so just pass a nil connection object for the
-      # time being.
       conn = PG.connect(conn_params)
       ConnectionAdapters::CockroachDBAdapter.new(conn, logger, conn_params, config)
-    rescue ::PG::Error => error
+    rescue ::PG::Error, ActiveRecord::ActiveRecordError  => error
       if error.message.include?("does not exist")
         raise ActiveRecord::NoDatabaseError
       else
@@ -41,15 +51,127 @@ end
 
 module ActiveRecord
   module ConnectionAdapters
+    module CockroachDBConnectionPool
+      def initialize(spec)
+        super(spec)
+        disable_telemetry = spec.config[:disable_cockroachdb_telemetry]
+        adapter = spec.config[:adapter]
+        return if disable_telemetry || adapter != "cockroachdb"
+
+        with_connection do |conn|
+          if conn.active?
+            begin
+              query = "SELECT crdb_internal.increment_feature_counter('ActiveRecord %d.%d')"
+              conn.execute(query % [ActiveRecord::VERSION::MAJOR, ActiveRecord::VERSION::MINOR])
+            rescue ActiveRecord::StatementInvalid
+              # The increment_feature_counter built-in is not supported on this
+              # CockroachDB version. Ignore.
+            rescue StandardError => e
+              conn.logger.warn "Unexpected error when incrementing feature counter: #{e}"
+            end
+          end
+        end
+      end
+    end
+    ConnectionPool.prepend(CockroachDBConnectionPool)
+
     class CockroachDBAdapter < PostgreSQLAdapter
       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
 
+      # override
+      # This method makes a sql 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.
+      #
+      # The issue with the default method is that the sql_type field is 
+      # retrieved with the `format_type` function, but this is implemented
+      # differently in CockroachDB than PostGIS, so geometry/geography
+      # types are missing information which makes parsing them impossible.
+      # Below is an example of what `format_type` returns for a geometry
+      # column.
+      #
+      # column_type: geometry(POINT, 4326)
+      # Expected: geometry(POINT, 4326)
+      # Actual: geometry
+      #
+      # The solution is to make the default query with super, then
+      # iterate through the columns and if it is a spatial type,
+      # access the proper column_type with the information_schema.columns 
+      # table.
+      #
+      # @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
+        # iterate through and identify all spatial fields based on format_type
+        # being geometry or geography, then query for the information_schema.column
+        # column_type because that contains the necessary information.
+        fields.map do |field|
+          dtype = field[1]
+          if dtype == 'geometry' || dtype == 'geography'
+            col_name = field[0]
+            data_type = \
+            query(<<~SQL, "SCHEMA")
+              SELECT c.data_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]
+          end
+          field
+        end
+      end
+
+      def arel_visitor
+        Arel::Visitors::CockroachDB.new(self)
+      end
+
+      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
@@ -120,6 +242,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
@@ -137,6 +263,7 @@ module ActiveRecord
 
       def initialize(connection, logger, conn_params, config)
         super(connection, logger, conn_params, config)
+
         crdb_version_string = query_value("SHOW crdb_version")
         if crdb_version_string.include? "v1."
           version_num = 1
@@ -154,10 +281,31 @@ module ActiveRecord
         @crdb_version = version_num
       end
 
+      def self.database_exists?(config)
+        !!ActiveRecord::Base.cockroachdb_connection(config)
+      rescue ActiveRecord::NoDatabaseError
+        false
+      end
+
       private
 
         def initialize_type_map(m = type_map)
-          super(m)
+          %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
+
           # NOTE(joey): PostgreSQL intervals have a precision.
           # CockroachDB intervals do not, so overide the type
           # definition. Returning a ArgumentError may not be correct.
@@ -169,6 +317,8 @@ module ActiveRecord
             end
             OID::SpecializedString.new(:interval, precision: precision)
           end
+
+          super(m)
         end
 
         # Configures the encoding, verbosity, schema search path, and time zone of the connection.
@@ -278,6 +428,100 @@ module ActiveRecord
           return "{}"
         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"))
+
+            # 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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: activerecord-cockroachdb-adapter
 version: !ruby/object:Gem::Version
-  version: 6.0.0beta1
+  version: 6.0.0
 platform: ruby
 authors:
 - Cockroach Labs
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2020-12-17 00:00:00.000000000 Z
+date: 2021-04-28 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -38,6 +38,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0.20'
+- !ruby/object:Gem::Dependency
+  name: rgeo-activerecord
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 7.0.0
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: 7.0.0
 description: Allows the use of CockroachDB as a backend for ActiveRecord and Rails
   apps.
 email:
@@ -48,6 +62,7 @@ extra_rdoc_files: []
 files:
 - ".gitignore"
 - ".gitmodules"
+- CHANGELOG.md
 - CONTRIBUTING.md
 - Gemfile
 - LICENSE
@@ -61,13 +76,20 @@ files:
 - build/local-test.sh
 - build/teamcity-test.sh
 - docker.sh
+- lib/active_record/connection_adapters/cockroachdb/arel_tosql.rb
 - lib/active_record/connection_adapters/cockroachdb/attribute_methods.rb
 - lib/active_record/connection_adapters/cockroachdb/column.rb
+- lib/active_record/connection_adapters/cockroachdb/column_methods.rb
 - lib/active_record/connection_adapters/cockroachdb/database_statements.rb
 - lib/active_record/connection_adapters/cockroachdb/database_tasks.rb
+- lib/active_record/connection_adapters/cockroachdb/oid/spatial.rb
+- lib/active_record/connection_adapters/cockroachdb/oid/type_map_initializer.rb
 - lib/active_record/connection_adapters/cockroachdb/quoting.rb
 - lib/active_record/connection_adapters/cockroachdb/referential_integrity.rb
 - lib/active_record/connection_adapters/cockroachdb/schema_statements.rb
+- lib/active_record/connection_adapters/cockroachdb/setup.rb
+- lib/active_record/connection_adapters/cockroachdb/spatial_column_info.rb
+- lib/active_record/connection_adapters/cockroachdb/table_definition.rb
 - lib/active_record/connection_adapters/cockroachdb/transaction_manager.rb
 - lib/active_record/connection_adapters/cockroachdb/type.rb
 - lib/active_record/connection_adapters/cockroachdb_adapter.rb
@@ -88,9 +110,9 @@ required_ruby_version: !ruby/object:Gem::Requirement
       version: '0'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
-  - - ">"
+  - - ">="
     - !ruby/object:Gem::Version
-      version: 1.3.1
+      version: '0'
 requirements: []
 rubygems_version: 3.1.4
 signing_key: