activerecord-cockroachdb-adapter 6.0.0beta1 → 6.1.0.pre.beta.2

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

@@ -2,8 +2,85 @@ module ActiveRecord
   module ConnectionAdapters
     module CockroachDB
       module PostgreSQLColumnMonkeyPatch
+        # most functions taken from activerecord-postgis-adapter spatial_column
+        # https://github.com/rgeo/activerecord-postgis-adapter/blob/master/lib/active_record/connection_adapters/postgis/spatial_column.rb
+        def initialize(name, default, sql_type_metadata = nil, null = true,
+                       default_function = nil, collation: nil, comment: nil,
+                       serial: nil, spatial: nil)
+          @sql_type_metadata = sql_type_metadata
+          @geographic = !!(sql_type_metadata.sql_type =~ /geography\(/i)
+
+          if spatial
+            # This case comes from an entry in the geometry_columns table
+            set_geometric_type_from_name(spatial[:type])
+            @srid = spatial[:srid].to_i
+            @has_z = !!spatial[:has_z]
+            @has_m = !!spatial[:has_m]
+          elsif @geographic
+            # Geographic type information is embedded in the SQL type
+            @srid = 4326
+            @has_z = @has_m = false
+            build_from_sql_type(sql_type_metadata.sql_type)
+          elsif sql_type =~ /geography|geometry|point|linestring|polygon/i
+            build_from_sql_type(sql_type_metadata.sql_type)
+          elsif sql_type_metadata.sql_type =~ /geography|geometry|point|linestring|polygon/i
+            # A geometry column with no geometry_columns entry.
+            # @geometric_type = geo_type_from_sql_type(sql_type)
+            build_from_sql_type(sql_type_metadata.sql_type)
+          end
+          super(name, default, sql_type_metadata, null, default_function,
+            collation: collation, comment: comment, serial: serial)
+          if spatial? && @srid
+            @limit = { srid: @srid, type: to_type_name(geometric_type) }
+            @limit[:has_z] = true if @has_z
+            @limit[:has_m] = true if @has_m
+            @limit[:geographic] = true if @geographic
+          end
+        end
+
+        attr_reader :geographic,
+                    :geometric_type,
+                    :has_m,
+                    :has_z,
+                    :srid
+
+        alias geographic? geographic
+        alias has_z? has_z
+        alias has_m? has_m
+
+        def limit
+          spatial? ? @limit : super
+        end
+
+        def spatial?
+          %i[geometry geography].include?(@sql_type_metadata.type)
+        end
+
         def serial?
-          default_function == "unique_rowid()"
+          default_function == 'unique_rowid()'
+        end
+
+        private
+
+        def set_geometric_type_from_name(name)
+          @geometric_type = RGeo::ActiveRecord.geometric_type_from_name(name) || RGeo::Feature::Geometry
+        end
+
+        def build_from_sql_type(sql_type)
+          geo_type, @srid, @has_z, @has_m = OID::Spatial.parse_sql_type(sql_type)
+          set_geometric_type_from_name(geo_type)
+        end
+
+        def to_type_name(geometric_type)
+          name = geometric_type.type_name.underscore
+          case name
+          when 'point'
+            'st_point'
+          when 'polygon'
+            'st_polygon'
+          else
+            name
+          end
         end
       end
     end

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

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module CockroachDB
+      module ColumnMethods
+        def spatial(name, options = {})
+          raise "You must set a type. For example: 't.spatial type: :st_point'" unless options[:type]
+
+          column(name, options[:type], **options)
+        end
+
+        def geography(name, options = {})
+          column(name, :geography, **options)
+        end
+
+        def geometry(name, options = {})
+          column(name, :geometry, **options)
+        end
+
+        def geometry_collection(name, options = {})
+          column(name, :geometry_collection, **options)
+        end
+
+        def line_string(name, options = {})
+          column(name, :line_string, **options)
+        end
+
+        def multi_line_string(name, options = {})
+          column(name, :multi_line_string, **options)
+        end
+
+        def multi_point(name, options = {})
+          column(name, :multi_point, **options)
+        end
+
+        def multi_polygon(name, options = {})
+          column(name, :multi_polygon, **options)
+        end
+
+        def st_point(name, options = {})
+          column(name, :st_point, **options)
+        end
+
+        def st_polygon(name, options = {})
+          column(name, :st_polygon, **options)
+        end
+      end
+    end
+
+    PostgreSQL::Table.include CockroachDB::ColumnMethods
+  end
+end

data/lib/active_record/connection_adapters/cockroachdb/oid/interval.rb

@@ -0,0 +1,126 @@
+# frozen-string-literal: true
+
+require "active_support/duration"
+
+module ActiveRecord
+  module ConnectionAdapters
+    module CockroachDB
+      module OID
+        module Interval # :nodoc:
+          DEFAULT_PRECISION = 6 # microseconds
+
+          def cast_value(value)
+            case value
+            when ::ActiveSupport::Duration
+              value
+            when ::String
+              begin
+                PostgresqlInterval::Parser.parse(value)
+              rescue PostgresqlInterval::ParseError
+                # Try ISO 8601
+                super
+              end
+            else
+              super
+            end
+          end
+
+          def serialize(value)
+            precision = self.precision || DEFAULT_PRECISION
+            case value
+            when ::ActiveSupport::Duration
+              serialize_duration(value, precision)
+            when ::Numeric
+              serialize_duration(value.seconds, precision)
+            else
+              super
+            end
+          end
+
+          def type_cast_for_schema(value)
+            serialize(value).inspect
+          end
+
+          private
+
+            # Convert an ActiveSupport::Duration to
+            # the postgres interval style
+            # ex. 1 year 2 mons 3 days 4 hours 5 minutes 6 seconds
+            def serialize_duration(value, precision)
+              yrs = value.parts.fetch(:years, 0)
+              mons = value.parts.fetch(:months, 0)
+              days = value.parts.fetch(:days, 0)
+              hrs = value.parts.fetch(:hours, 0)
+              mins = value.parts.fetch(:minutes, 0)
+              secs = value.parts.fetch(:seconds, 0).round(precision)
+
+              "#{yrs} years #{mons} mons #{days} days #{hrs} hours #{mins} minutes #{secs} seconds"
+            end
+        end
+
+        PostgreSQL::OID::Interval.prepend(Interval)
+      end
+
+      module PostgresqlInterval
+        class Parser
+          PARTS = ActiveSupport::Duration::PARTS
+          PARTS_IN_SECONDS = ActiveSupport::Duration::PARTS_IN_SECONDS
+
+          # modified regex from https://github.com/jeremyevans/sequel/blob/master/lib/sequel/extensions/pg_interval.rb#L86
+          REGEX = /\A([+-]?\d+ years?\s?)?([+-]?\d+ mons?\s?)?([+-]?\d+ days?\s?)?(?:([+-])?(\d{2,10}):(\d\d):(\d\d(\.\d+)?))?\z/
+
+          def self.parse(string)
+            matches = REGEX.match(string)
+            raise(ParseError) unless matches
+
+            # 1 => years, 2 => months, 3 => days, 4 => nil, 5 => hours,
+            # 6 => minutes, 7 => seconds with fraction digits, 8 => fractional portion of 7
+            duration = 0
+            parts = {}
+
+            if matches[1]
+              val = matches[1].to_i
+              duration += val * PARTS_IN_SECONDS[:years]
+              parts[:years] = val
+            end
+
+            if matches[2]
+              val = matches[2].to_i
+              duration += val * PARTS_IN_SECONDS[:months]
+              parts[:months] = val
+            end
+
+            if matches[3]
+              val = matches[3].to_i
+              duration += val * PARTS_IN_SECONDS[:days]
+              parts[:days] = val
+            end
+
+            if matches[5]
+              val = matches[5].to_i
+              duration += val * PARTS_IN_SECONDS[:hours]
+              parts[:hours] = val
+            end
+
+            if matches[6]
+              val = matches[6].to_i
+              duration += val * PARTS_IN_SECONDS[:minutes]
+              parts[:minutes] = val
+            end
+
+            if matches[7]
+              val = matches[7].to_f
+              duration += val * PARTS_IN_SECONDS[:seconds]
+              parts[:seconds] = val
+            end
+
+            ActiveSupport::Duration.new(duration, parts)
+          end
+        end
+
+        class ParseError < StandardError
+        end
+      end
+    end
+  end
+end

data/lib/active_record/connection_adapters/cockroachdb/oid/spatial.rb

@@ -0,0 +1,121 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+  module ConnectionAdapters
+    module CockroachDB
+      module OID
+        class Spatial < Type::Value
+          # sql_type is a string that comes from the database definition
+          # examples:
+          #   "geometry(Point,4326)"
+          #   "geography(Point,4326)"
+          #   "geometry(Polygon,4326) NOT NULL"
+          #   "geometry(Geography,4326)"
+          def initialize(oid, sql_type)
+            @sql_type = sql_type
+            @geo_type, @srid, @has_z, @has_m = self.class.parse_sql_type(sql_type)
+          end
+
+          # sql_type: geometry, geometry(Point), geometry(Point,4326), ...
+          #
+          # returns [geo_type, srid, has_z, has_m]
+          #   geo_type: geography, geometry, point, line_string, polygon, ...
+          #   srid:     1234
+          #   has_z:    false
+          #   has_m:    false
+          def self.parse_sql_type(sql_type)
+            geo_type = nil
+            srid = 0
+            has_z = false
+            has_m = false
+
+            if sql_type =~ /(geography|geometry)\((.*)\)$/i
+              # geometry(Point)
+              # geometry(Point,4326)
+              params = Regexp.last_match(2).split(',')
+              if params.first =~ /([a-z]+[^zm])(z?)(m?)/i
+                has_z = Regexp.last_match(2).length > 0
+                has_m = Regexp.last_match(3).length > 0
+                geo_type = Regexp.last_match(1)
+              end
+              srid = Regexp.last_match(1).to_i if params.last =~ /(\d+)/
+            else
+              geo_type = sql_type
+            end
+            [geo_type, srid, has_z, has_m]
+          end
+
+          def spatial_factory
+            @spatial_factory ||=
+              RGeo::ActiveRecord::SpatialFactoryStore.instance.factory(
+                factory_attrs
+              )
+          end
+
+          def geographic?
+            @sql_type =~ /geography/
+          end
+
+          def spatial?
+            true
+          end
+
+          def type
+            geographic? ? :geography : :geometry
+          end
+
+          # support setting an RGeo object or a WKT string
+          def serialize(value)
+            return if value.nil?
+
+            geo_value = cast_value(value)
+
+            # TODO: - only valid types should be allowed
+            # e.g. linestring is not valid for point column
+            # raise "maybe should raise" unless RGeo::Feature::Geometry.check_type(geo_value)
+
+            RGeo::WKRep::WKBGenerator.new(hex_format: true, type_format: :ewkb, emit_ewkb_srid: true)
+                                     .generate(geo_value)
+          end
+
+          private
+
+          def cast_value(value)
+            return if value.nil?
+
+            value.is_a?(String) ? parse_wkt(value) : value
+          end
+
+          # convert WKT string into RGeo object
+          def parse_wkt(string)
+            wkt_parser(string).parse(string)
+          rescue RGeo::Error::ParseError
+            nil
+          end
+
+          def binary_string?(string)
+            string[0] == "\x00" || string[0] == "\x01" || string[0, 4] =~ /[0-9a-fA-F]{4}/
+          end
+
+          def wkt_parser(string)
+            if binary_string?(string)
+              RGeo::WKRep::WKBParser.new(spatial_factory, support_ewkb: true, default_srid: @srid)
+            else
+              RGeo::WKRep::WKTParser.new(spatial_factory, support_ewkt: true, default_srid: @srid)
+            end
+          end
+
+          def factory_attrs
+            {
+              geo_type: @geo_type.underscore,
+              has_m: @has_m,
+              has_z: @has_z,
+              srid: @srid,
+              sql_type: type.to_s
+            }
+          end
+        end
+      end
+    end
+  end
+end

data/lib/active_record/connection_adapters/cockroachdb/oid/type_map_initializer.rb

@@ -0,0 +1,26 @@
+module ActiveRecord
+  module ConnectionAdapters
+    module CockroachDB
+      module OID
+        module TypeMapInitializer
+          # override
+          # Replaces the query with a faster version that doesn't rely on the
+          # use of 'array_in(cstring,oid,integer)'::regprocedure.
+          def query_conditions_for_initial_load
+            known_type_names = @store.keys.map { |n| "'#{n}'" }
+            known_type_types = %w('r' 'e' 'd')
+            <<~SQL % [known_type_names.join(", "), known_type_types.join(", ")]
+              WHERE
+                t.typname IN (%s)
+                OR t.typtype IN (%s)
+                OR (t.typarray = 0 AND t.typcategory='A')
+                OR t.typelem != 0
+            SQL
+          end
+        end
+
+        PostgreSQL::OID::TypeMapInitializer.prepend(TypeMapInitializer)
+      end
+    end
+  end
+end

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

@@ -14,10 +14,18 @@ module ActiveRecord
         # always be strings. Then, we won't have to make any additional changes
         # to ActiveRecord to support inserting integer values into string
         # columns.
+        #
+        # For spatial types, data is stored as Well-known Binary (WKB) strings
+        # (https://en.wikipedia.org/wiki/Well-known_text_representation_of_geometry#Well-known_binary)
+        # but when creating objects, using RGeo features is more convenient than
+        # converting to WKB, so this does it automatically.
         def _quote(value)
-          case value
-          when Numeric
+          if value.is_a?(Numeric)
             "'#{quote_string(value.to_s)}'"
+          elsif RGeo::Feature::Geometry.check_type(value)
+            "'#{RGeo::WKRep::WKBGenerator.new(hex_format: true, type_format: :ewkb, emit_ewkb_srid: true).generate(value)}'"
+          elsif value.is_a?(RGeo::Cartesian::BoundingBox)
+            "'#{value.min_x},#{value.min_y},#{value.max_x},#{value.max_y}'::box"
           else
             super
           end

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

@@ -4,7 +4,7 @@ module ActiveRecord
       module SchemaStatements
         include ActiveRecord::ConnectionAdapters::PostgreSQL::SchemaStatements
 
-        def add_index(table_name, column_name, options = {})
+        def add_index(table_name, column_name, **options)
           super
         rescue ActiveRecord::StatementInvalid => error
           if debugging? && error.cause.class == PG::FeatureNotSupported
@@ -39,8 +39,43 @@ module ActiveRecord
           nil
         end
 
+        # override
+        # https://github.com/rails/rails/blob/6-0-stable/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb#L624
+        def new_column_from_field(table_name, field)
+          column_name, type, default, notnull, oid, fmod, collation, comment = field
+          type_metadata = fetch_type_metadata(column_name, type, oid.to_i, fmod.to_i)
+          default_value = extract_value_from_default(default)
+          default_function = extract_default_function(default_value, default)
+
+          serial =
+            if (match = default_function&.match(/\Anextval\('"?(?<sequence_name>.+_(?<suffix>seq\d*))"?'::regclass\)\z/))
+              sequence_name_from_parts(table_name, column_name, match[:suffix]) == match[:sequence_name]
+            end
+
+          # {:dimension=>2, :has_m=>false, :has_z=>false, :name=>"latlon", :srid=>0, :type=>"GEOMETRY"}
+          spatial = spatial_column_info(table_name).get(column_name, type_metadata.sql_type)
+
+          PostgreSQL::Column.new(
+            column_name,
+            default_value,
+            type_metadata,
+            !notnull,
+            default_function,
+            collation: collation,
+            comment: comment.presence,
+            serial: serial,
+            spatial: spatial
+          )
+        end
+
         # CockroachDB will use INT8 if the SQL type is INTEGER, so we make it use
         # INT4 explicitly when needed.
+        #
+        # For spatial columns, include the limit to properly format the column name
+        # since type alone is not enough to format the column.
+        # Ex. type_to_sql(:geography, limit: "Point,4326")
+        # => "geography(Point,4326)"
+        # 
         def type_to_sql(type, limit: nil, precision: nil, scale: nil, array: nil, **) # :nodoc:
           sql = \
             case type.to_s
@@ -52,6 +87,8 @@ module ActiveRecord
               when 5..8; "int8"
               else super
               end
+            when "geometry", "geography"
+              "#{type}(#{limit})"
             else
               super
             end
@@ -86,6 +123,34 @@ module ActiveRecord
             query_value("SELECT setval(#{quote(quoted_sequence)}, #{max_pk ? max_pk : minvalue}, #{max_pk ? true : false})", "SCHEMA")
           end
         end
+
+        # override
+        def native_database_types
+          # Add spatial types
+          super.merge(
+            geography:           { name: "geography" },
+            geometry:            { name: "geometry" },
+            geometry_collection: { name: "geometry_collection" },
+            line_string:         { name: "line_string" },
+            multi_line_string:   { name: "multi_line_string" },
+            multi_point:         { name: "multi_point" },
+            multi_polygon:       { name: "multi_polygon" },
+            spatial:             { name: "geometry" },
+            st_point:            { name: "st_point" },
+            st_polygon:          { name: "st_polygon" }
+          )
+        end
+
+        # override
+        def create_table_definition(*args, **kwargs)
+          CockroachDB::TableDefinition.new(self, *args, **kwargs)
+        end
+
+        # memoize hash of column infos for tables
+        def spatial_column_info(table_name)
+          @spatial_column_info ||= {}
+          @spatial_column_info[table_name.to_sym] ||= SpatialColumnInfo.new(self, table_name.to_s)
+        end
       end
     end
   end