fix(schema_statements): add unique_rowid to pk_and_sequence_for

Allow to get the primary key in the default case.

Author: BuonOmo

lib/active_record/connection_adapters/cockroachdb/schema_statements.rb

@@ -55,6 +55,63 @@ def primary_key(table_name)
           end
         end
 
+        # OVERRIDE: Added `unique_rowid` to the last line of the second query.
+        #   This is a CockroachDB-specific function used for primary keys.
+        #
+        # Returns a table's primary key and belonging sequence.
+        def pk_and_sequence_for(table) # :nodoc:
+          # First try looking for a sequence with a dependency on the
+          # given table's primary key.
+          result = query(<<~SQL, "SCHEMA")[0]
+            SELECT attr.attname, nsp.nspname, seq.relname
+            FROM pg_class      seq,
+                 pg_attribute  attr,
+                 pg_depend     dep,
+                 pg_constraint cons,
+                 pg_namespace  nsp
+            WHERE seq.oid           = dep.objid
+              AND seq.relkind       = 'S'
+              AND attr.attrelid     = dep.refobjid
+              AND attr.attnum       = dep.refobjsubid
+              AND attr.attrelid     = cons.conrelid
+              AND attr.attnum       = cons.conkey[1]
+              AND seq.relnamespace  = nsp.oid
+              AND cons.contype      = 'p'
+              AND dep.classid       = 'pg_class'::regclass
+              AND dep.refobjid      = #{quote(quote_table_name(table))}::regclass
+          SQL
+
+          if result.nil? || result.empty?
+            result = query(<<~SQL, "SCHEMA")[0]
+              SELECT attr.attname, nsp.nspname,
+                CASE
+                  WHEN pg_get_expr(def.adbin, def.adrelid) !~* 'nextval' THEN NULL
+                  WHEN split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2) ~ '.' THEN
+                    substr(split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2),
+                           strpos(split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2), '.')+1)
+                  ELSE split_part(pg_get_expr(def.adbin, def.adrelid), '''', 2)
+                END
+              FROM pg_class       t
+              JOIN pg_attribute   attr ON (t.oid = attrelid)
+              JOIN pg_attrdef     def  ON (adrelid = attrelid AND adnum = attnum)
+              JOIN pg_constraint  cons ON (conrelid = adrelid AND adnum = conkey[1])
+              JOIN pg_namespace   nsp  ON (t.relnamespace = nsp.oid)
+              WHERE t.oid = #{quote(quote_table_name(table))}::regclass
+                AND cons.contype = 'p'
+                AND pg_get_expr(def.adbin, def.adrelid) ~* 'nextval|uuid_generate|gen_random_uuid|unique_rowid'
+            SQL
+          end
+
+          pk = result.shift
+          if result.last
+            [pk, PostgreSQL::Name.new(*result)]
+          else
+            [pk, nil]
+          end
+        rescue
+          nil
+        end
+
         # override
         # Modified version of the postgresql foreign_keys method.
         # Replaces t2.oid::regclass::text with t2.relname since this is

test/cases/adapters/postgresql/postgresql_adapter_test.rb

@@ -2,6 +2,7 @@
 require "cases/helper"
 require "support/ddl_helper"
 require "support/connection_helper"
+require "support/copy_cat"
 
 module CockroachDB
   module ConnectionAdapters
@@ -12,7 +13,6 @@ class PostgreSQLAdapterTest < ActiveRecord::PostgreSQLTestCase
 
       def setup
         @connection = ActiveRecord::Base.connection
-        @connection_handler = ActiveRecord::Base.connection_handler
       end
 
       def teardown
@@ -87,11 +87,27 @@ def test_invalid_index
         end
       end
 
+      # OVERRIDE: the `default_sequence_name` is `nil`, let's `to_s` it
+      #   for a fair comparison.
+      CopyCat.copy_methods(self, ActiveRecord::ConnectionAdapters::PostgreSQLAdapterTest,
+        :test_pk_and_sequence_for,
+        :test_pk_and_sequence_for_with_non_standard_primary_key
+      ) do
+        attr_accessor :already_updated
+        def on_send(node)
+          return super unless node in [:send, _, :default_sequence_name, [:str, "ex"], [:str, "id"|"code"]]
+
+          raise "The source code must have changed" if already_updated
+          already_updated ||= :yes
+          insert_after(node.loc.expression, ".to_s")
+        end
+      end
+
       private
 
-      def with_example_table(definition = "id serial primary key, number integer, data character varying(255)", &block)
-        super(@connection, "ex", definition, &block)
-      end
+      CopyCat.copy_methods(self, ActiveRecord::ConnectionAdapters::PostgreSQLAdapterTest,
+        :with_example_table
+      )
     end
   end
 end

test/cases/migration_test.rb

@@ -1,5 +1,6 @@
 require "cases/helper_cockroachdb"
 require "models/person"
+require "support/copy_cat"
 
 class Reminder < ActiveRecord::Base; end unless Object.const_defined?(:Reminder)
 class Thing < ActiveRecord::Base; end unless Object.const_defined?(:Thing)
@@ -92,7 +93,6 @@ def migrate(x)
 
   class BulkAlterTableMigrationsTest < ActiveRecord::TestCase
     def setup
-      @original_test = ::BulkAlterTableMigrationsTest.new(name)
       @connection = Person.connection
       @connection.create_table(:delete_me, force: true) { |t| }
       Person.reset_column_information
@@ -103,38 +103,28 @@ def setup
       Person.connection.drop_table(:delete_me) rescue nil
     end
 
-    %i(
-      test_adding_indexes
-      test_removing_index
-      test_adding_multiple_columns
-      test_changing_index
-    ).each do |method_name|
-      file, line = ::BulkAlterTableMigrationsTest.instance_method(method_name).source_location
-      iter = File.foreach(file)
-      (line - 1).times { iter.next }
-      indent = iter.next[/\A\s*/]
-      content = +""
-      content << iter.next until iter.peek == indent + "end\n"
-      content['"PostgreSQLAdapter" =>'] = '"CockroachDBAdapter" =>'
-      eval(<<-RUBY, binding, __FILE__, __LINE__ + 1)
-        def #{method_name}
-          #{content}
-        end
-      RUBY
+    # Change expected query count from PostgreSQLAdapter to CockroachDBAdapter.
+    CopyCat.copy_methods(self, ::BulkAlterTableMigrationsTest,
+      :test_adding_indexes,
+      :test_removing_index,
+      :test_adding_multiple_columns,
+      :test_changing_index
+    ) do
+      def on_str(node)
+        return unless node in [:str, "PostgreSQLAdapter"]
+
+        replace(node.loc.expression, '"CockroachDBAdapter"')
+      end
     end
 
     private
 
-    %i[
-      with_bulk_change_table
-      column
-      columns
-      index
-      indexes
-    ].each do |method|
-      define_method(method) do |*args, &block|
-        @original_test.send(method, *args, &block)
-      end
-    end
+    CopyCat.copy_methods(self, ::BulkAlterTableMigrationsTest,
+      :with_bulk_change_table,
+      :column,
+      :columns,
+      :index,
+      :indexes
+    )
   end
 end

test/excludes/ActiveRecord/ConnectionAdapters/PostgreSQLAdapterTest.rb

@@ -17,6 +17,8 @@
 exclude :test_partial_index_on_column_named_like_keyword, "CockroachDB adds parentheses around the WHERE definition."
 exclude :test_only_check_for_insensitive_comparison_capability_once, "the DROP DOMAIN syntax is not supported by CRDB"
 
+exclude :test_pk_and_sequence_for, "The test needs a little rework since the sequence is empty in CRDB"
+exclude :test_pk_and_sequence_for_with_non_standard_primary_key, "The test needs a little rework since the sequence is empty in CRDB"
 
 plpgsql_needed = "Will be testable with CRDB 23.2, introducing PL-PGSQL"
 exclude :test_ignores_warnings_when_behaviour_ignore, plpgsql_needed