From 9976570d2fa5a5b9b1ad790b7871ff5984f88de2 Mon Sep 17 00:00:00 2001
From: Joseph Sumabat <joseph@mercury.com>
Date: Sat, 30 Sep 2023 07:56:55 -0400
Subject: [PATCH] Add keyAndEntityFieldsDatabase function to return key with
 migration only fields

---
 persistent/Database/Persist/Types/Base.hs | 25 ++++++++++++++++++++---
 1 file changed, 22 insertions(+), 3 deletions(-)

diff --git a/persistent/Database/Persist/Types/Base.hs b/persistent/Database/Persist/Types/Base.hs
index b17def38b..7f82f38a5 100644
--- a/persistent/Database/Persist/Types/Base.hs
+++ b/persistent/Database/Persist/Types/Base.hs
@@ -201,7 +201,28 @@ entityKeyFields =
 -- columns for an 'EntityDef'.
 keyAndEntityFields :: EntityDef -> NonEmpty FieldDef
 keyAndEntityFields ent =
-    case entityId ent of
+    keyWithFields (entityId ent) fields
+  where
+    fields = entityFields ent
+
+-- | Returns a 'NonEmpty' list of 'FieldDef' that correspond with the key
+-- columns for an 'EntityDef' including those fields that are marked as
+-- 'MigrationOnly' (and therefore only present in the database) or
+-- 'SafeToRemove' (and a migration will drop the column if it exists in the
+-- database).
+--
+-- For fields on the Haskell type use 'keyAndEntityFieldsDatabase'
+--
+-- @since 2.14.7.0
+keyAndEntityFieldsDatabase :: EntityDef -> NonEmpty FieldDef
+keyAndEntityFieldsDatabase ent =
+    keyWithFields (entityId ent) fields
+  where
+    fields = filter isHaskellField $ entityFields ent
+
+keyWithFields :: EntityIdDef -> [FieldDef] -> NonEmpty FieldDef
+keyWithFields entId fields =
+    case entId of
         EntityIdField fd ->
             fd :| fields
         EntityIdNaturalKey _ ->
@@ -214,8 +235,6 @@ keyAndEntityFields ent =
                         ]
                 Just xs ->
                     xs
-  where
-    fields = filter isHaskellField $ entityFields ent
 
 type ExtraLine = [Text]