Add support for spatial features.

Author: keithdoggett

activerecord-cockroachdb-adapter.gemspec

@@ -15,6 +15,7 @@ Gem::Specification.new do |spec|
 
   spec.add_dependency "activerecord", "~> 6.0.3"
   spec.add_dependency "pg", ">= 0.20"
+  spec.add_dependency "rgeo-activerecord", "~> 7.0.0"
 
   # Prevent pushing this gem to RubyGems.org. To allow pushes either set the 'allowed_push_host'
   # to allow pushing to a single host or delete this section to allow pushing to any host.

lib/active_record/connection_adapters/cockroachdb/arel_tosql.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module RGeo
+  module ActiveRecord
+    ##
+    # Extend rgeo-activerecord visitors to use PostGIS specific functionality
+    module SpatialToPostGISSql
+      def visit_in_spatial_context(node, collector)
+        # Use ST_GeomFromEWKT for EWKT geometries
+        if node.is_a?(String) && node =~ /SRID=[\d+]{0,};/
+          collector << "#{st_func('ST_GeomFromEWKT')}(#{quote(node)})"
+        else
+          super(node, collector)
+        end
+      end
+    end
+  end
+end
+RGeo::ActiveRecord::SpatialToSql.prepend RGeo::ActiveRecord::SpatialToPostGISSql
+
+module Arel # :nodoc:
+  module Visitors # :nodoc:
+    class CockroachDB < PostgreSQL  # :nodoc:
+      include RGeo::ActiveRecord::SpatialToSql
+    end
+  end
+end

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

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

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

lib/active_record/connection_adapters/cockroachdb/quoting.rb

@@ -14,10 +14,18 @@ module Quoting
         # 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