activerecord-cockroachdb-adapter 5.2.1 → 6.1.0beta1

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/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,15 +87,70 @@ module ActiveRecord
               when 5..8; "int8"
               else super
               end
+            when "geometry", "geography"
+              "#{type}(#{limit})"
             else
               super
             end
           # The call to super might have appeneded [] already.
           if array && type != :primary_key && !sql.end_with?("[]")
-            sql = "#{sql}[]" 
+            sql = "#{sql}[]"
           end
           sql
         end
+
+        # This overrides the method from PostegreSQL adapter
+        # Resets the sequence of a table's primary key to the maximum value.
+        def reset_pk_sequence!(table, pk = nil, sequence = nil)
+          unless pk && sequence
+            default_pk, default_sequence = pk_and_sequence_for(table)
+
+            pk ||= default_pk
+            sequence ||= default_sequence
+          end
+
+          if @logger && pk && !sequence
+            @logger.warn "#{table} has primary key #{pk} with no default sequence."
+          end
+
+          if pk && sequence
+            quoted_sequence = quote_table_name(sequence)
+            max_pk = query_value("SELECT MAX(#{quote_column_name pk}) FROM #{quote_table_name(table)}", "SCHEMA")
+            if max_pk.nil?
+              minvalue = query_value("SELECT seqmin FROM pg_sequence WHERE seqrelid = #{quote(quoted_sequence)}::regclass", "SCHEMA")
+            end
+
+            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