module Hasql.Comms.Recv
  ( Recv,
    singleResult,
    allResults,
    toHandler,
    Error (..),
  )
where

import Hasql.Comms.ResultDecoder qualified as ResultDecoder
import Hasql.Platform.Prelude
import Hasql.Pq qualified as Pq

newtype Recv context a
  = Recv (Pq.Connection -> IO (Either (Error context) a))
  deriving stock ((forall a b. (a -> b) -> Recv context a -> Recv context b)
-> (forall a b. a -> Recv context b -> Recv context a)
-> Functor (Recv context)
forall a b. a -> Recv context b -> Recv context a
forall a b. (a -> b) -> Recv context a -> Recv context b
forall context a b. a -> Recv context b -> Recv context a
forall context a b. (a -> b) -> Recv context a -> Recv context b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall context a b. (a -> b) -> Recv context a -> Recv context b
fmap :: forall a b. (a -> b) -> Recv context a -> Recv context b
$c<$ :: forall context a b. a -> Recv context b -> Recv context a
<$ :: forall a b. a -> Recv context b -> Recv context a
Functor)

instance Applicative (Recv context) where
  {-# INLINE pure #-}
  pure :: forall a. a -> Recv context a
pure a
x = (Connection -> IO (Either (Error context) a)) -> Recv context a
forall context a.
(Connection -> IO (Either (Error context) a)) -> Recv context a
Recv \Connection
_ -> Either (Error context) a -> IO (Either (Error context) a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a -> Either (Error context) a
forall a b. b -> Either a b
Right a
x)
  {-# INLINE (<*>) #-}
  Recv Connection -> IO (Either (Error context) (a -> b))
recv1 <*> :: forall a b.
Recv context (a -> b) -> Recv context a -> Recv context b
<*> Recv Connection -> IO (Either (Error context) a)
recv2 =
    (Connection -> IO (Either (Error context) b)) -> Recv context b
forall context a.
(Connection -> IO (Either (Error context) a)) -> Recv context a
Recv \Connection
cs -> do
      Either (Error context) (a -> b)
ef <- Connection -> IO (Either (Error context) (a -> b))
recv1 Connection
cs
      Either (Error context) a
eg <- Connection -> IO (Either (Error context) a)
recv2 Connection
cs
      pure (Either (Error context) (a -> b)
ef Either (Error context) (a -> b)
-> Either (Error context) a -> Either (Error context) b
forall a b.
Either (Error context) (a -> b)
-> Either (Error context) a -> Either (Error context) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Either (Error context) a
eg)

instance Bifunctor Recv where
  {-# INLINE bimap #-}
  bimap :: forall a b c d. (a -> b) -> (c -> d) -> Recv a c -> Recv b d
bimap a -> b
f c -> d
g (Recv Connection -> IO (Either (Error a) c)
recv) = (Connection -> IO (Either (Error b) d)) -> Recv b d
forall context a.
(Connection -> IO (Either (Error context) a)) -> Recv context a
Recv ((Either (Error a) c -> Either (Error b) d)
-> IO (Either (Error a) c) -> IO (Either (Error b) d)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((Error a -> Error b)
-> (c -> d) -> Either (Error a) c -> Either (Error b) d
forall a b c d. (a -> b) -> (c -> d) -> Either a c -> Either b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap ((a -> b) -> Error a -> Error b
forall a b. (a -> b) -> Error a -> Error b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f) c -> d
g) (IO (Either (Error a) c) -> IO (Either (Error b) d))
-> (Connection -> IO (Either (Error a) c))
-> Connection
-> IO (Either (Error b) d)
forall b c a. (b -> c) -> (a -> b) -> a -> c
forall {k} (cat :: k -> k -> *) (b :: k) (c :: k) (a :: k).
Category cat =>
cat b c -> cat a b -> cat a c
. Connection -> IO (Either (Error a) c)
recv)

toHandler :: Recv context a -> Pq.Connection -> IO (Either (Error context) a)
toHandler :: forall context a.
Recv context a -> Connection -> IO (Either (Error context) a)
toHandler (Recv Connection -> IO (Either (Error context) a)
recv) = Connection -> IO (Either (Error context) a)
recv

-- | Exactly one result.
singleResult :: context -> ResultDecoder.ResultDecoder a -> Recv context a
singleResult :: forall context a. context -> ResultDecoder a -> Recv context a
singleResult context
context ResultDecoder a
handler = (Connection -> IO (Either (Error context) a)) -> Recv context a
forall context a.
(Connection -> IO (Either (Error context) a)) -> Recv context a
Recv \Connection
connection -> ExceptT (Error context) IO a -> IO (Either (Error context) a)
forall e (m :: * -> *) a. ExceptT e m a -> m (Either e a)
runExceptT do
  Result
result <- IO (Either (Error context) Result)
-> ExceptT (Error context) IO Result
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT do
    Maybe Result
result <- Connection -> IO (Maybe Result)
Pq.getResult Connection
connection
    case Maybe Result
result of
      Maybe Result
Nothing -> do
        Maybe ByteString
errorMessage <- Connection -> IO (Maybe ByteString)
Pq.errorMessage Connection
connection
        pure (Error context -> Either (Error context) Result
forall a b. a -> Either a b
Left (context -> Maybe ByteString -> Error context
forall context. context -> Maybe ByteString -> Error context
NoResultsError context
context Maybe ByteString
errorMessage))
      Just Result
result -> Either (Error context) Result -> IO (Either (Error context) Result)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Result -> Either (Error context) Result
forall a b. b -> Either a b
Right Result
result)
  IO (Either (Error context) (Maybe Result))
-> ExceptT (Error context) IO (Maybe Result)
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT do
    Maybe Result
result <- Connection -> IO (Maybe Result)
Pq.getResult Connection
connection
    case Maybe Result
result of
      Maybe Result
Nothing -> Either (Error context) (Maybe Result)
-> IO (Either (Error context) (Maybe Result))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Result -> Either (Error context) (Maybe Result)
forall a b. b -> Either a b
Right Maybe Result
result)
      Just Result
_ -> Either (Error context) (Maybe Result)
-> IO (Either (Error context) (Maybe Result))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Error context -> Either (Error context) (Maybe Result)
forall a b. a -> Either a b
Left (context -> Int -> Error context
forall context. context -> Int -> Error context
TooManyResultsError context
context Int
1))
  a
result <- IO (Either (Error context) a) -> ExceptT (Error context) IO a
forall e (m :: * -> *) a. m (Either e a) -> ExceptT e m a
ExceptT do
    Either Error a
result <- ResultDecoder a -> Handler a
forall a. ResultDecoder a -> Handler a
ResultDecoder.toHandler ResultDecoder a
handler Result
result
    pure ((Error -> Error context)
-> Either Error a -> Either (Error context) a
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (context -> Int -> Error -> Error context
forall context. context -> Int -> Error -> Error context
ResultError context
context Int
0) Either Error a
result)
  pure a
result

-- | Consume all results from a multi-statement query (e.g., scripts).
-- Each result is decoded using the provided handler.
-- This is useful for scripts that may contain multiple statements,
-- where each statement produces a result that needs to be validated.
-- All results are consumed even if an error occurs, to leave the connection
-- in a clean state.
allResults :: context -> ResultDecoder.ResultDecoder a -> Recv context ()
allResults :: forall context a. context -> ResultDecoder a -> Recv context ()
allResults context
context ResultDecoder a
handler = (Connection -> IO (Either (Error context) ())) -> Recv context ()
forall context a.
(Connection -> IO (Either (Error context) a)) -> Recv context a
Recv \Connection
connection -> do
  let loop :: Int -> Maybe (Error context) -> IO (Maybe (Error context))
loop Int
resultIndex Maybe (Error context)
maybeError = do
        Maybe Result
result <- Connection -> IO (Maybe Result)
Pq.getResult Connection
connection
        case Maybe Result
result of
          Maybe Result
Nothing -> Maybe (Error context) -> IO (Maybe (Error context))
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe (Error context)
maybeError
          Just Result
result -> do
            Either Error a
decodedResult <- ResultDecoder a -> Handler a
forall a. ResultDecoder a -> Handler a
ResultDecoder.toHandler ResultDecoder a
handler Result
result
            case Either Error a
decodedResult of
              Left Error
err ->
                -- Continue consuming results even after error to clean up connection
                Int -> Maybe (Error context) -> IO (Maybe (Error context))
loop (Int
resultIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) (Error context -> Maybe (Error context)
forall a. a -> Maybe a
Just (context -> Int -> Error -> Error context
forall context. context -> Int -> Error -> Error context
ResultError context
context Int
resultIndex Error
err))
              Right a
_ ->
                Int -> Maybe (Error context) -> IO (Maybe (Error context))
loop (Int
resultIndex Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) Maybe (Error context)
maybeError
  Maybe (Error context)
errorOrUnit <- Int -> Maybe (Error context) -> IO (Maybe (Error context))
loop Int
0 Maybe (Error context)
forall a. Maybe a
Nothing
  pure (Either (Error context) ()
-> (Error context -> Either (Error context) ())
-> Maybe (Error context)
-> Either (Error context) ()
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (() -> Either (Error context) ()
forall a b. b -> Either a b
Right ()) Error context -> Either (Error context) ()
forall a b. a -> Either a b
Left Maybe (Error context)
errorOrUnit)

-- * Errors

data Error context
  = ResultError
      context
      -- | Offset of the result in the series.
      Int
      -- | Underlying error.
      ResultDecoder.Error
  | NoResultsError
      context
      -- | Details about the error. Possibly empty.
      (Maybe ByteString)
  | TooManyResultsError
      context
      -- | Expected count.
      Int
  deriving stock (Int -> Error context -> ShowS
[Error context] -> ShowS
Error context -> String
(Int -> Error context -> ShowS)
-> (Error context -> String)
-> ([Error context] -> ShowS)
-> Show (Error context)
forall context. Show context => Int -> Error context -> ShowS
forall context. Show context => [Error context] -> ShowS
forall context. Show context => Error context -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall context. Show context => Int -> Error context -> ShowS
showsPrec :: Int -> Error context -> ShowS
$cshow :: forall context. Show context => Error context -> String
show :: Error context -> String
$cshowList :: forall context. Show context => [Error context] -> ShowS
showList :: [Error context] -> ShowS
Show, Error context -> Error context -> Bool
(Error context -> Error context -> Bool)
-> (Error context -> Error context -> Bool) -> Eq (Error context)
forall context.
Eq context =>
Error context -> Error context -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall context.
Eq context =>
Error context -> Error context -> Bool
== :: Error context -> Error context -> Bool
$c/= :: forall context.
Eq context =>
Error context -> Error context -> Bool
/= :: Error context -> Error context -> Bool
Eq, (forall a b. (a -> b) -> Error a -> Error b)
-> (forall a b. a -> Error b -> Error a) -> Functor Error
forall a b. a -> Error b -> Error a
forall a b. (a -> b) -> Error a -> Error b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Error a -> Error b
fmap :: forall a b. (a -> b) -> Error a -> Error b
$c<$ :: forall a b. a -> Error b -> Error a
<$ :: forall a b. a -> Error b -> Error a
Functor)