-- | Reusable scenario fragments shared by the per-operation spec modules.
--
-- These are the recurring building blocks — running a query and observing its
-- result, draining the asynchronous result stream, collecting a @COPY OUT@
-- stream, driving a polling loop to its terminal status, and the handful of
-- well-known type OIDs — factored out so each operation module stays focused on
-- the one operation it covers.
module Pqi.Conformance.Scenario
  ( -- * Running queries
    execScenario,
    execAllScenario,
    observed,

    -- * Asynchronous result collection
    drainResults,
    takeResult,
    takeCommandResults,

    -- * Copy and polling loops
    collectCopyOut,
    pollUntilDone,
    flushUntilDone,

    -- * Transactions
    inTransaction,

    -- * Well-known type OIDs
    boolOid,
    byteaOid,
    int2Oid,
    int4Oid,
    int8Oid,
    textOid,
    float8Oid,
  )
where

import Pqi (CopyOutResult (..), FlushStatus (..), PollingStatus (..))
import qualified Pqi
import Pqi.Conformance.Observation
import Pqi.Conformance.Prelude

-- | Run 'Pqi.exec' and observe its result (if any).
execScenario :: ByteString -> Pqi.Connection -> IO (Maybe ResultObservation)
execScenario :: ByteString -> Connection -> IO (Maybe ResultObservation)
execScenario ByteString
sql Connection
connection = Connection -> ByteString -> IO (Maybe Result)
Pqi.exec Connection
connection ByteString
sql IO (Maybe Result)
-> (Maybe Result -> IO (Maybe ResultObservation))
-> IO (Maybe ResultObservation)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Result -> IO ResultObservation)
-> Maybe Result -> IO (Maybe ResultObservation)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Result -> IO ResultObservation
observeResult

-- | Run a sequence of statements with 'Pqi.exec', observing every result.
execAllScenario :: [ByteString] -> Pqi.Connection -> IO [Maybe ResultObservation]
execAllScenario :: [ByteString] -> Connection -> IO [Maybe ResultObservation]
execAllScenario [ByteString]
sqls Connection
connection = (ByteString -> IO (Maybe ResultObservation))
-> [ByteString] -> IO [Maybe ResultObservation]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (ByteString -> Connection -> IO (Maybe ResultObservation)
`execScenario` Connection
connection) [ByteString]
sqls

-- | Run 'Pqi.execParams' and observe its result (if any).
observed ::
  ByteString ->
  [Maybe (Word32, ByteString, Pqi.Format)] ->
  Pqi.Format ->
  Pqi.Connection ->
  IO (Maybe ResultObservation)
observed :: ByteString
-> [Maybe (Word32, ByteString, Format)]
-> Format
-> Connection
-> IO (Maybe ResultObservation)
observed ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat Connection
connection =
  Connection
-> ByteString
-> [Maybe (Word32, ByteString, Format)]
-> Format
-> IO (Maybe Result)
Pqi.execParams Connection
connection ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat IO (Maybe Result)
-> (Maybe Result -> IO (Maybe ResultObservation))
-> IO (Maybe ResultObservation)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Result -> IO ResultObservation)
-> Maybe Result -> IO (Maybe ResultObservation)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Result -> IO ResultObservation
observeResult

-- | Collect and observe results with 'Pqi.getResult' until it reports
-- completion with 'Nothing'.
drainResults :: Pqi.Connection -> IO [ResultObservation]
drainResults :: Connection -> IO [ResultObservation]
drainResults Connection
connection = [ResultObservation] -> IO [ResultObservation]
go []
  where
    go :: [ResultObservation] -> IO [ResultObservation]
go [ResultObservation]
acc =
      Connection -> IO (Maybe Result)
Pqi.getResult Connection
connection IO (Maybe Result)
-> (Maybe Result -> IO [ResultObservation])
-> IO [ResultObservation]
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe Result
Nothing -> [ResultObservation] -> IO [ResultObservation]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([ResultObservation] -> [ResultObservation]
forall a. [a] -> [a]
reverse [ResultObservation]
acc)
        Just Result
result -> do
          observation <- Result -> IO ResultObservation
observeResult Result
result
          go (observation : acc)

-- | One 'Pqi.getResult' step, observed.
takeResult :: Pqi.Connection -> IO (Maybe ResultObservation)
takeResult :: Connection -> IO (Maybe ResultObservation)
takeResult Connection
connection = Connection -> IO (Maybe Result)
Pqi.getResult Connection
connection IO (Maybe Result)
-> (Maybe Result -> IO (Maybe ResultObservation))
-> IO (Maybe ResultObservation)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Result -> IO ResultObservation)
-> Maybe Result -> IO (Maybe ResultObservation)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Result -> IO ResultObservation
observeResult

-- | The results of one pipelined command: its result and the 'Nothing'
-- separator that ends it.
takeCommandResults ::
  Pqi.Connection -> IO (Maybe ResultObservation, Maybe ResultObservation)
takeCommandResults :: Connection -> IO (Maybe ResultObservation, Maybe ResultObservation)
takeCommandResults Connection
connection = do
  result <- Connection -> IO (Maybe ResultObservation)
takeResult Connection
connection
  separator <- takeResult connection
  pure (result, separator)

-- | Collect blocking 'Pqi.getCopyData' outcomes until the stream reports
-- anything other than a row (normally 'CopyOutDone'), including that final
-- outcome.
collectCopyOut :: Pqi.Connection -> IO [CopyOutResult]
collectCopyOut :: Connection -> IO [CopyOutResult]
collectCopyOut Connection
connection = Int -> [CopyOutResult] -> IO [CopyOutResult]
go (Int
1000 :: Int) []
  where
    go :: Int -> [CopyOutResult] -> IO [CopyOutResult]
go Int
0 [CopyOutResult]
acc = [CopyOutResult] -> IO [CopyOutResult]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([CopyOutResult] -> [CopyOutResult]
forall a. [a] -> [a]
reverse [CopyOutResult]
acc)
    go Int
n [CopyOutResult]
acc =
      Connection -> Bool -> IO CopyOutResult
Pqi.getCopyData Connection
connection Bool
False IO CopyOutResult
-> (CopyOutResult -> IO [CopyOutResult]) -> IO [CopyOutResult]
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        CopyOutRow ByteString
row -> Int -> [CopyOutResult] -> IO [CopyOutResult]
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (ByteString -> CopyOutResult
CopyOutRow ByteString
row CopyOutResult -> [CopyOutResult] -> [CopyOutResult]
forall a. a -> [a] -> [a]
: [CopyOutResult]
acc)
        CopyOutResult
terminal -> [CopyOutResult] -> IO [CopyOutResult]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([CopyOutResult] -> [CopyOutResult]
forall a. [a] -> [a]
reverse (CopyOutResult
terminal CopyOutResult -> [CopyOutResult] -> [CopyOutResult]
forall a. a -> [a] -> [a]
: [CopyOutResult]
acc))

-- | Drive a polling loop to its terminal status, spinning with a small delay
-- instead of waiting on the socket. Bails out as failed after ten seconds.
pollUntilDone :: IO PollingStatus -> IO PollingStatus
pollUntilDone :: IO PollingStatus -> IO PollingStatus
pollUntilDone IO PollingStatus
poll = Int -> IO PollingStatus
go (Int
10000 :: Int)
  where
    go :: Int -> IO PollingStatus
go Int
0 = PollingStatus -> IO PollingStatus
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PollingStatus
PollingFailed
    go Int
n =
      IO PollingStatus
poll IO PollingStatus
-> (PollingStatus -> IO PollingStatus) -> IO PollingStatus
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        PollingStatus
PollingReading -> Int -> IO ()
threadDelay Int
1000 IO () -> IO PollingStatus -> IO PollingStatus
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO PollingStatus
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        PollingStatus
PollingWriting -> Int -> IO ()
threadDelay Int
1000 IO () -> IO PollingStatus -> IO PollingStatus
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO PollingStatus
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        PollingStatus
terminal -> PollingStatus -> IO PollingStatus
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PollingStatus
terminal

-- | Drive 'Pqi.flush' to a terminal status, spinning while it reports
-- 'FlushWriting'. Bails out as failed after ten seconds.
flushUntilDone :: IO FlushStatus -> IO FlushStatus
flushUntilDone :: IO FlushStatus -> IO FlushStatus
flushUntilDone IO FlushStatus
doFlush = Int -> IO FlushStatus
go (Int
10000 :: Int)
  where
    go :: Int -> IO FlushStatus
go Int
0 = FlushStatus -> IO FlushStatus
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FlushStatus
FlushFailed
    go Int
n =
      IO FlushStatus
doFlush IO FlushStatus -> (FlushStatus -> IO FlushStatus) -> IO FlushStatus
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        FlushStatus
FlushWriting -> Int -> IO ()
threadDelay Int
1000 IO () -> IO FlushStatus -> IO FlushStatus
forall a b. IO a -> IO b -> IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Int -> IO FlushStatus
go (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
        FlushStatus
terminal -> FlushStatus -> IO FlushStatus
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure FlushStatus
terminal

-- | Run an action between @begin@ and @commit@. Large-object operations must
-- run inside a transaction block.
inTransaction :: Pqi.Connection -> IO a -> IO a
inTransaction :: forall a. Connection -> IO a -> IO a
inTransaction Connection
connection IO a
action = do
  _ <- Connection -> ByteString -> IO (Maybe Result)
Pqi.exec Connection
connection ByteString
"begin"
  result <- action
  _ <- Pqi.exec connection "commit"
  pure result

boolOid :: Word32
boolOid :: Word32
boolOid = Word32
16

byteaOid :: Word32
byteaOid :: Word32
byteaOid = Word32
17

int2Oid :: Word32
int2Oid :: Word32
int2Oid = Word32
21

int4Oid :: Word32
int4Oid :: Word32
int4Oid = Word32
23

int8Oid :: Word32
int8Oid :: Word32
int8Oid = Word32
20

textOid :: Word32
textOid :: Word32
textOid = Word32
25

float8Oid :: Word32
float8Oid :: Word32
float8Oid = Word32
701