module Database.PostgreSQL.PQTypes.Fold (
    queryResult
  , foldrDB
  , foldlDB
  , mapDB_
  , fetchMany
  , fetchMaybe
  , fetchOne
  ) where

import Control.Monad.Catch
import qualified Data.Foldable as F

import Database.PostgreSQL.PQTypes.Class
import Database.PostgreSQL.PQTypes.FromRow
import Database.PostgreSQL.PQTypes.Internal.Error
import Database.PostgreSQL.PQTypes.Internal.QueryResult
import Database.PostgreSQL.PQTypes.Utils

-- | Get current 'QueryResult' or throw an exception if there isn't one.
{-# INLINABLE queryResult #-}
queryResult :: (MonadDB m, MonadThrow m, FromRow row) => m (QueryResult row)
queryResult :: m (QueryResult row)
queryResult = m (Maybe (QueryResult row))
forall (m :: * -> *) row.
(MonadDB m, FromRow row) =>
m (Maybe (QueryResult row))
getQueryResult
  m (Maybe (QueryResult row))
-> (Maybe (QueryResult row) -> m (QueryResult row))
-> m (QueryResult row)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= m (QueryResult row)
-> (QueryResult row -> m (QueryResult row))
-> Maybe (QueryResult row)
-> m (QueryResult row)
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (HPQTypesError -> m (QueryResult row)
forall e (m :: * -> *) a.
(Exception e, MonadDB m, MonadThrow m) =>
e -> m a
throwDB (HPQTypesError -> m (QueryResult row))
-> (String -> HPQTypesError) -> String -> m (QueryResult row)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> HPQTypesError
HPQTypesError (String -> m (QueryResult row)) -> String -> m (QueryResult row)
forall a b. (a -> b) -> a -> b
$ String
"queryResult: no query result") QueryResult row -> m (QueryResult row)
forall (m :: * -> *) a. Monad m => a -> m a
return

----------------------------------------

-- | Specialization of 'F.foldrM' for convenient query results fetching.
{-# INLINABLE foldrDB #-}
foldrDB :: (MonadDB m, FromRow row) => (row -> acc -> m acc) -> acc -> m acc
foldrDB :: (row -> acc -> m acc) -> acc -> m acc
foldrDB row -> acc -> m acc
f acc
acc = m acc
-> (QueryResult row -> m acc) -> Maybe (QueryResult row) -> m acc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (acc -> m acc
forall (m :: * -> *) a. Monad m => a -> m a
return acc
acc) ((row -> acc -> m acc) -> acc -> QueryResult row -> m acc
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> b -> m b) -> b -> t a -> m b
F.foldrM row -> acc -> m acc
f acc
acc) (Maybe (QueryResult row) -> m acc)
-> m (Maybe (QueryResult row)) -> m acc
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Maybe (QueryResult row))
forall (m :: * -> *) row.
(MonadDB m, FromRow row) =>
m (Maybe (QueryResult row))
getQueryResult

-- | Specialization of 'F.foldlM' for convenient query results fetching.
{-# INLINABLE foldlDB #-}
foldlDB :: (MonadDB m, FromRow row) => (acc -> row -> m acc) -> acc -> m acc
foldlDB :: (acc -> row -> m acc) -> acc -> m acc
foldlDB acc -> row -> m acc
f acc
acc = m acc
-> (QueryResult row -> m acc) -> Maybe (QueryResult row) -> m acc
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (acc -> m acc
forall (m :: * -> *) a. Monad m => a -> m a
return acc
acc) ((acc -> row -> m acc) -> acc -> QueryResult row -> m acc
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
F.foldlM acc -> row -> m acc
f acc
acc) (Maybe (QueryResult row) -> m acc)
-> m (Maybe (QueryResult row)) -> m acc
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Maybe (QueryResult row))
forall (m :: * -> *) row.
(MonadDB m, FromRow row) =>
m (Maybe (QueryResult row))
getQueryResult

-- | Specialization of 'F.mapM_' for convenient mapping over query results.
{-# INLINABLE mapDB_ #-}
mapDB_ :: (MonadDB m, FromRow row) => (row -> m t) -> m ()
mapDB_ :: (row -> m t) -> m ()
mapDB_ row -> m t
f = m ()
-> (QueryResult row -> m ()) -> Maybe (QueryResult row) -> m ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> m ()
forall (m :: * -> *) a. Monad m => a -> m a
return ()) ((row -> m t) -> QueryResult row -> m ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
F.mapM_ row -> m t
f) (Maybe (QueryResult row) -> m ())
-> m (Maybe (QueryResult row)) -> m ()
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< m (Maybe (QueryResult row))
forall (m :: * -> *) row.
(MonadDB m, FromRow row) =>
m (Maybe (QueryResult row))
getQueryResult

----------------------------------------

-- | Specialization of 'foldrDB' that fetches a list of rows.
{-# INLINABLE fetchMany #-}
fetchMany :: (MonadDB m, FromRow row) => (row -> t) -> m [t]
fetchMany :: (row -> t) -> m [t]
fetchMany row -> t
f = (row -> [t] -> m [t]) -> [t] -> m [t]
forall (m :: * -> *) row acc.
(MonadDB m, FromRow row) =>
(row -> acc -> m acc) -> acc -> m acc
foldrDB (\row
row [t]
acc -> [t] -> m [t]
forall (m :: * -> *) a. Monad m => a -> m a
return ([t] -> m [t]) -> [t] -> m [t]
forall a b. (a -> b) -> a -> b
$ row -> t
f row
row t -> [t] -> [t]
forall a. a -> [a] -> [a]
: [t]
acc) []

-- | Specialization of 'foldlDB' that fetches one or zero rows. If
-- more rows are delivered, 'AffectedRowsMismatch' exception is thrown.
{-# INLINABLE fetchMaybe #-}
fetchMaybe :: (MonadDB m, MonadThrow m, FromRow row) => (row -> t) -> m (Maybe t)
fetchMaybe :: (row -> t) -> m (Maybe t)
fetchMaybe row -> t
f = m (Maybe (QueryResult row))
forall (m :: * -> *) row.
(MonadDB m, FromRow row) =>
m (Maybe (QueryResult row))
getQueryResult m (Maybe (QueryResult row))
-> (Maybe (QueryResult row) -> m (Maybe t)) -> m (Maybe t)
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Maybe (QueryResult row)
mqr -> case Maybe (QueryResult row)
mqr of
  Maybe (QueryResult row)
Nothing -> Maybe t -> m (Maybe t)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe t
forall a. Maybe a
Nothing
  Just QueryResult row
qr -> (Maybe t, QueryResult t) -> Maybe t
forall a b. (a, b) -> a
fst ((Maybe t, QueryResult t) -> Maybe t)
-> m (Maybe t, QueryResult t) -> m (Maybe t)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((Maybe t, QueryResult t) -> row -> m (Maybe t, QueryResult t))
-> (Maybe t, QueryResult t) -> m (Maybe t, QueryResult t)
forall (m :: * -> *) row acc.
(MonadDB m, FromRow row) =>
(acc -> row -> m acc) -> acc -> m acc
foldlDB (Maybe t, QueryResult t) -> row -> m (Maybe t, QueryResult t)
go (Maybe t
forall a. Maybe a
Nothing, row -> t
f (row -> t) -> QueryResult row -> QueryResult t
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QueryResult row
qr)
  where
    go :: (Maybe t, QueryResult t) -> row -> m (Maybe t, QueryResult t)
go (Maybe t
Nothing, QueryResult t
qr) row
row = (Maybe t, QueryResult t) -> m (Maybe t, QueryResult t)
forall (m :: * -> *) a. Monad m => a -> m a
return (t -> Maybe t
forall a. a -> Maybe a
Just (t -> Maybe t) -> t -> Maybe t
forall a b. (a -> b) -> a -> b
$ row -> t
f row
row, QueryResult t
qr)
    go (Just t
_, QueryResult t
qr) row
_ = AffectedRowsMismatch -> m (Maybe t, QueryResult t)
forall e (m :: * -> *) a.
(Exception e, MonadDB m, MonadThrow m) =>
e -> m a
throwDB AffectedRowsMismatch :: [(Int, Int)] -> Int -> AffectedRowsMismatch
AffectedRowsMismatch {
        rowsExpected :: [(Int, Int)]
rowsExpected  = [(Int
0, Int
1)]
      , rowsDelivered :: Int
rowsDelivered = QueryResult t -> Int
forall t. QueryResult t -> Int
ntuples QueryResult t
qr
      }

-- | Specialization of 'fetchMaybe' that fetches exactly one row. If
-- no row is delivered, 'AffectedRowsMismatch' exception is thrown.
{-# INLINABLE fetchOne #-}
fetchOne :: (MonadDB m, MonadThrow m, FromRow row) => (row -> t) -> m t
fetchOne :: (row -> t) -> m t
fetchOne row -> t
f = do
  Maybe t
mt <- (row -> t) -> m (Maybe t)
forall (m :: * -> *) row t.
(MonadDB m, MonadThrow m, FromRow row) =>
(row -> t) -> m (Maybe t)
fetchMaybe row -> t
f
  case Maybe t
mt of
    Just t
t  -> t -> m t
forall (m :: * -> *) a. Monad m => a -> m a
return t
t
    Maybe t
Nothing -> AffectedRowsMismatch -> m t
forall e (m :: * -> *) a.
(Exception e, MonadDB m, MonadThrow m) =>
e -> m a
throwDB AffectedRowsMismatch :: [(Int, Int)] -> Int -> AffectedRowsMismatch
AffectedRowsMismatch {
      rowsExpected :: [(Int, Int)]
rowsExpected = [(Int
1, Int
1)]
    , rowsDelivered :: Int
rowsDelivered = Int
0
    }