From 24570678d6845368074ebea3072195d6b027a212 Mon Sep 17 00:00:00 2001
From: Nikita Vasilevsky <nikita.vasilevsky@shopify.com>
Date: Thu, 18 Jan 2024 17:20:01 +0000
Subject: [PATCH] [Tests only] Expand trigger populated PK test case to run
 against more DBs

---
 activerecord/test/cases/persistence_test.rb |  2 +-
 activerecord/test/schema/schema.rb          | 51 +++++++++++++--------
 2 files changed, 32 insertions(+), 21 deletions(-)

diff --git a/activerecord/test/cases/persistence_test.rb b/activerecord/test/cases/persistence_test.rb
index 44bf27af00170..a14d3d283d444 100644
--- a/activerecord/test/cases/persistence_test.rb
+++ b/activerecord/test/cases/persistence_test.rb
@@ -1581,7 +1581,7 @@ def test_model_with_no_auto_populated_fields_still_returns_primary_key_after_ins
 
     assert_not_nil record.id
     assert record.id > 0
-  end if current_adapter?(:PostgreSQLAdapter)
+  end if supports_insert_returning? && !current_adapter?(:SQLite3Adapter)
 end
 
 class QueryConstraintsTest < ActiveRecord::TestCase
diff --git a/activerecord/test/schema/schema.rb b/activerecord/test/schema/schema.rb
index a8bdb4e8f466a..fa6e0bb8256d7 100644
--- a/activerecord/test/schema/schema.rb
+++ b/activerecord/test/schema/schema.rb
@@ -1483,30 +1483,41 @@
   end
 end
 
-if ActiveRecord::TestCase.current_adapter?(:PostgreSQLAdapter)
+if ActiveRecord::Base.connection.supports_insert_returning? && !ActiveRecord::TestCase.current_adapter?(:SQLite3Adapter)
   ActiveRecord::Base.connection.create_table :pk_autopopulated_by_a_trigger_records, force: true, id: false do |t|
     t.integer :id, null: false
   end
 
-  ActiveRecord::Base.connection.execute(
-    <<-SQL
-      CREATE OR REPLACE FUNCTION populate_column()
-      RETURNS TRIGGER AS $$
-      DECLARE
-        max_value INTEGER;
-      BEGIN
-          SELECT MAX(id) INTO max_value FROM pk_autopopulated_by_a_trigger_records;
-          NEW.id = COALESCE(max_value, 0) + 1;
-          RETURN NEW;
-      END;
-      $$ LANGUAGE plpgsql;
-
-      CREATE TRIGGER before_insert_trigger
-      BEFORE INSERT ON pk_autopopulated_by_a_trigger_records
-      FOR EACH ROW
-      EXECUTE FUNCTION populate_column();
-    SQL
-  )
+  if ActiveRecord::TestCase.current_adapter?(:PostgreSQLAdapter)
+    ActiveRecord::Base.connection.execute(
+      <<-SQL
+        CREATE OR REPLACE FUNCTION populate_column()
+        RETURNS TRIGGER AS $$
+        DECLARE
+          max_value INTEGER;
+        BEGIN
+            SELECT MAX(id) INTO max_value FROM pk_autopopulated_by_a_trigger_records;
+            NEW.id = COALESCE(max_value, 0) + 1;
+            RETURN NEW;
+        END;
+        $$ LANGUAGE plpgsql;
+
+        CREATE TRIGGER before_insert_trigger
+        BEFORE INSERT ON pk_autopopulated_by_a_trigger_records
+        FOR EACH ROW
+        EXECUTE FUNCTION populate_column();
+      SQL
+    )
+  elsif ActiveRecord::TestCase.current_adapter?(:Mysql2Adapter, :TrilogyAdapter)
+    ActiveRecord::Base.connection.execute(
+      <<-SQL
+        CREATE TRIGGER before_insert_trigger
+        BEFORE INSERT ON pk_autopopulated_by_a_trigger_records
+        FOR EACH ROW
+        SET NEW.id = (SELECT COALESCE(MAX(id), 0) + 1 FROM pk_autopopulated_by_a_trigger_records);
+      SQL
+    )
+  end
 end
 
 Course.connection.create_table :courses, force: true do |t|