From 7e04d0b2d05414ea01e3929841e010de4b7e13ef Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Fri, 24 Jan 2025 21:49:39 +0200
Subject: [PATCH] Support composite primary keys

---
 lib/active_record_doctor/detectors/base.rb    |  6 ++---
 .../detectors/extraneous_indexes.rb           | 18 +++++++++------
 .../detectors/missing_presence_validation.rb  |  2 +-
 .../detectors/short_primary_key_type.rb       |  6 +++--
 .../detectors/extraneous_indexes_test.rb      | 22 +++++++++++++++++++
 .../missing_presence_validation_test.rb       |  9 ++++++++
 .../detectors/short_primary_key_type_test.rb  | 11 ++++++++++
 .../table_without_primary_key_test.rb         | 10 +++++++++
 test/setup.rb                                 |  4 ++++
 9 files changed, 75 insertions(+), 13 deletions(-)

diff --git a/lib/active_record_doctor/detectors/base.rb b/lib/active_record_doctor/detectors/base.rb
index b02630ba..278e4700 100644
--- a/lib/active_record_doctor/detectors/base.rb
+++ b/lib/active_record_doctor/detectors/base.rb
@@ -118,10 +118,10 @@ def indexes(table_name)
       end
 
       def primary_key(table_name)
-        primary_key_name = connection.primary_key(table_name)
-        return nil if primary_key_name.nil?
+        primary_key_names = Array(connection.primary_key(table_name))
+        return nil if primary_key_names.empty?
 
-        column(table_name, primary_key_name)
+        primary_key_names.map { |column_name| column(table_name, column_name) }
       end
 
       def column(table_name, column_name)
diff --git a/lib/active_record_doctor/detectors/extraneous_indexes.rb b/lib/active_record_doctor/detectors/extraneous_indexes.rb
index ea1574bd..aeca93ef 100644
--- a/lib/active_record_doctor/detectors/extraneous_indexes.rb
+++ b/lib/active_record_doctor/detectors/extraneous_indexes.rb
@@ -62,7 +62,7 @@ def indexed_primary_keys
           each_table(except: config(:ignore_tables)) do |table|
             each_index(table, except: config(:ignore_indexes), multicolumn_only: true) do |index|
               primary_key = connection.primary_key(table)
-              if index.columns == [primary_key] && index.where.nil?
+              if replaceable_columns?(index, primary_key, true, []) && index.where.nil?
                 problem!(table: table, extraneous_index: index.name, replacement_indexes: nil)
               end
             end
@@ -76,17 +76,21 @@ def replaceable_with?(index1, index2)
         return false if index1.where != index2.where
         return false if opclasses(index1) != opclasses(index2)
 
-        index1_columns = Array(index1.columns)
-        index2_columns = Array(index2.columns)
+        replaceable_columns?(index1, index2.columns, index2.unique, includes(index2))
+      end
+
+      def replaceable_columns?(index, columns2, unique2, includes2)
+        columns1 = Array(index.columns)
+        columns2 = Array(columns2)
 
-        case [index1.unique, index2.unique]
+        case [index.unique, unique2]
         when [true, true]
-          contains_all?(index1_columns, index2_columns)
+          contains_all?(columns1, columns2)
         when [true, false]
           false
         else
-          prefix?(index1_columns, index2_columns) &&
-            contains_all?(index2_columns + includes(index2), includes(index1))
+          prefix?(columns1, columns2) &&
+            contains_all?(columns2 + includes2, includes(index))
         end
       end
 
diff --git a/lib/active_record_doctor/detectors/missing_presence_validation.rb b/lib/active_record_doctor/detectors/missing_presence_validation.rb
index 4924116b..40350b38 100644
--- a/lib/active_record_doctor/detectors/missing_presence_validation.rb
+++ b/lib/active_record_doctor/detectors/missing_presence_validation.rb
@@ -38,7 +38,7 @@ def detect
       end
 
       def validator_needed?(model, column)
-        ![model.primary_key, "created_at", "updated_at", "created_on", "updated_on"].include?(column.name) &&
+        ![*Array(model.primary_key), "created_at", "updated_at", "created_on", "updated_on"].include?(column.name) &&
           (!column.null || not_null_check_constraint_exists?(model.table_name, column)) &&
           !default_value_instead_of_validation?(column)
       end
diff --git a/lib/active_record_doctor/detectors/short_primary_key_type.rb b/lib/active_record_doctor/detectors/short_primary_key_type.rb
index a1a481b2..4b9ec59b 100644
--- a/lib/active_record_doctor/detectors/short_primary_key_type.rb
+++ b/lib/active_record_doctor/detectors/short_primary_key_type.rb
@@ -23,8 +23,10 @@ def detect
         return if ActiveRecordDoctor::Utils.sqlite?
 
         each_table(except: config(:ignore_tables)) do |table|
-          column = primary_key(table)
-          next if column.nil?
+          columns = primary_key(table)
+          next if columns.nil? || columns.size > 1
+
+          column = columns.first
           next if !integer?(column) || bigint?(column)
 
           problem!(table: table, column: column.name)
diff --git a/test/active_record_doctor/detectors/extraneous_indexes_test.rb b/test/active_record_doctor/detectors/extraneous_indexes_test.rb
index 4592a1d0..2cf2be59 100644
--- a/test/active_record_doctor/detectors/extraneous_indexes_test.rb
+++ b/test/active_record_doctor/detectors/extraneous_indexes_test.rb
@@ -22,6 +22,28 @@ def test_partial_index_on_primary_key
     refute_problems
   end
 
+  def test_index_on_composite_primary_key_columns
+    require_composite_primary_keys!
+
+    Context.create_table(:users, primary_key: [:company_id, :id]) do |t|
+      t.bigint :company_id
+      t.bigint :id
+      t.string :email
+      t.index [:company_id, :id], name: "index_1"
+      t.index [:company_id], name: "index_2"
+      t.index [:id], name: "index_3"
+      t.index [:company_id, :id, :email], unique: true, name: "index_4"
+    end
+
+    assert_problems(<<OUTPUT)
+remove index_1 from users - coincides with the primary key on the table
+remove index_2 from users - coincides with the primary key on the table
+remove index_4 from users - coincides with the primary key on the table
+remove the index index_1 from the table users - queries should be able to use the following index instead: index_4
+remove the index index_2 from the table users - queries should be able to use the following indices instead: index_1 or index_4
+OUTPUT
+  end
+
   def test_index_on_non_standard_primary_key
     Context.create_table(:profiles, primary_key: :user_id) do |t|
       t.index :user_id
diff --git a/test/active_record_doctor/detectors/missing_presence_validation_test.rb b/test/active_record_doctor/detectors/missing_presence_validation_test.rb
index 1572a5fa..bd1f8418 100644
--- a/test/active_record_doctor/detectors/missing_presence_validation_test.rb
+++ b/test/active_record_doctor/detectors/missing_presence_validation_test.rb
@@ -167,6 +167,15 @@ def test_counter_caches_with_custom_names_are_not_reported
     end.define_model do
       belongs_to :company, counter_cache: :custom_users_count
     end
+  end
+
+  def test_composite_primary_key_is_not_reported
+    require_composite_primary_keys!
+
+    Context.create_table(:users, primary_key: [:company_id, :id]) do |t|
+      t.bigint :company_id
+      t.bigint :id
+    end.define_model
 
     refute_problems
   end
diff --git a/test/active_record_doctor/detectors/short_primary_key_type_test.rb b/test/active_record_doctor/detectors/short_primary_key_type_test.rb
index b66b9647..a0435112 100644
--- a/test/active_record_doctor/detectors/short_primary_key_type_test.rb
+++ b/test/active_record_doctor/detectors/short_primary_key_type_test.rb
@@ -63,6 +63,17 @@ def test_no_primary_key_is_not_reported
     refute_problems
   end
 
+  def test_composite_primary_key_is_not_reported
+    require_composite_primary_keys!
+
+    Context.create_table(:companies, primary_key: [:country_id, :id]) do |t|
+      t.integer :country_id
+      t.integer :id
+    end
+
+    refute_problems
+  end
+
   def test_config_ignore_tables
     skip if sqlite?
 
diff --git a/test/active_record_doctor/detectors/table_without_primary_key_test.rb b/test/active_record_doctor/detectors/table_without_primary_key_test.rb
index 0804607d..1c2fcea2 100644
--- a/test/active_record_doctor/detectors/table_without_primary_key_test.rb
+++ b/test/active_record_doctor/detectors/table_without_primary_key_test.rb
@@ -16,6 +16,16 @@ def test_table_with_primary_key_is_not_reported
     refute_problems
   end
 
+  def test_table_with_composite_primary_key_is_not_reported
+    require_composite_primary_keys!
+
+    Context.create_table(:orders, primary_key: [:shop_id, :id]) do |t|
+      t.bigint :shop_id
+      t.bigint :id
+    end
+    refute_problems
+  end
+
   def test_config_ignore_tables
     Context.create_table(:companies, id: false) do |t|
       t.string :name
diff --git a/test/setup.rb b/test/setup.rb
index 63822336..e82542c2 100644
--- a/test/setup.rb
+++ b/test/setup.rb
@@ -177,6 +177,10 @@ def require_foreign_keys_of_different_type!
     skip("#{current_adapter} doesn't support foreign keys of different type than the referenced column") if mysql?
   end
 
+  def require_composite_primary_keys!
+    skip("ActiveRecord < 7.1 doesn't support composite primary keys") if ActiveRecord::VERSION::STRING < "7.1"
+  end
+
   def current_adapter
     ActiveRecord::Base.connection.adapter_name
   end