fix: handle indexes with a hidden compound

This fixes regional related parts, hence fixes #344

Author: BuonOmo

bin/console

@@ -12,20 +12,22 @@ require "activerecord-cockroachdb-adapter"
 #       structure_load(Post.connection_db_config, "awesome-file.sql")
 require "active_record/connection_adapters/cockroachdb/database_tasks"
 
+DB_NAME = "ar_crdb_console"
+
 schema_kind = ENV.fetch("SCHEMA_KIND", ENV.fetch("SCHEMA", "default"))
 
-system("cockroach sql --insecure --host=localhost:26257 --execute='drop database if exists ar_crdb_console'",
+system("cockroach sql --insecure --host=localhost:26257 --execute='drop database if exists #{DB_NAME}'",
   exception: true)
-system("cockroach sql --insecure --host=localhost:26257 --execute='create database ar_crdb_console'",
+system("cockroach sql --insecure --host=localhost:26257 --execute='create database #{DB_NAME}'",
   exception: true)
 
 ActiveRecord::Base.establish_connection(
-  #Alternative version:  "cockroachdb://root@localhost:26257/ar_crdb_console"
+  #Alternative version:  "cockroachdb://root@localhost:26257/#{DB_NAME}"
   adapter: "cockroachdb",
   host: "localhost",
   port: 26257,
   user: "root",
-  database: "ar_crdb_console"
+  database: DB_NAME
 )
 
 load "#{__dir__}/console_schemas/#{schema_kind}.rb"

bin/console_schemas/default.rb

@@ -6,4 +6,6 @@ class Post < ActiveRecord::Base
     t.string :title
     t.text :body
   end
+
+  add_index("posts", ["title"], name: "index_posts_on_title", unique: true)
 end

lib/active_record/connection_adapters/cockroachdb/schema_statements.rb

@@ -55,6 +55,37 @@ def primary_key(table_name)
           end
         end
 
+        def primary_keys(table_name)
+          return super unless database_version >= 24_02_02
+
+          query_values(<<~SQL, "SCHEMA")
+            SELECT a.attname
+              FROM (
+                     SELECT indrelid, indkey, generate_subscripts(indkey, 1) idx
+                       FROM pg_index
+                      WHERE indrelid = #{quote(quote_table_name(table_name))}::regclass
+                        AND indisprimary
+                   ) i
+              JOIN pg_attribute a
+                ON a.attrelid = i.indrelid
+               AND a.attnum = i.indkey[i.idx]
+               AND NOT a.attishidden
+             ORDER BY i.idx
+          SQL
+        end
+
+        def column_names_from_column_numbers(table_oid, column_numbers)
+          return super unless database_version >= 24_02_02
+
+          Hash[query(<<~SQL, "SCHEMA")].values_at(*column_numbers).compact
+            SELECT a.attnum, a.attname
+            FROM pg_attribute a
+            WHERE a.attrelid = #{table_oid}
+            AND a.attnum IN (#{column_numbers.join(", ")})
+            AND NOT a.attishidden
+          SQL
+        end
+
         # OVERRIDE: CockroachDB does not support deferrable constraints.
         #   See: https://go.crdb.dev/issue-v/31632/v23.1
         def foreign_key_options(from_table, to_table, options)

test/cases/migration/hidden_column_test.rb

@@ -53,6 +53,17 @@ def test_add_hidden_column
           output = dump_table_schema "rockets"
           assert_match %r{t.uuid "new_col", hidden: true$}, output
         end
+
+        # Since 24.2.2, hash sharded indexes add a hidden column to the table.
+        # This tests ensure that the user can still drop the index even if they
+        # call `#remove_index` with the column name rather than the index name.
+        def test_remove_index_with_a_hidden_column
+          @connection.execute <<-SQL
+            CREATE INDEX hash_idx ON rockets (name) USING HASH WITH (bucket_count=8);
+          SQL
+          @connection.remove_index :rockets, :name
+          # assert :ok
+        end
       end
     end
   end

test/cases/primary_keys_test.rb

@@ -97,4 +97,38 @@ def test_schema_dump_primary_key_integer_with_default_nil
       assert_match %r{create_table "int_defaults", id: :bigint, default: nil}, schema
     end
   end
+
+  class PrimaryKeyHiddenColumnTest < ActiveRecord::TestCase
+    class Rocket < ActiveRecord::Base
+    end
+
+    def setup
+      connection = ActiveRecord::Base.lease_connection
+      connection.execute <<-SQL
+        CREATE TABLE rockets(
+          id SERIAL PRIMARY KEY USING HASH WITH (bucket_count=4)
+        )
+      SQL
+    end
+
+    def teardown
+      ActiveRecord::Base.connection.drop_table :rockets
+    end
+
+    def test_to_key_with_hidden_primary_key_part
+      rocket = Rocket.new
+      assert_nil rocket.to_key
+      rocket.save
+      assert_equal rocket.to_key, [rocket.id]
+    end
+
+    def test_read_attribute_with_hidden_primary_key_part
+      rocket = Rocket.create!
+      id = assert_not_deprecated(ActiveRecord.deprecator) do
+        rocket.read_attribute(:id)
+      end
+
+      assert_equal rocket.id, id
+    end
+  end
 end