From d6529af2954a67bd57fda45286fa9cfd0ff6b5ac Mon Sep 17 00:00:00 2001
From: Jordan Lewis <jordanthelewis@gmail.com>
Date: Wed, 8 Feb 2017 17:46:54 -0500
Subject: [PATCH] Simplify and speed up Postgres query for primary_keys

primary_keys(table) needs to query various metadata tables in Postgres to
determine the primary key for the table. Previously, it did so using a
complex common table expression against pg_constraint and pg_attribute.

This patch simplifies the query by joining pg_index against pg_attribute
instead of going through pg_constraint. This avoids an expensive unnest,
window function query, and common table expression.

EXPLAINing these queries in Postgres against a database with a single
table with a composite primary key shows a 66% reduction in the plan and
execute latencies. This is significant during application startup time,
especially against very large schemas, where these queries would be even
slower and more numerous.

Closes #27949
---
 .../postgresql/schema_statements.rb            | 18 ++++++++----------
 activerecord/test/cases/primary_keys_test.rb   |  9 +++++++++
 2 files changed, 17 insertions(+), 10 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 a4b1723fc7..bd8aab526e 100644
--- a/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
+++ b/activerecord/lib/active_record/connection_adapters/postgresql/schema_statements.rb
@@ -426,16 +426,14 @@ def pk_and_sequence_for(table) #:nodoc:
 
         def primary_keys(table_name) # :nodoc:
           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 a.attname
+                  FROM pg_index i
+            CROSS JOIN unnest(i.indkey) as k
+                  JOIN pg_attribute a
+                    ON a.attrelid = i.indrelid
+                   AND a.attnum = k
+                 WHERE i.indrelid = #{quote(quote_table_name(table_name))}::regclass
+                   AND i.indisprimary
           SQL
         end
 
diff --git a/activerecord/test/cases/primary_keys_test.rb b/activerecord/test/cases/primary_keys_test.rb
index 1baf7c4902..0c6799c015 100644
--- a/activerecord/test/cases/primary_keys_test.rb
+++ b/activerecord/test/cases/primary_keys_test.rb
@@ -298,6 +298,10 @@ def setup
       t.string :region
       t.integer :code
     end
+    @connection.create_table(:barcodes_reverse, primary_key: ["code", "region"], force: true) do |t|
+      t.string :region
+      t.integer :code
+    end
   end
 
   def teardown
@@ -308,6 +312,11 @@ def test_composite_primary_key
     assert_equal ["region", "code"], @connection.primary_keys("barcodes")
   end
 
+  def test_composite_primary_key_out_of_order
+    skip if current_adapter?(:SQLite3Adapter)
+    assert_equal ["region", "code"], @connection.primary_keys("barcodes")
+  end
+
   def test_primary_key_issues_warning
     model = Class.new(ActiveRecord::Base) do
       def self.table_name
-- 
GitLab