1 files changed, 17 insertions, 4 deletions
diff --git a/src/server/Controller/Payment.hs b/src/server/Controller/Payment.hs
index 9155a78..e3f1082 100644
--- a/src/server/Controller/Payment.hs
+++ b/src/server/Controller/Payment.hs
@@ -23,6 +23,7 @@ import Json (jsonId)
 
 import Model.Database
 import qualified Model.Payment as Payment
+import qualified Model.PaymentCategory as PaymentCategory
 import qualified Model.Json.CreatePayment as Json
 import qualified Model.Json.EditPayment as Json
 
@@ -33,15 +34,27 @@ list =
   )
 
 create :: Json.CreatePayment -> ActionM ()
-create (Json.CreatePayment name cost date frequency) =
+create (Json.CreatePayment name cost date category frequency) =
   Secure.loggedAction (\user ->
-    (liftIO . runDb $ Payment.create (entityKey user) name cost date frequency) >>= jsonId
+    (liftIO . runDb $ do
+      PaymentCategory.set name category
+      Payment.create (entityKey user) name cost date frequency
+    ) >>= jsonId
   )
 
 editOwn :: Json.EditPayment -> ActionM ()
-editOwn (Json.EditPayment paymentId name cost date frequency) =
+editOwn (Json.EditPayment paymentId name cost date category frequency) =
   Secure.loggedAction (\user -> do
-    updated <- liftIO . runDb $ Payment.editOwn (entityKey user) paymentId name cost date frequency
+    updated <- liftIO . runDb $ do
+      mbPayment <- fmap entityVal <$> Payment.find paymentId
+      case mbPayment of
+        Just payment -> do
+          edited <- Payment.editOwn (entityKey user) paymentId name cost date frequency
+          if edited
+            then PaymentCategory.edit (paymentName payment) name category >> return True
+            else return edited
+        _ ->
+          return False
     if updated
       then status ok200
       else status badRequest400