Merge pull request #27961 from jordanlewis/simplify-postgres-primary-keys-v2

Simplify and optimize Postgres query for primary_keys()

3 files changed, 10 insertions, 24 deletions

diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
index eebc688686..863a0dd997 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
@@ -430,17 +430,17 @@ module ActiveRecord
         end
 
         def primary_keys(table_name) # :nodoc:
+          name = Utils.extract_schema_qualified_name(table_name.to_s)
           select_values(<<-SQL.strip_heredoc, "SCHEMA")
-            WITH pk_constraint AS (
-              SELECT conrelid, unnest(conkey) AS connum FROM pg_constraint
-              WHERE contype = 'p'
-                AND conrelid = #{quote(quote_table_name(table_name))}::regclass
-            ), cons AS (
-              SELECT conrelid, connum, row_number() OVER() AS rownum FROM pk_constraint
-            )
-            SELECT attr.attname FROM pg_attribute attr
-            INNER JOIN cons ON attr.attrelid = cons.conrelid AND attr.attnum = cons.connum
-            ORDER BY cons.rownum
+            SELECT column_name
+              FROM information_schema.key_column_usage kcu
+              JOIN information_schema.table_constraints tc
+                ON kcu.table_name = tc.table_name
+               AND kcu.table_schema = tc.table_schema
+               AND kcu.constraint_name = tc.constraint_name
+             WHERE constraint_type = 'PRIMARY KEY'
+               AND kcu.table_name = #{quote(name.identifier)}
+               AND kcu.table_schema = #{name.schema ? quote(name.schema) : "ANY (current_schemas(false))"}
           SQL
         end
 
diff --git a/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb b/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
index 4de04646ed..3054f0271f 100644
--- a/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/postgresql_adapter_test.rb
@@ -54,12 +54,6 @@ module ActiveRecord
         end
       end
 
-      def test_primary_key_raises_error_if_table_not_found
-        assert_raises(ActiveRecord::StatementInvalid) do
-          @connection.primary_key("unobtainium")
-        end
-      end
-
       def test_exec_insert_with_returning_disabled
         connection = connection_without_insert_returning
         result = connection.exec_insert("insert into postgresql_partitioned_table_parent (number) VALUES (1)", nil, [], "id", "postgresql_partitioned_table_parent_id_seq")
diff --git a/activerecord/test/cases/adapters/postgresql/schema_test.rb b/activerecord/test/cases/adapters/postgresql/schema_test.rb
index 8756507531..7b065ff320 100644
--- a/activerecord/test/cases/adapters/postgresql/schema_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/schema_test.rb
@@ -366,14 +366,6 @@ class SchemaTest < ActiveRecord::PostgreSQLTestCase
     end
   end
 
-  def test_primary_key_raises_error_if_table_not_found_on_schema_search_path
-    with_schema_search_path(SCHEMA2_NAME) do
-      assert_raises(ActiveRecord::StatementInvalid) do
-        @connection.primary_key(PK_TABLE_NAME)
-      end
-    end
-  end
-
   def test_pk_and_sequence_for_with_schema_specified
     pg_name = ActiveRecord::ConnectionAdapters::PostgreSQL::Name
     [