From 19d420302237a1a6e0382cd881dc2810ece1ce63 Mon Sep 17 00:00:00 2001
From: Mona Mayrhofer <mona.mayrhofer@proton.me>
Date: Sat, 2 Mar 2024 16:27:32 +0100
Subject: [PATCH] Fix SchemaCompiler to support non-id-pks

---
 IHP/SchemaCompiler.hs | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/IHP/SchemaCompiler.hs b/IHP/SchemaCompiler.hs
index 12df92462..280f98dab 100644
--- a/IHP/SchemaCompiler.hs
+++ b/IHP/SchemaCompiler.hs
@@ -581,11 +581,21 @@ compileUpdate table@(CreateTable { name, columns }) =
         columnNames = writableColumns
                 |> map (.name)
                 |> intercalate ", "
+
+        primaryKeyPattern = case primaryKeyColumns table of
+                                [] -> error $ "Impossible happened in compileUpdate. No primary keys found for table " <> cs name <> ". At least one primary key is required."
+                                [col] -> col.name
+                                cols -> "(" <> commaSep (map (\col -> col.name) cols) <> ")"
+
+        primaryKeyParameters = case primaryKeyColumns table of
+                                [] -> error $ "Impossible happened in compileUpdate. No primary keys found for table " <> cs name <> ". At least one primary key is required."
+                                [col] -> "?"
+                                cols -> "(" <> commaSep (map (const "?") (primaryKeyColumns table)) <> ")"
     in
         "instance CanUpdate " <> modelName <> " where\n"
         <> indent ("updateRecord model = do\n"
                 <> indent (
-                    "List.head <$> sqlQuery \"UPDATE " <> name <> " SET " <> updates <> " WHERE id = ? RETURNING " <> columnNames <> "\" (" <> bindings <> ")\n"
+                    "List.head <$> sqlQuery \"UPDATE " <> name <> " SET " <> updates <> " WHERE " <> primaryKeyPattern <> " = "<> primaryKeyParameters <> " RETURNING " <> columnNames <> "\" (" <> bindings <> ")\n"
                 )
             )
 
@@ -607,19 +617,25 @@ instance FromRow #{modelName} where
 
         compileField (fieldName, _)
             | isColumn fieldName = fieldName
-            | isManyToManyField fieldName = let (Just ref) = find (\(n, _) -> columnNameToFieldName n == fieldName) referencing in compileSetQueryBuilder ref
+            | isOneToManyField fieldName = let (Just ref) = find (\(n, _) -> columnNameToFieldName n == fieldName) referencing in compileSetQueryBuilder ref
             | fieldName == "meta" = "def { originalDatabaseRecord = Just (Data.Dynamic.toDyn theRecord) }"
             | otherwise = "def"
 
         isPrimaryKey name = name `elem` primaryKeyColumnNames table.primaryKeyConstraint
         isColumn name = name `elem` columnNames
-        isManyToManyField fieldName = fieldName `elem` (referencing |> map (columnNameToFieldName . fst))
+        isOneToManyField fieldName = fieldName `elem` (referencing |> map (columnNameToFieldName . fst))
 
         compileSetQueryBuilder (refTableName, refFieldName) = "(QueryBuilder.filterWhere (#" <> columnNameToFieldName refFieldName <> ", " <> primaryKeyField <> ") (QueryBuilder.query @" <> tableNameToModelName refTableName <> "))"
             where
                 -- | When the referenced column is nullable, we have to wrap the @Id@ in @Just@
                 primaryKeyField :: Text
-                primaryKeyField = if refColumn.notNull then "id" else "Just id"
+                primaryKeyField = if refColumn.notNull then actualPrimaryKeyField else "Just " <> actualPrimaryKeyField
+                actualPrimaryKeyField :: Text
+                actualPrimaryKeyField = case primaryKeyColumns table of
+                        [] -> error $ "Impossible happened in compilePrimaryKeyInstance. No primary keys found for table " <> cs name <> ". At least one primary key is required."
+                        [pk] -> columnNameToFieldName pk.name
+                        pks -> error $ "No support yet for composite foreign keys. Tables cannot have foreign keys to table '" <> cs name <> "' which has more than one column as its primary key."
+
 
                 (Just refTable) = let (Schema statements) = ?schema in
                         statements