From a8783b85b761262fd4d6b3efd71964ad55684f25 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Jano=C5=A1=C3=ADk?=
 <ondrej.janosik@scrive.com>
Date: Fri, 10 Jan 2025 16:02:34 +0100
Subject: [PATCH] Allow extra DB enums for rollback compatibility

---
 src/Database/PostgreSQL/PQTypes/Checks.hs | 57 ++++++++++++++++-------
 test/Main.hs                              | 20 ++++++--
 2 files changed, 58 insertions(+), 19 deletions(-)

diff --git a/src/Database/PostgreSQL/PQTypes/Checks.hs b/src/Database/PostgreSQL/PQTypes/Checks.hs
index 398a6c9..80407f4 100644
--- a/src/Database/PostgreSQL/PQTypes/Checks.hs
+++ b/src/Database/PostgreSQL/PQTypes/Checks.hs
@@ -27,7 +27,6 @@ import Control.Arrow ((&&&))
 import Control.Concurrent (threadDelay)
 import Control.Monad
 import Control.Monad.Catch
-import Control.Monad.Except
 import Control.Monad.Writer as W
 import Data.Foldable (foldMap')
 import Data.Function
@@ -386,34 +385,60 @@ checkEnumTypes
   :: (MonadDB m, MonadThrow m)
   => [EnumType]
   -> m ValidationResult
-checkEnumTypes defs = fmap mconcat . forM defs $ \def -> do
+checkEnumTypes defs = fmap mconcat . forM defs $ \defEnum -> do
   runQuery_ . sqlSelect "pg_catalog.pg_type t" $ do
     sqlResult "t.typname::text" -- name
     sqlResult
       "ARRAY(SELECT e.enumlabel::text FROM pg_catalog.pg_enum e WHERE e.enumtypid = t.oid ORDER BY e.enumsortorder)" -- values
-    sqlWhereEq "t.typname" $ unRawSQL $ etName def
+    sqlWhereEq "t.typname" $ unRawSQL $ etName defEnum
   enum <- fetchMaybe $
     \(enumName, enumValues) ->
       EnumType
         { etName = unsafeSQL enumName
         , etValues = map unsafeSQL $ unArray1 enumValues
         }
-  return $ case enum of
-    Just e
-      | e /= def ->
-          validationError $
-            "Enum '"
-              <> unRawSQL (etName e)
-              <> "' does not match (database:"
-              <+> T.pack (show . map unRawSQL $ etValues e)
-              <> ", definition:"
-              <+> T.pack (show . map unRawSQL $ etValues def)
-              <> ")"
-      | otherwise -> mempty
+  pure $ case enum of
+    Just dbEnum -> do
+      let enumName = unRawSQL $ etName defEnum
+          dbValues = map unRawSQL $ etValues dbEnum
+          defValues = map unRawSQL $ etValues defEnum
+          dbSet = S.fromList dbValues
+          defSet = S.fromList defValues
+      if
+        | dbValues == defValues -> mempty
+        | L.sort dbValues == L.sort defValues ->
+            validationInfo $
+              "Enum '"
+                <> enumName
+                <> "' has same values, but differs in order (database:"
+                <+> T.pack (show dbValues)
+                <> ", definition:"
+                <+> T.pack (show defValues)
+                <> ")."
+                <+> "This isn't usually a problem, unless the enum is used for ordering."
+        | S.isSubsetOf defSet dbSet ->
+            validationInfo $
+              "Enum '"
+                <> enumName
+                <> "' has all necessary values, but the database has additional ones"
+                <+> "(database:"
+                <+> T.pack (show dbValues)
+                <> ", definition:"
+                <+> T.pack (show defValues)
+                <> ")"
+        | otherwise ->
+            validationError $
+              "Enum '"
+                <> enumName
+                <> "' does not match (database:"
+                <+> T.pack (show dbValues)
+                <> ", definition:"
+                <+> T.pack (show defValues)
+                <> ")"
     Nothing ->
       validationError $
         "Enum '"
-          <> unRawSQL (etName def)
+          <> unRawSQL (etName defEnum)
           <> "' doesn't exist in the database"
 
 -- | Check that the tables that must have been dropped are actually
diff --git a/test/Main.hs b/test/Main.hs
index fa63fef..9a44ffe 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -2334,8 +2334,12 @@ enumTest connSource =
       report <- checkDatabaseWithReport defaultExtrasOptions (emptyDbDefinitions {dbEnums = [enum1misorder]})
       liftIO $
         assertEqual
-          "Missing enum2 should be reported"
-          (validationError "Enum 'enum1' does not match (database: [\"enum-100\",\"enum-101\"], definition: [\"enum-101\",\"enum-100\"])")
+          "Order mismatch should be reported"
+          ( validationInfo
+              "Enum 'enum1' has same values, but differs in order \
+              \(database: [\"enum-100\",\"enum-101\"], definition: [\"enum-101\",\"enum-100\"]). \
+              \This isn't usually a problem, unless the enum is used for ordering."
+          )
           report
 
     step "Check the database with mismatching enum"
@@ -2343,14 +2347,24 @@ enumTest connSource =
       report <- checkDatabaseWithReport defaultExtrasOptions (emptyDbDefinitions {dbEnums = [enum1mismatch]})
       liftIO $
         assertEqual
-          "Missing enum2 should be reported"
+          "DB mismatch should be reported"
           (validationError "Enum 'enum1' does not match (database: [\"enum-100\",\"enum-101\"], definition: [\"enum-100\",\"enum-102\"])")
           report
+
+    step "Check the database with extra enum values"
+    do
+      report <- checkDatabaseWithReport defaultExtrasOptions (emptyDbDefinitions {dbEnums = [enum1missing]})
+      liftIO $
+        assertEqual
+          "Extra values in the DB enum should be reported"
+          (validationInfo "Enum 'enum1' has all necessary values, but the database has additional ones (database: [\"enum-100\",\"enum-101\"], definition: [\"enum-100\"])")
+          report
   where
     enum1 = EnumType {etName = "enum1", etValues = ["enum-100", "enum-101"]}
     enum2 = EnumType {etName = "enum2", etValues = ["enum-200", "enum-201", "enum-202"]}
     enum1misorder = EnumType {etName = "enum1", etValues = ["enum-101", "enum-100"]}
     enum1mismatch = EnumType {etName = "enum1", etValues = ["enum-100", "enum-102"]}
+    enum1missing = EnumType {etName = "enum1", etValues = ["enum-100"]}
 
 assertNoException :: String -> TestM () -> TestM ()
 assertNoException t =