From 97e5d9e61c8ad34680729a41251d806aa6bf1645 Mon Sep 17 00:00:00 2001
From: Paul Brinkmeier <hallo@pbrinkmeier.de>
Date: Wed, 3 Jul 2024 11:12:07 +0200
Subject: [PATCH] Add Maybe as possible row container type

---
 lib/Database/PostgreSQL/Opium.hs       | 17 +++++++-----
 lib/Database/PostgreSQL/Opium/Error.hs |  1 +
 test/Database/PostgreSQL/OpiumSpec.hs  | 36 +++++++++++++++++---------
 3 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/lib/Database/PostgreSQL/Opium.hs b/lib/Database/PostgreSQL/Opium.hs
index 6184d24..4e33796 100644
--- a/lib/Database/PostgreSQL/Opium.hs
+++ b/lib/Database/PostgreSQL/Opium.hs
@@ -47,16 +47,20 @@ import Database.PostgreSQL.Opium.ToField (ToField (..))
 import Database.PostgreSQL.Opium.ToParamList (ToParamList (..))
 
 class RowContainer c where
-  extract :: FromRow a => Result -> ColumnTable -> ExceptT Error IO (c a)
+  extract :: FromRow a => Result -> LibPQ.Row -> ColumnTable -> ExceptT Error IO (c a)
 
 instance RowContainer [] where
-  extract result columnTable = do
-    nRows <- liftIO $ LibPQ.ntuples result
+  extract result nRows columnTable = do
     mapM (ExceptT . fromRow result columnTable) [0..nRows - 1]
 
+instance RowContainer Maybe where
+  extract result nRows columnTable
+    | nRows == 0 = pure Nothing
+    | nRows == 1 = Just <$> ExceptT (fromRow result columnTable 0)
+    | otherwise = throwE ErrorMoreThanOneRow
+
 instance RowContainer Identity where
-  extract result columnTable = do
-    nRows <- liftIO $ LibPQ.ntuples result
+  extract result nRows columnTable = do
     unless (nRows == 1) $ throwE ErrorNotExactlyOneRow
     Identity <$> ExceptT (fromRow result columnTable 0)
 
@@ -69,8 +73,9 @@ fetch
   -> IO (Either Error (b a))
 fetch query params conn = runExceptT $ do
   result <- execParams conn query params
+  nRows <- liftIO $ LibPQ.ntuples result
   columnTable <- ExceptT $ getColumnTable @a Proxy result
-  extract result columnTable
+  extract result nRows columnTable
 
 fetch_ :: forall a c. (FromRow a, RowContainer c) => Text -> Connection -> IO (Either Error (c a))
 fetch_ query = fetch query ()
diff --git a/lib/Database/PostgreSQL/Opium/Error.hs b/lib/Database/PostgreSQL/Opium/Error.hs
index 57e93f9..2b6e3a8 100644
--- a/lib/Database/PostgreSQL/Opium/Error.hs
+++ b/lib/Database/PostgreSQL/Opium/Error.hs
@@ -18,6 +18,7 @@ data Error
   | ErrorUnexpectedNull ErrorPosition
   | ErrorInvalidField ErrorPosition Oid ByteString String
   | ErrorNotExactlyOneRow
+  | ErrorMoreThanOneRow
   deriving (Eq, Show)
 
 instance Exception Error where
diff --git a/test/Database/PostgreSQL/OpiumSpec.hs b/test/Database/PostgreSQL/OpiumSpec.hs
index 487d43c..f62a62f 100644
--- a/test/Database/PostgreSQL/OpiumSpec.hs
+++ b/test/Database/PostgreSQL/OpiumSpec.hs
@@ -131,18 +131,6 @@ spec = do
       rows <- Opium.fetch "SELECT count(*) AS only FROM person WHERE name = $1" (Identity ("paul" :: Text)) conn
       rows `shouldBe` Right [Only (1 :: Int)]
 
-    it "Accepts exactly one row when Identity is the return type" $ \conn -> do
-      row <- Opium.fetch "SELECT ($1 + $2) AS only" (17 :: Int, 25 :: Int) conn
-      row `shouldBe` Right (Identity (Only (42 :: Int)))
-
-    it "Does not accept zero rows when Identity is the return type" $ \conn -> do
-      row <- Opium.fetch @(Only Int) @Identity "SELECT ($1 + $2) AS only WHERE false" (17 :: Int, 25 :: Int) conn
-      row `shouldSatisfy` isLeft
-
-    it "Does not accept two rows when Identity is the return type" $ \conn -> do
-      row <- Opium.fetch @(Only Int) @Identity "SELECT $1 AS only UNION ALL SELECT $2 AS only" (17 :: Int, 25 :: Int) conn
-      row `shouldSatisfy` isLeft
-
   describe "fetch_" $ do
     it "Retrieves a list of rows" $ \conn -> do
       rows <- Opium.fetch_ "SELECT * FROM person" conn
@@ -163,3 +151,27 @@ spec = do
     it "Works for the readme regression example" $ \conn -> do
       rows <- Opium.fetch_ @ScoreByAge @[] "SELECT regr_intercept(score, age) AS t, regr_slope(score, age) AS m FROM person" conn
       rows `shouldSatisfy` \case { (Right [ScoreByAge _ _]) -> True; _ -> False }
+
+    it "Accepts exactly one row when Identity is the row container type" $ \conn -> do
+      row <- Opium.fetch_ "SELECT 42 AS only" conn
+      row `shouldBe` Right (Identity (Only (42 :: Int)))
+
+    it "Does not accept zero rows when Identity is the row container type" $ \conn -> do
+      row <- Opium.fetch_ @(Only Int) @Identity "SELECT 42 AS only WHERE false" conn
+      row `shouldSatisfy` isLeft
+
+    it "Does not accept two rows when Identity is the row container type" $ \conn -> do
+      row <- Opium.fetch_ @(Only Int) @Identity "SELECT 17 AS only UNION ALL SELECT 25 AS only" conn
+      row `shouldSatisfy` isLeft
+
+    it "Accepts zero rows when Maybe is the row container type" $ \conn -> do
+      row <- Opium.fetch_ @(Only Int) @Maybe "SELECT 17 AS only WHERE false" conn
+      row `shouldBe` Right Nothing
+
+    it "Accepts one row when Maybe is the row container type" $ \conn -> do
+      row <- Opium.fetch_ @(Only Int) @Maybe "SELECT 42 AS only" conn
+      row `shouldBe` Right (Just (Only 42))
+
+    it "Does not accept two rows when Maybe is the row container type" $ \conn -> do
+      row <- Opium.fetch_ @(Only Int) @Maybe "SELECT 17 AS only UNION ALL SELECT 25 AS only" conn
+      row `shouldSatisfy` isLeft