Add PGUpdateReturning facility

src/Gargantext/Database/Transactional.hs

@@ -15,6 +15,7 @@ module Gargantext.Database.Transactional (
   -- * Smart constructors
   , mkPGQuery
   , mkPGUpdate
+  , mkPGUpdateReturning
   , mkOpaQuery
   , mkOpaUpdate
   , mkOpaInsert
@@ -35,6 +36,7 @@ import Database.PostgreSQL.Simple.Transaction qualified as PG
 import Gargantext.Database.Prelude
 import Opaleye
 import Prelude
+import qualified Control.Exception.Safe as Safe
 
 data DBOperation = DBRead | DBWrite
 
@@ -48,6 +50,10 @@ data DBTransactionOp err (r :: DBOperation) next where
   -- | A Postgres /write/, returning the number of affected rows. It can be used only in
   -- 'DBWrite' transactions.
   PGUpdate  :: PG.ToRow a => PG.Query -> a -> (Int -> next) -> DBTransactionOp err DBWrite next
+  -- | Unlike a 'PGUpdate' that returns the list of affected rows, this can be used
+  -- to write updates that returns a value via the \"RETURNING\" directive. It's the programmer's
+  -- responsibility to ensure that the SQL fragment contains it.
+  PGUpdateReturning  :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> (a -> next) -> DBTransactionOp err DBWrite next
   -- | An Opaleye /read/, returning a list of results. The 'r' in the result is polymorphic
   -- so that reads can be embedded in updates transactions.
   OpaQuery  :: Default FromFields fields a => Select fields -> ([a] -> next) -> DBTransactionOp err r next
@@ -74,12 +80,13 @@ type DBTxCmd err a = forall m env. (IsCmd env err m, HasConnectionPool env) => m
 
 instance Functor (DBTransactionOp err r) where
   fmap f = \case
-    PGQuery   q params cont -> PGQuery   q params (f . cont)
-    PGUpdate  q a cont      -> PGUpdate  q a      (f . cont)
-    OpaQuery  sel cont      -> OpaQuery  sel      (f . cont)
-    OpaInsert ins cont      -> OpaInsert ins      (f . cont)
-    OpaUpdate upd cont      -> OpaUpdate upd      (f . cont)
-    DBFail    err           -> DBFail err
+    PGQuery   q params cont     -> PGQuery   q params (f . cont)
+    PGUpdate  q a cont          -> PGUpdate  q a      (f . cont)
+    PGUpdateReturning  q a cont -> PGUpdateReturning q a      (f . cont)
+    OpaQuery  sel cont          -> OpaQuery  sel      (f . cont)
+    OpaInsert ins cont          -> OpaInsert ins      (f . cont)
+    OpaUpdate upd cont          -> OpaUpdate upd      (f . cont)
+    DBFail    err               -> DBFail err
 
 -- | Generalised version of 'withResource' to work over any unlifted monad.
 -- For some reason 'resource-pool' removed this from version 0.3.0.0 onwards.
@@ -125,12 +132,20 @@ runDBQuery (DBTx m) = do
 -- 'DBCmd'.
 evalOp :: PG.Connection -> DBTransactionOp err r a -> DBTxCmd err a
 evalOp conn = \case
-  PGQuery   qr q cc -> cc <$> liftBase (PG.query conn qr q)
-  PGUpdate  qr a cc -> cc <$> liftBase (fromIntegral <$> PG.execute conn qr a)
-  OpaQuery  sel  cc -> cc <$> liftBase (runSelect conn sel)
-  OpaInsert ins  cc -> cc <$> liftBase (runInsert conn ins)
-  OpaUpdate upd  cc -> cc <$> liftBase (runUpdate conn upd)
-  DBFail    err  -> throwError err
+  PGQuery   qr q cc          -> cc <$> liftBase (PG.query conn qr q)
+  PGUpdate  qr a cc          -> cc <$> liftBase (fromIntegral <$> PG.execute conn qr a)
+  PGUpdateReturning  qr a cc -> cc <$> liftBase (queryOne conn qr a)
+  OpaQuery  sel  cc          -> cc <$> liftBase (runSelect conn sel)
+  OpaInsert ins  cc          -> cc <$> liftBase (runInsert conn ins)
+  OpaUpdate upd  cc          -> cc <$> liftBase (runUpdate conn upd)
+  DBFail    err              -> throwError err
+
+queryOne :: (PG.ToRow q, PG.FromRow r) => PG.Connection -> PG.Query -> q -> IO r
+queryOne conn q v = do
+  rs <- PG.query conn q v
+  case rs of
+    [x] -> pure x
+    _   -> Safe.throwIO $ userError "queryOne: more than one result returned. Have you used the 'RETURNING' directive?"
 
 --
 -- Smart constructors
@@ -148,6 +163,9 @@ mkPGQuery q a = DBTx $ liftF (PGQuery q a id)
 mkPGUpdate :: PG.ToRow a => PG.Query -> a -> DBUpdate err Int
 mkPGUpdate q a = DBTx $ liftF (PGUpdate q a id)
 
+mkPGUpdateReturning :: (PG.ToRow q, PG.FromRow a) => PG.Query -> q -> DBUpdate err a
+mkPGUpdateReturning q a = DBTx $ liftF (PGUpdateReturning q a id)
+
 mkOpaQuery :: Default FromFields fields a
            => Select fields
            -> DBQuery err x [a]

test/Test/Database/Transactions.hs

@@ -140,6 +140,10 @@ getCounterById (CounterId cid) = do
     [c] -> pure c
     rst -> dbFail $ Prelude.userError ("getCounterId returned more than one result: " <> show rst)
 
+insertCounter :: DBUpdate IOException Counter
+insertCounter = do
+  mkPGUpdateReturning [sql| INSERT INTO public.ggtx_test_counter_table(counter_value) VALUES(0) RETURNING id, counter_value|] ()
+
 --
 -- MAIN TESTS
 --
@@ -147,10 +151,17 @@ getCounterById (CounterId cid) = do
 tests :: Spec
 tests = parallel $ around withTestCounterDB $
   describe "Database Transactions" $ do
-    describe "Pure Queries" $ do
-      it "Simple query works" simpleQueryWorks
+    describe "Pure PG Queries" $ do
+      it "Simple query works" simplePGQueryWorks
+    describe "Pure PG Inserts" $ do
+      it "Simple insert works" simplePGInsertWorks
 
-simpleQueryWorks :: DBHandle -> Assertion
-simpleQueryWorks env = flip runReaderT env $ runTestMonad $ do
+simplePGQueryWorks :: DBHandle -> Assertion
+simplePGQueryWorks env = flip runReaderT env $ runTestMonad $ do
   x <- runDBQuery $ getCounterById (CounterId 1)
   liftIO $ counterValue x `shouldBe` 42
+
+simplePGInsertWorks :: DBHandle -> Assertion
+simplePGInsertWorks env = flip runReaderT env $ runTestMonad $ do
+  x <- runDBTx $ insertCounter
+  liftIO $ x `shouldBe` (Counter (CounterId 2) 0)