diff --git a/lib/Database/PostgreSQL/Opium.hs b/lib/Database/PostgreSQL/Opium.hs index 74b78a7..6184d24 100644 --- a/lib/Database/PostgreSQL/Opium.hs +++ b/lib/Database/PostgreSQL/Opium.hs @@ -26,9 +26,10 @@ module Database.PostgreSQL.Opium ) where -import Control.Monad (void) +import Control.Monad (unless, void) import Control.Monad.IO.Class (liftIO) -import Control.Monad.Trans.Except (ExceptT (..), except, runExceptT) +import Control.Monad.Trans.Except (ExceptT (..), except, runExceptT, throwE) +import Data.Functor.Identity (Identity (..)) import Data.Proxy (Proxy (..)) import Data.Text (Text) import Database.PostgreSQL.LibPQ @@ -41,24 +42,37 @@ import qualified Database.PostgreSQL.LibPQ as LibPQ import Database.PostgreSQL.Opium.Error (Error (..), ErrorPosition (..)) import Database.PostgreSQL.Opium.FromField (FromField (..), RawField (..)) -import Database.PostgreSQL.Opium.FromRow (FromRow (..)) +import Database.PostgreSQL.Opium.FromRow (FromRow (..), ColumnTable) import Database.PostgreSQL.Opium.ToField (ToField (..)) import Database.PostgreSQL.Opium.ToParamList (ToParamList (..)) --- The order of the type parameters is important, because it is more common to use type applications for providing the row type. +class RowContainer c where + extract :: FromRow a => Result -> ColumnTable -> ExceptT Error IO (c a) + +instance RowContainer [] where + extract result columnTable = do + nRows <- liftIO $ LibPQ.ntuples result + mapM (ExceptT . fromRow result columnTable) [0..nRows - 1] + +instance RowContainer Identity where + extract result columnTable = do + nRows <- liftIO $ LibPQ.ntuples result + unless (nRows == 1) $ throwE ErrorNotExactlyOneRow + Identity <$> ExceptT (fromRow result columnTable 0) + +-- The order of the type parameters is important, because it is more common to use type applications for providing the row type and row container type. fetch - :: forall a b. (ToParamList b, FromRow a) + :: forall a b c. (ToParamList c, FromRow a, RowContainer b) => Text - -> b + -> c -> Connection - -> IO (Either Error [a]) + -> IO (Either Error (b a)) fetch query params conn = runExceptT $ do result <- execParams conn query params columnTable <- ExceptT $ getColumnTable @a Proxy result - nRows <- liftIO $ LibPQ.ntuples result - mapM (ExceptT . fromRow result columnTable) [0..nRows - 1] + extract result columnTable -fetch_ :: forall a. FromRow a => Text -> Connection -> IO (Either Error [a]) +fetch_ :: forall a c. (FromRow a, RowContainer c) => Text -> Connection -> IO (Either Error (c a)) fetch_ query = fetch query () execute diff --git a/lib/Database/PostgreSQL/Opium/Error.hs b/lib/Database/PostgreSQL/Opium/Error.hs index 1ac3531..57e93f9 100644 --- a/lib/Database/PostgreSQL/Opium/Error.hs +++ b/lib/Database/PostgreSQL/Opium/Error.hs @@ -17,6 +17,7 @@ data Error | ErrorInvalidOid Text Oid | ErrorUnexpectedNull ErrorPosition | ErrorInvalidField ErrorPosition Oid ByteString String + | ErrorNotExactlyOneRow deriving (Eq, Show) instance Exception Error where diff --git a/lib/Database/PostgreSQL/Opium/FromRow.hs b/lib/Database/PostgreSQL/Opium/FromRow.hs index de5a4e6..60bcc06 100644 --- a/lib/Database/PostgreSQL/Opium/FromRow.hs +++ b/lib/Database/PostgreSQL/Opium/FromRow.hs @@ -14,6 +14,7 @@ module Database.PostgreSQL.Opium.FromRow -- * FromRow ( FromRow (..) -- * Internal + , ColumnTable , toListColumnTable ) where diff --git a/test/Database/PostgreSQL/OpiumSpec.hs b/test/Database/PostgreSQL/OpiumSpec.hs index 532d2e3..487d43c 100644 --- a/test/Database/PostgreSQL/OpiumSpec.hs +++ b/test/Database/PostgreSQL/OpiumSpec.hs @@ -131,23 +131,35 @@ 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 rows `shouldBe` Right [Person "paul" 25, Person "albus" 103] it "Fails for invalid queries" $ \conn -> do - rows <- Opium.fetch_ @Person "MRTLBRNFT" conn + rows <- Opium.fetch_ @Person @[] "MRTLBRNFT" conn rows `shouldSatisfy` isLeft it "Fails for unexpected NULLs" $ \conn -> do - rows <- Opium.fetch_ @Person "SELECT NULL AS name, 0 AS age" conn + rows <- Opium.fetch_ @Person @[] "SELECT NULL AS name, 0 AS age" conn rows `shouldBe` Left (Opium.ErrorUnexpectedNull (Opium.ErrorPosition 0 "name")) it "Fails for the wrong column type" $ \conn -> do - rows <- Opium.fetch_ @Person "SELECT 'quby' AS name, 'indeterminate' AS age" conn + rows <- Opium.fetch_ @Person @[] "SELECT 'quby' AS name, 'indeterminate' AS age" conn rows `shouldBe` Left (Opium.ErrorInvalidOid "age" $ LibPQ.Oid 25) 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 <- 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 }