postgresql: Add foreign key constraints the first time when migrating

Author: hdgarrood

persistent-postgresql/Database/Persist/Postgresql/Internal/Migration.hs

@@ -1097,42 +1097,33 @@ findAlters
     -> EntityDef
     -- ^ The entity definition for the entity that we're working on.
     -> Column
-    -- ^ The column that we're searching for potential alterations for.
+    -- ^ The column that we're searching for potential alterations for, derived
+    -- from the Persistent EntityDef. That is: this is how we _want_ the column
+    -- to look, and not necessarily how it actually looks in the database right
+    -- now.
     -> [Column]
+    -- ^ The columns for this table, as they currently exist in the database.
     -> ([AlterColumn], [Column])
 findAlters defs edef col@(Column name isNull sqltype def _gen _defConstraintName _maxLen ref) cols =
     case List.find (\c -> cName c == name) cols of
         Nothing ->
-            ([AddColumn col], cols)
+            ([AddColumn col] ++ refAdd ref, cols)
         Just
-            (Column _oldName isNull' sqltype' def' _gen' _defConstraintName' _maxLen' ref') ->
+            (Column oldName isNull' sqltype' def' _gen' _defConstraintName' _maxLen' ref') ->
                 let
                     refDrop Nothing = []
                     refDrop (Just ColumnReference{crConstraintName = cname}) =
                         [DropReference cname]
 
-                    refAdd Nothing = []
-                    refAdd (Just colRef) =
-                        case find ((== crTableName colRef) . getEntityDBName) defs of
-                            Just refdef
-                                | Just _oldName /= fmap fieldDB (getEntityIdField edef) ->
-                                    [ AddReference
-                                        (crTableName colRef)
-                                        (crConstraintName colRef)
-                                        (name NEL.:| [])
-                                        (NEL.toList $ Util.dbIdColumnsEsc escapeF refdef)
-                                        (crFieldCascade colRef)
-                                    ]
-                            Just _ -> []
-                            Nothing ->
-                                error $
-                                    "could not find the entityDef for reftable["
-                                        ++ show (crTableName colRef)
-                                        ++ "]"
                     modRef =
                         if equivalentRef ref ref'
                             then []
-                            else refDrop ref' ++ refAdd ref
+                            else
+                                refDrop ref'
+                                    ++ ( do
+                                            guard $ Just oldName /= fmap fieldDB (getEntityIdField edef)
+                                            refAdd ref
+                                       )
                     modNull = case (isNull, isNull') of
                         (True, False) -> do
                             guard $ Just name /= fmap fieldDB (getEntityIdField edef)
@@ -1174,3 +1165,20 @@ findAlters defs edef col@(Column name isNull sqltype def _gen _defConstraintName
                     ( modRef ++ modDef ++ modNull ++ modType ++ dropSafe
                     , filter (\c -> cName c /= name) cols
                     )
+  where
+    refAdd Nothing = []
+    refAdd (Just colRef) =
+        case find ((== crTableName colRef) . getEntityDBName) defs of
+            Just refdef ->
+                [ AddReference
+                    (crTableName colRef)
+                    (crConstraintName colRef)
+                    (name NEL.:| [])
+                    (NEL.toList $ Util.dbIdColumnsEsc escapeF refdef)
+                    (crFieldCascade colRef)
+                ]
+            Nothing ->
+                error $
+                    "could not find the entityDef for reftable["
+                        ++ show (crTableName colRef)
+                        ++ "]"

persistent-postgresql/test/MigrationSpec.hs

@@ -60,6 +60,14 @@ AdminUser sql=admin_users
 
     promotedByUserId UserId
     UniquePromotedByUserId promotedByUserId
+
+FKParent sql=migration_fk_parent
+
+FKChildV1 sql=migration_fk_child
+
+-- Simulate creating a new FK field on an existing table
+FKChildV2 sql=migration_fk_child
+    parentId FKParentId
 |]
 
 userEntityDef :: EntityDef
@@ -77,15 +85,28 @@ password2EntityDef = entityDef (Proxy :: Proxy Password2)
 adminUserEntityDef :: EntityDef
 adminUserEntityDef = entityDef (Proxy :: Proxy AdminUser)
 
+fkParentEntityDef :: EntityDef
+fkParentEntityDef = entityDef (Proxy :: Proxy FKParent)
+
+fkChildV1EntityDef :: EntityDef
+fkChildV1EntityDef = entityDef (Proxy :: Proxy FKChildV1)
+
+fkChildV2EntityDef :: EntityDef
+fkChildV2EntityDef = entityDef (Proxy :: Proxy FKChildV2)
+
+-- Note that FKChild is deliberately omitted here because we have two
+-- versions of it
 allEntityDefs :: [EntityDef]
 allEntityDefs =
     [ userEntityDef
     , userFriendshipEntityDef
     , passwordEntityDef
     , password2EntityDef
     , adminUserEntityDef
+    , fkParentEntityDef
     ]
 
+-- Note that this function migrates to the schema expected by FKChildV1
 migrateManually :: (HasCallStack, MonadIO m) => SqlPersistT m ()
 migrateManually = do
     cleanDB
@@ -150,6 +171,8 @@ migrateManually = do
             , "  ADD CONSTRAINT unique_promoted_by_user_id"
             , "  UNIQUE(promoted_by_user_id);"
             ]
+    rawEx "CREATE TABLE migration_fk_parent(id int8 primary key);"
+    rawEx "CREATE TABLE migration_fk_child(id int8 primary key);"
     rawEx "CREATE TABLE ignored(id int8 primary key);"
 
 cleanDB :: (HasCallStack, MonadIO m) => SqlPersistT m ()
@@ -162,6 +185,8 @@ cleanDB = do
     rawEx "DROP TABLE IF EXISTS ignored;"
     rawEx "DROP TABLE IF EXISTS admin_users;"
     rawEx "DROP TABLE IF EXISTS users;"
+    rawEx "DROP TABLE IF EXISTS migration_fk_child;"
+    rawEx "DROP TABLE IF EXISTS migration_fk_parent;"
 
 spec :: Spec
 spec = describe "MigrationSpec" $ do
@@ -582,3 +607,27 @@ spec = describe "MigrationSpec" $ do
                 result2 <-
                     liftIO $ migrateEntitiesStructured getter allEntityDefs allEntityDefs
                 result2 `shouldBe` Right []
+
+    it "suggests FK constraints for new fields first time" $ runConnAssert $ do
+        migrateManually
+
+        getter <- getStmtGetter
+        result <-
+            liftIO $
+                migrateEntitiesStructured
+                    getter
+                    (fkChildV2EntityDef : allEntityDefs)
+                    [fkChildV2EntityDef]
+
+        cleanDB
+
+        case result of
+            Right [] ->
+                pure ()
+            Left err ->
+                expectationFailure $ show err
+            Right alters ->
+                map (snd . showAlterDb) alters
+                    `shouldBe` [ "ALTER TABLE \"migration_fk_child\" ADD COLUMN \"parent_id\" INT8 NOT NULL"
+                               , "ALTER TABLE \"migration_fk_child\" ADD CONSTRAINT \"migration_fk_child_parent_id_fkey\" FOREIGN KEY(\"parent_id\") REFERENCES \"migration_fk_parent\"(\"id\") ON DELETE RESTRICT  ON UPDATE RESTRICT"
+                               ]