-- | Command execution: the simple- and extended-query flows, and the
-- materialization of the backend message stream into a 'NativeResult'.
module Pqi.Native.Query
  ( exec,
    execParams,
    prepare,
    execPrepared,
    describePrepared,
    describePortal,
    sendQuery,
    sendQueryParams,
    sendPrepare,
    sendQueryPrepared,
    sendDescribePrepared,
    sendDescribePortal,
    getNextResult,
  )
where

import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import Pqi (ConnStatus (..), ExecStatus (..), Format (..), PipelineStatus (..))
import Pqi.Native.Connection
import Pqi.Native.Prelude
import Pqi.Native.Transport.Message
import Pqi.Native.Types (NativeResult (..), formatResultError)
import qualified PtrPoker.Write as Poker

-- * Message construction

-- Sync-inclusive variants used by the synchronous exec* functions.

paramsWrite :: ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> Poker.Write
paramsWrite :: ByteString
-> [Maybe (Word32, ByteString, Format)] -> Format -> Write
paramsWrite ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat =
  ByteString
-> [Maybe (Word32, ByteString, Format)] -> Format -> Write
asyncParamsWrite ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage

preparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write
preparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Write
preparedWrite ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat =
  ByteString -> [Maybe (ByteString, Format)] -> Format -> Write
asyncPreparedWrite ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage

prepareWrite :: ByteString -> ByteString -> Maybe [Word32] -> Poker.Write
prepareWrite :: ByteString -> ByteString -> Maybe [Word32] -> Write
prepareWrite ByteString
name ByteString
sql Maybe [Word32]
parameterTypes =
  ByteString -> ByteString -> [Word32] -> Write
parseMessage ByteString
name ByteString
sql ([Word32] -> Maybe [Word32] -> [Word32]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Word32]
parameterTypes) Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage

-- Sync-free variants used by the async send* functions.
-- In non-pipeline mode sendAsync appends syncMessage; in pipeline mode it does not.

asyncParamsWrite :: ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> Poker.Write
asyncParamsWrite :: ByteString
-> [Maybe (Word32, ByteString, Format)] -> Format -> Write
asyncParamsWrite ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat =
  ByteString -> ByteString -> [Word32] -> Write
parseMessage ByteString
"" ByteString
sql ((Maybe (Word32, ByteString, Format) -> Word32)
-> [Maybe (Word32, ByteString, Format)] -> [Word32]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe (Word32, ByteString, Format) -> Word32
paramOid [Maybe (Word32, ByteString, Format)]
params)
    Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString
-> ByteString -> [Int16] -> [Maybe ByteString] -> [Int16] -> Write
bindMessage ByteString
"" ByteString
"" ((Maybe (Word32, ByteString, Format) -> Int16)
-> [Maybe (Word32, ByteString, Format)] -> [Int16]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe (Word32, ByteString, Format) -> Int16
paramFormat [Maybe (Word32, ByteString, Format)]
params) ((Maybe (Word32, ByteString, Format) -> Maybe ByteString)
-> [Maybe (Word32, ByteString, Format)] -> [Maybe ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe (Word32, ByteString, Format) -> Maybe ByteString
paramValue [Maybe (Word32, ByteString, Format)]
params) [Format -> Int16
formatCodeOf Format
resultFormat]
    Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString -> Write
describePortalMessage ByteString
""
    Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString -> Int32 -> Write
executeMessage ByteString
"" Int32
0

asyncPreparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write
asyncPreparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Write
asyncPreparedWrite ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat =
  ByteString
-> ByteString -> [Int16] -> [Maybe ByteString] -> [Int16] -> Write
bindMessage ByteString
"" ByteString
name ((Maybe (ByteString, Format) -> Int16)
-> [Maybe (ByteString, Format)] -> [Int16]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe (ByteString, Format) -> Int16
boundFormat [Maybe (ByteString, Format)]
params) ((Maybe (ByteString, Format) -> Maybe ByteString)
-> [Maybe (ByteString, Format)] -> [Maybe ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Maybe (ByteString, Format) -> Maybe ByteString
boundValue [Maybe (ByteString, Format)]
params) [Format -> Int16
formatCodeOf Format
resultFormat]
    Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString -> Write
describePortalMessage ByteString
""
    Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> ByteString -> Int32 -> Write
executeMessage ByteString
"" Int32
0

-- * Synchronous flows

-- | Simple query. Returns the last result, mirroring @PQexec@.
exec :: Connection -> ByteString -> IO (Maybe NativeResult)
exec :: Connection -> ByteString -> IO (Maybe NativeResult)
exec Connection
connection ByteString
sql = Connection -> IO (Maybe NativeResult) -> IO (Maybe NativeResult)
forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection do
  Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> Write
queryMessage ByteString
sql)
  [NativeResult] -> Maybe NativeResult
forall a. [a] -> Maybe a
lastMaybe ([NativeResult] -> Maybe NativeResult)
-> IO [NativeResult] -> IO (Maybe NativeResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> ByteString -> IO [NativeResult]
collectSimple Connection
connection ByteString
sql

-- | Parameterized query via the extended protocol.
execParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO (Maybe NativeResult)
execParams :: Connection
-> ByteString
-> [Maybe (Word32, ByteString, Format)]
-> Format
-> IO (Maybe NativeResult)
execParams Connection
connection ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat = Connection -> IO (Maybe NativeResult) -> IO (Maybe NativeResult)
forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection do
  Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString
-> [Maybe (Word32, ByteString, Format)] -> Format -> Write
paramsWrite ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat)
  NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (NativeResult -> Maybe NativeResult)
-> IO NativeResult -> IO (Maybe NativeResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> ByteString -> IO NativeResult
collectExtended Connection
connection ByteString
sql

-- | Prepare a named statement.
prepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO (Maybe NativeResult)
prepare :: Connection
-> ByteString
-> ByteString
-> Maybe [Word32]
-> IO (Maybe NativeResult)
prepare Connection
connection ByteString
name ByteString
sql Maybe [Word32]
parameterTypes = Connection -> IO (Maybe NativeResult) -> IO (Maybe NativeResult)
forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection do
  Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> ByteString -> Maybe [Word32] -> Write
prepareWrite ByteString
name ByteString
sql Maybe [Word32]
parameterTypes)
  NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (NativeResult -> Maybe NativeResult)
-> IO NativeResult -> IO (Maybe NativeResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> ByteString -> IO NativeResult
collectExtended Connection
connection ByteString
sql

-- | Execute a previously prepared statement.
execPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO (Maybe NativeResult)
execPrepared :: Connection
-> ByteString
-> [Maybe (ByteString, Format)]
-> Format
-> IO (Maybe NativeResult)
execPrepared Connection
connection ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat = Connection -> IO (Maybe NativeResult) -> IO (Maybe NativeResult)
forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection do
  Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> [Maybe (ByteString, Format)] -> Format -> Write
preparedWrite ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat)
  NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (NativeResult -> Maybe NativeResult)
-> IO NativeResult -> IO (Maybe NativeResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> ByteString -> IO NativeResult
collectExtended Connection
connection ByteString
""

-- * Asynchronous flows

-- | Send a write in async mode, tracking pending commands for pipeline abort.
sendAsync :: Connection -> ByteString -> Poker.Write -> IO Bool
sendAsync :: Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
sql Write
write = do
  ConnStatus
status <- IORef ConnStatus -> IO ConnStatus
forall a. IORef a -> IO a
readIORef (Connection -> IORef ConnStatus
connStatus Connection
connection)
  case ConnStatus
status of
    ConnStatus
ConnectionOk -> do
      Connection -> Write -> IO ()
sendMessage Connection
connection Write
write
      IORef ByteString -> ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef ByteString
currentQuery Connection
connection) ByteString
sql
      IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
asyncPending Connection
connection) Bool
True
      PipelineStatus
pipeStatus <- IORef PipelineStatus -> IO PipelineStatus
forall a. IORef a -> IO a
readIORef (Connection -> IORef PipelineStatus
pipelineStatus Connection
connection)
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (PipelineStatus
pipeStatus PipelineStatus -> PipelineStatus -> Bool
forall a. Eq a => a -> a -> Bool
/= PipelineStatus
PipelineOff) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef Int
pendingCommands Connection
connection) (Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
      Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
    ConnStatus
_ -> Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False

-- | Whether the connection is in pipeline mode.
inPipeline :: Connection -> IO Bool
inPipeline :: Connection -> IO Bool
inPipeline Connection
connection = (PipelineStatus -> PipelineStatus -> Bool
forall a. Eq a => a -> a -> Bool
/= PipelineStatus
PipelineOff) (PipelineStatus -> Bool) -> IO PipelineStatus -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IORef PipelineStatus -> IO PipelineStatus
forall a. IORef a -> IO a
readIORef (Connection -> IORef PipelineStatus
pipelineStatus Connection
connection)

-- Simple query protocol: no Sync needed (server sends ReadyForQuery on its own).
sendQuery :: Connection -> ByteString -> IO Bool
sendQuery :: Connection -> ByteString -> IO Bool
sendQuery Connection
connection ByteString
sql = Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
sql (ByteString -> Write
queryMessage ByteString
sql)

-- Extended query: include Sync when not in pipeline mode; omit Sync in
-- pipeline mode (the caller drives sync boundaries via 'pipelineSync').
sendQueryParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO Bool
sendQueryParams :: Connection
-> ByteString
-> [Maybe (Word32, ByteString, Format)]
-> Format
-> IO Bool
sendQueryParams Connection
connection ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat = do
  Bool
pipeline <- Connection -> IO Bool
inPipeline Connection
connection
  Bool
ok <-
    Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
sql
      (Write -> IO Bool) -> Write -> IO Bool
forall a b. (a -> b) -> a -> b
$ if Bool
pipeline
        then ByteString
-> [Maybe (Word32, ByteString, Format)] -> Format -> Write
asyncParamsWrite ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat
        else ByteString
-> [Maybe (Word32, ByteString, Format)] -> Format -> Write
paramsWrite ByteString
sql [Maybe (Word32, ByteString, Format)]
params Format
resultFormat
  -- 'sendQueryParams' drives the extended protocol and so always sends an
  -- unnamed @Parse@, whose @ParseComplete@ must fold into this command's own
  -- result (like 'sendQueryPrepared'). Record its origin so it is not mistaken
  -- for the terminal @ParseComplete@ of a pipelined 'sendPrepare'.
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
ok Bool -> Bool -> Bool
&& Bool
pipeline) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef (Seq Bool) -> (Seq Bool -> Seq Bool) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef (Seq Bool)
pendingParseOrigins Connection
connection) (Seq Bool -> Bool -> Seq Bool
forall a. Seq a -> a -> Seq a
Seq.|> Bool
False)
  Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
ok

sendPrepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO Bool
sendPrepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO Bool
sendPrepare Connection
connection ByteString
name ByteString
sql Maybe [Word32]
parameterTypes = do
  Bool
pipeline <- Connection -> IO Bool
inPipeline Connection
connection
  Bool
ok <-
    Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
sql
      (Write -> IO Bool) -> Write -> IO Bool
forall a b. (a -> b) -> a -> b
$ if Bool
pipeline
        then ByteString -> ByteString -> [Word32] -> Write
parseMessage ByteString
name ByteString
sql ([Word32] -> Maybe [Word32] -> [Word32]
forall a. a -> Maybe a -> a
fromMaybe [] Maybe [Word32]
parameterTypes)
        else ByteString -> ByteString -> Maybe [Word32] -> Write
prepareWrite ByteString
name ByteString
sql Maybe [Word32]
parameterTypes
  Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Bool
ok Bool -> Bool -> Bool
&& Bool
pipeline) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef (Seq Bool) -> (Seq Bool -> Seq Bool) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef (Seq Bool)
pendingParseOrigins Connection
connection) (Seq Bool -> Bool -> Seq Bool
forall a. Seq a -> a -> Seq a
Seq.|> Bool
True)
  Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
ok

sendQueryPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO Bool
sendQueryPrepared :: Connection
-> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO Bool
sendQueryPrepared Connection
connection ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat = do
  Bool
pipeline <- Connection -> IO Bool
inPipeline Connection
connection
  Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
""
    (Write -> IO Bool) -> Write -> IO Bool
forall a b. (a -> b) -> a -> b
$ if Bool
pipeline
      then ByteString -> [Maybe (ByteString, Format)] -> Format -> Write
asyncPreparedWrite ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat
      else ByteString -> [Maybe (ByteString, Format)] -> Format -> Write
preparedWrite ByteString
name [Maybe (ByteString, Format)]
params Format
resultFormat

sendDescribePrepared :: Connection -> ByteString -> IO Bool
sendDescribePrepared :: Connection -> ByteString -> IO Bool
sendDescribePrepared Connection
connection ByteString
name = do
  Bool
pipeline <- Connection -> IO Bool
inPipeline Connection
connection
  Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
""
    (Write -> IO Bool) -> Write -> IO Bool
forall a b. (a -> b) -> a -> b
$ if Bool
pipeline
      then ByteString -> Write
describeStatementMessage ByteString
name
      else ByteString -> Write
describeStatementMessage ByteString
name Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage

sendDescribePortal :: Connection -> ByteString -> IO Bool
sendDescribePortal :: Connection -> ByteString -> IO Bool
sendDescribePortal Connection
connection ByteString
name = do
  Bool
pipeline <- Connection -> IO Bool
inPipeline Connection
connection
  Connection -> ByteString -> Write -> IO Bool
sendAsync Connection
connection ByteString
""
    (Write -> IO Bool) -> Write -> IO Bool
forall a b. (a -> b) -> a -> b
$ if Bool
pipeline
      then ByteString -> Write
describePortalMessage ByteString
name
      else ByteString -> Write
describePortalMessage ByteString
name Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage

-- | Read the next result of an in-flight asynchronous command, or 'Nothing'
-- once @ReadyForQuery@ is reached (clearing the pending flag), mirroring
-- @PQgetResult@.
--
-- In pipeline mode a separator 'Nothing' is returned between each command's
-- result set, and a 'PipelineSync' result is returned for each @Sync@
-- boundary. In single-row mode each data row is delivered as a separate
-- 'SingleTuple' result followed by a final 'TuplesOk' with no rows.
getNextResult :: Connection -> IO (Maybe NativeResult)
getNextResult :: Connection -> IO (Maybe NativeResult)
getNextResult Connection
connection = do
  Bool
pending <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
asyncPending Connection
connection)
  if Bool -> Bool
not Bool
pending
    then Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NativeResult
forall a. Maybe a
Nothing
    else do
      Bool
sepPending <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
pipelineSeparatorPending Connection
connection)
      if Bool
sepPending
        then do
          IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
pipelineSeparatorPending Connection
connection) Bool
False
          Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NativeResult
forall a. Maybe a
Nothing
        else do
          Bool
singleRow <- IORef Bool -> IO Bool
forall a. IORef a -> IO a
readIORef (Connection -> IORef Bool
singleRowMode Connection
connection)
          [FieldDescription]
cachedFields <- IORef [FieldDescription] -> IO [FieldDescription]
forall a. IORef a -> IO a
readIORef (Connection -> IORef [FieldDescription]
singleRowFields Connection
connection)
          let initBuilder :: Builder
initBuilder =
                if Bool
singleRow Bool -> Bool -> Bool
&& Bool -> Bool
not ([FieldDescription] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [FieldDescription]
cachedFields)
                  then Builder
emptyBuilder {accFields = cachedFields, accSawRowDescription = True}
                  else Builder
emptyBuilder
          Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
initBuilder
  where
    -- Decrement the pending-command counter and set the separator flag when in
    -- pipeline mode.  Called when a "terminal" result is about to be returned.
    finishCommand :: PipelineStatus -> IO ()
finishCommand PipelineStatus
pipeStatus = do
      Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (PipelineStatus
pipeStatus PipelineStatus -> PipelineStatus -> Bool
forall a. Eq a => a -> a -> Bool
/= PipelineStatus
PipelineOff) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
        IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef Int
pendingCommands Connection
connection) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
        IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
pipelineSeparatorPending Connection
connection) Bool
True

    go :: Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder = do
      PipelineStatus
pipeStatus <- IORef PipelineStatus -> IO PipelineStatus
forall a. IORef a -> IO a
readIORef (Connection -> IORef PipelineStatus
pipelineStatus Connection
connection)
      -- In aborted pipeline mode, if the server has already sent nothing for
      -- the remaining commands (it discards them after the first error), we
      -- generate synthetic PipelineAbort results for each outstanding command
      -- rather than blocking on a wire read that will never come.
      Int
pending <- IORef Int -> IO Int
forall a. IORef a -> IO a
readIORef (Connection -> IORef Int
pendingCommands Connection
connection)
      if PipelineStatus
pipeStatus PipelineStatus -> PipelineStatus -> Bool
forall a. Eq a => a -> a -> Bool
== PipelineStatus
PipelineAborted Bool -> Bool -> Bool
&& Int
pending Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
        then do
          IORef Int -> (Int -> Int) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' (Connection -> IORef Int
pendingCommands Connection
connection) (Int -> Int -> Int
forall a. Num a => a -> a -> a
subtract Int
1)
          IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
pipelineSeparatorPending Connection
connection) Bool
True
          Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
PipelineAbort [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
        else Bool -> Builder -> PipelineStatus -> IO (Maybe NativeResult)
readAndProcess Bool
singleRow Builder
builder PipelineStatus
pipeStatus

    readAndProcess :: Bool -> Builder -> PipelineStatus -> IO (Maybe NativeResult)
readAndProcess Bool
singleRow Builder
builder PipelineStatus
pipeStatus = do
      BackendMessage
message <- Connection -> IO BackendMessage
nextMessage Connection
connection
      case BackendMessage
message of
        RowDescription [FieldDescription]
fs ->
          Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accFields = fs, accSawRowDescription = True, accHadResponse = True}
        ParameterDescription [Word32]
oids ->
          Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accParamOids = oids, accHadResponse = True}
        BackendMessage
NoData ->
          Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accHadResponse = True}
        DataRow [Maybe ByteString]
values ->
          if Bool
singleRow
            then do
              IORef [FieldDescription] -> [FieldDescription] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef [FieldDescription]
singleRowFields Connection
connection) (Builder -> [FieldDescription]
accFields Builder
builder)
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
SingleTuple (Builder -> [FieldDescription]
accFields Builder
builder) [[Maybe ByteString]
values] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
            else Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accRevRows = values : (accRevRows builder)}
        BackendMessage
ParseComplete -> do
          -- Charge the @ParseComplete@ to the command that produced it by
          -- popping the oldest recorded origin. Only a 'sendPrepare' origin
          -- (@True@) terminates the command as 'CommandOk'; a 'sendQueryParams'
          -- origin (@False@) folds into the accumulating result, and so does a
          -- @ParseComplete@ seen outside pipeline mode (e.g. a non-pipelined
          -- async 'sendQueryParams', whose result is collected by 'collectExtended').
          Maybe Bool
origin <-
            if PipelineStatus
pipeStatus PipelineStatus -> PipelineStatus -> Bool
forall a. Eq a => a -> a -> Bool
/= PipelineStatus
PipelineOff
              then Connection -> IO (Maybe Bool)
popPendingParseOrigin Connection
connection
              else Maybe Bool -> IO (Maybe Bool)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Bool
forall a. Maybe a
Nothing
          case Maybe Bool
origin of
            Just Bool
True -> do
              PipelineStatus -> IO ()
finishCommand PipelineStatus
pipeStatus
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
CommandOk [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
            Maybe Bool
_ -> Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accHadResponse = True}
        BackendMessage
BindComplete ->
          Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accHadResponse = True}
        BackendMessage
CloseComplete ->
          Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder {accHadResponse = True}
        CommandComplete ByteString
tag -> do
          if Bool
singleRow
            then do
              IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
singleRowMode Connection
connection) Bool
False
              IORef [FieldDescription] -> [FieldDescription] -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef [FieldDescription]
singleRowFields Connection
connection) []
              IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
TuplesOk (Builder -> [FieldDescription]
accFields Builder
builder) [] (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
tag) Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
            else do
              IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
              PipelineStatus -> IO ()
finishCommand PipelineStatus
pipeStatus
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (Builder -> Maybe ByteString -> NativeResult
commandResult Builder
builder (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
tag)))
        BackendMessage
EmptyQueryResponse -> do
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
          PipelineStatus -> IO ()
finishCommand PipelineStatus
pipeStatus
          Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
EmptyQuery [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
        ErrorResponse [(Word8, ByteString)]
fs -> do
          let errMap :: Map Word8 ByteString
errMap = [(Word8, ByteString)] -> Map Word8 ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Word8, ByteString)]
fs
          case PipelineStatus
pipeStatus of
            PipelineStatus
PipelineAborted -> do
              -- Should not normally happen (server discards commands in abort
              -- mode) but handle defensively.
              PipelineStatus -> IO ()
finishCommand PipelineStatus
pipeStatus
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
PipelineAbort [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
            PipelineStatus
PipelineOn -> do
              IORef PipelineStatus -> PipelineStatus -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef PipelineStatus
pipelineStatus Connection
connection) PipelineStatus
PipelineAborted
              PipelineStatus -> IO ()
finishCommand PipelineStatus
PipelineOn
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
FatalError [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
errMap [] ByteString
""))
            PipelineStatus
PipelineOff -> do
              ByteString
sql <- IORef ByteString -> IO ByteString
forall a. IORef a -> IO a
readIORef (Connection -> IORef ByteString
currentQuery Connection
connection)
              IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Map Word8 ByteString -> ByteString
formatResultError ByteString
sql Map Word8 ByteString
errMap))
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
FatalError [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
errMap [] ByteString
sql))
        BackendMessage
PortalSuspended ->
          Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (Builder -> Maybe ByteString -> NativeResult
commandResult Builder
builder Maybe ByteString
forall a. Maybe a
Nothing))
        ReadyForQuery Word8
txState -> do
          IORef Word8 -> Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Word8
txStatus Connection
connection) Word8
txState
          case PipelineStatus
pipeStatus of
            PipelineStatus
PipelineOff -> do
              IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
asyncPending Connection
connection) Bool
False
              if (Builder -> Bool
accHadResponse Builder
builder)
                then Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (Builder -> NativeResult
describeResult Builder
builder))
                else Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe NativeResult
forall a. Maybe a
Nothing
            PipelineStatus
_ -> do
              IORef PipelineStatus -> PipelineStatus -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef PipelineStatus
pipelineStatus Connection
connection) PipelineStatus
PipelineOn
              -- A PipelineSync result is its own command boundary: unlike a
              -- normal command result, libpq does not emit a separating NULL
              -- after it, so consecutive syncs are reported back-to-back. We
              -- therefore never set 'pipelineSeparatorPending' here. Only the
              -- final sync clears 'asyncPending'; an earlier one leaves it set
              -- so the next 'getNextResult' reads straight on to the next sync.
              Int
remaining <- IORef Int -> (Int -> (Int, Int)) -> IO Int
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef' (Connection -> IORef Int
pendingSyncs Connection
connection) (\Int
n -> (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1))
              Bool -> IO () -> IO ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when (Int
remaining Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ IORef Bool -> Bool -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Bool
asyncPending Connection
connection) Bool
False
              Maybe NativeResult -> IO (Maybe NativeResult)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
PipelineSync [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
        BackendMessage
_ -> Bool -> Builder -> IO (Maybe NativeResult)
go Bool
singleRow Builder
builder

-- | Pop the origin of the next in-flight @ParseComplete@ in pipeline mode:
-- @True@ if it is the terminal @ParseComplete@ of a 'sendPrepare' (to be
-- materialized as 'CommandOk'), @False@ if it belongs to a 'sendQueryParams'
-- and must fold into that command's accumulating result. @Nothing@ when no
-- @Parse@ is pending (defensive: the @ParseComplete@ is then folded).
popPendingParseOrigin :: Connection -> IO (Maybe Bool)
popPendingParseOrigin :: Connection -> IO (Maybe Bool)
popPendingParseOrigin Connection
connection =
  IORef (Seq Bool)
-> (Seq Bool -> (Seq Bool, Maybe Bool)) -> IO (Maybe Bool)
forall a b. IORef a -> (a -> (a, b)) -> IO b
atomicModifyIORef'
    (Connection -> IORef (Seq Bool)
pendingParseOrigins Connection
connection)
    ( \Seq Bool
queue -> case Seq Bool -> ViewL Bool
forall a. Seq a -> ViewL a
Seq.viewl Seq Bool
queue of
        Bool
origin Seq.:< Seq Bool
rest -> (Seq Bool
rest, Bool -> Maybe Bool
forall a. a -> Maybe a
Just Bool
origin)
        ViewL Bool
Seq.EmptyL -> (Seq Bool
queue, Maybe Bool
forall a. Maybe a
Nothing)
    )

-- | Describe a prepared statement.
describePrepared :: Connection -> ByteString -> IO (Maybe NativeResult)
describePrepared :: Connection -> ByteString -> IO (Maybe NativeResult)
describePrepared Connection
connection ByteString
name = Connection -> IO (Maybe NativeResult) -> IO (Maybe NativeResult)
forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection do
  Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> Write
describeStatementMessage ByteString
name Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage)
  NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (NativeResult -> Maybe NativeResult)
-> IO NativeResult -> IO (Maybe NativeResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> ByteString -> IO NativeResult
collectExtended Connection
connection ByteString
""

-- | Describe a portal.
describePortal :: Connection -> ByteString -> IO (Maybe NativeResult)
describePortal :: Connection -> ByteString -> IO (Maybe NativeResult)
describePortal Connection
connection ByteString
name = Connection -> IO (Maybe NativeResult) -> IO (Maybe NativeResult)
forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection do
  Connection -> Write -> IO ()
sendMessage Connection
connection (ByteString -> Write
describePortalMessage ByteString
name Write -> Write -> Write
forall a. Semigroup a => a -> a -> a
<> Write
syncMessage)
  NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (NativeResult -> Maybe NativeResult)
-> IO NativeResult -> IO (Maybe NativeResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> ByteString -> IO NativeResult
collectExtended Connection
connection ByteString
""

-- * Parameter projections

paramOid :: Maybe (Word32, ByteString, Format) -> Word32
paramOid :: Maybe (Word32, ByteString, Format) -> Word32
paramOid = Word32
-> ((Word32, ByteString, Format) -> Word32)
-> Maybe (Word32, ByteString, Format)
-> Word32
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Word32
0 (\(Word32
oid, ByteString
_, Format
_) -> Word32
oid)

paramFormat :: Maybe (Word32, ByteString, Format) -> Int16
paramFormat :: Maybe (Word32, ByteString, Format) -> Int16
paramFormat = Int16
-> ((Word32, ByteString, Format) -> Int16)
-> Maybe (Word32, ByteString, Format)
-> Int16
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int16
0 (\(Word32
_, ByteString
_, Format
format) -> Format -> Int16
formatCodeOf Format
format)

paramValue :: Maybe (Word32, ByteString, Format) -> Maybe ByteString
paramValue :: Maybe (Word32, ByteString, Format) -> Maybe ByteString
paramValue = ((Word32, ByteString, Format) -> ByteString)
-> Maybe (Word32, ByteString, Format) -> Maybe ByteString
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\(Word32
_, ByteString
value, Format
_) -> ByteString
value)

boundFormat :: Maybe (ByteString, Format) -> Int16
boundFormat :: Maybe (ByteString, Format) -> Int16
boundFormat = Int16
-> ((ByteString, Format) -> Int16)
-> Maybe (ByteString, Format)
-> Int16
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int16
0 (Format -> Int16
formatCodeOf (Format -> Int16)
-> ((ByteString, Format) -> Format)
-> (ByteString, Format)
-> Int16
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString, Format) -> Format
forall a b. (a, b) -> b
snd)

boundValue :: Maybe (ByteString, Format) -> Maybe ByteString
boundValue :: Maybe (ByteString, Format) -> Maybe ByteString
boundValue = ((ByteString, Format) -> ByteString)
-> Maybe (ByteString, Format) -> Maybe ByteString
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (ByteString, Format) -> ByteString
forall a b. (a, b) -> a
fst

-- | Renamed from @formatCode@ to avoid clashing with 'FieldDescription's
-- @formatCode@ field now that 'DuplicateRecordFields' is no longer enabled.
formatCodeOf :: Format -> Int16
formatCodeOf :: Format -> Int16
formatCodeOf = \case
  Format
Text -> Int16
0
  Format
Binary -> Int16
1

-- * Result collection

-- | Only run a flow on a ready connection; mirror libpq returning no result
-- when the connection is not usable.
withReady :: Connection -> IO (Maybe a) -> IO (Maybe a)
withReady :: forall a. Connection -> IO (Maybe a) -> IO (Maybe a)
withReady Connection
connection IO (Maybe a)
action = do
  ConnStatus
status <- IORef ConnStatus -> IO ConnStatus
forall a. IORef a -> IO a
readIORef (Connection -> IORef ConnStatus
connStatus Connection
connection)
  case ConnStatus
status of
    ConnStatus
ConnectionOk -> IO (Maybe a)
action
    ConnStatus
_ -> Maybe a -> IO (Maybe a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe a
forall a. Maybe a
Nothing

-- accumulator for a result under construction
data Builder = Builder
  { Builder -> [FieldDescription]
accFields :: [FieldDescription],
    Builder -> [[Maybe ByteString]]
accRevRows :: [[Maybe ByteString]],
    Builder -> [Word32]
accParamOids :: [Word32],
    Builder -> Bool
accSawRowDescription :: Bool,
    Builder -> Bool
accHadResponse :: Bool
  }

emptyBuilder :: Builder
emptyBuilder :: Builder
emptyBuilder = [FieldDescription]
-> [[Maybe ByteString]] -> [Word32] -> Bool -> Bool -> Builder
Builder [] [] [] Bool
False Bool
False

-- | Collect the (possibly several) results of a simple query, up to
-- @ReadyForQuery@. The last is what @PQexec@ returns.
-- @CopyInResponse@ and @CopyOutResponse@ terminate the loop immediately,
-- returning a synthetic result so the caller can enter the copy sub-protocol.
collectSimple :: Connection -> ByteString -> IO [NativeResult]
collectSimple :: Connection -> ByteString -> IO [NativeResult]
collectSimple Connection
connection ByteString
sql = Builder -> [NativeResult] -> IO [NativeResult]
go Builder
emptyBuilder []
  where
    go :: Builder -> [NativeResult] -> IO [NativeResult]
go Builder
builder [NativeResult]
acc = do
      BackendMessage
message <- Connection -> IO BackendMessage
nextMessage Connection
connection
      case BackendMessage
message of
        RowDescription [FieldDescription]
fs -> Builder -> [NativeResult] -> IO [NativeResult]
go Builder
builder {accFields = fs, accSawRowDescription = True} [NativeResult]
acc
        DataRow [Maybe ByteString]
values -> Builder -> [NativeResult] -> IO [NativeResult]
go Builder
builder {accRevRows = values : (accRevRows builder)} [NativeResult]
acc
        CommandComplete ByteString
tag -> do
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
          Builder -> [NativeResult] -> IO [NativeResult]
go Builder
emptyBuilder (Builder -> Maybe ByteString -> NativeResult
commandResult Builder
builder (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
tag) NativeResult -> [NativeResult] -> [NativeResult]
forall a. a -> [a] -> [a]
: [NativeResult]
acc)
        BackendMessage
EmptyQueryResponse -> do
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
          Builder -> [NativeResult] -> IO [NativeResult]
go Builder
emptyBuilder (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
EmptyQuery [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
"" NativeResult -> [NativeResult] -> [NativeResult]
forall a. a -> [a] -> [a]
: [NativeResult]
acc)
        ErrorResponse [(Word8, ByteString)]
fs -> do
          let errMap :: Map Word8 ByteString
errMap = [(Word8, ByteString)] -> Map Word8 ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Word8, ByteString)]
fs
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Map Word8 ByteString -> ByteString
formatResultError ByteString
sql Map Word8 ByteString
errMap))
          Builder -> [NativeResult] -> IO [NativeResult]
go Builder
emptyBuilder (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
FatalError [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
errMap [] ByteString
sql NativeResult -> [NativeResult] -> [NativeResult]
forall a. a -> [a] -> [a]
: [NativeResult]
acc)
        CopyInResponse Int16
_ [Int16]
formats ->
          let fields :: [FieldDescription]
fields = (Int16 -> FieldDescription) -> [Int16] -> [FieldDescription]
forall a b. (a -> b) -> [a] -> [b]
map Int16 -> FieldDescription
copyField [Int16]
formats
           in [NativeResult] -> IO [NativeResult]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([NativeResult] -> [NativeResult]
forall a. [a] -> [a]
reverse (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
CopyIn [FieldDescription]
fields [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
"" NativeResult -> [NativeResult] -> [NativeResult]
forall a. a -> [a] -> [a]
: [NativeResult]
acc))
        CopyOutResponse Int16
_ [Int16]
formats ->
          let fields :: [FieldDescription]
fields = (Int16 -> FieldDescription) -> [Int16] -> [FieldDescription]
forall a b. (a -> b) -> [a] -> [b]
map Int16 -> FieldDescription
copyField [Int16]
formats
           in [NativeResult] -> IO [NativeResult]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([NativeResult] -> [NativeResult]
forall a. [a] -> [a]
reverse (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
CopyOut [FieldDescription]
fields [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
"" NativeResult -> [NativeResult] -> [NativeResult]
forall a. a -> [a] -> [a]
: [NativeResult]
acc))
        ReadyForQuery Word8
txState -> do
          IORef Word8 -> Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Word8
txStatus Connection
connection) Word8
txState
          [NativeResult] -> IO [NativeResult]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([NativeResult] -> [NativeResult]
forall a. [a] -> [a]
reverse [NativeResult]
acc)
        BackendMessage
_ -> Builder -> [NativeResult] -> IO [NativeResult]
go Builder
builder [NativeResult]
acc

-- | Collect the single result of an extended-protocol command.
collectExtended :: Connection -> ByteString -> IO NativeResult
collectExtended :: Connection -> ByteString -> IO NativeResult
collectExtended Connection
connection ByteString
sql = Builder -> Maybe NativeResult -> IO NativeResult
go Builder
emptyBuilder Maybe NativeResult
forall a. Maybe a
Nothing
  where
    go :: Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder Maybe NativeResult
finished = do
      BackendMessage
message <- Connection -> IO BackendMessage
nextMessage Connection
connection
      case BackendMessage
message of
        RowDescription [FieldDescription]
fs -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder {accFields = fs, accSawRowDescription = True} Maybe NativeResult
finished
        ParameterDescription [Word32]
oids -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder {accParamOids = oids} Maybe NativeResult
finished
        BackendMessage
NoData -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder Maybe NativeResult
finished
        DataRow [Maybe ByteString]
values -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder {accRevRows = values : (accRevRows builder)} Maybe NativeResult
finished
        BackendMessage
ParseComplete -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder Maybe NativeResult
finished
        BackendMessage
BindComplete -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder Maybe NativeResult
finished
        BackendMessage
CloseComplete -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder Maybe NativeResult
finished
        BackendMessage
PortalSuspended -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
emptyBuilder (Maybe NativeResult
finished Maybe NativeResult -> Maybe NativeResult -> Maybe NativeResult
forall a. Maybe a -> Maybe a -> Maybe a
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (Builder -> Maybe ByteString -> NativeResult
commandResult Builder
builder Maybe ByteString
forall a. Maybe a
Nothing))
        CommandComplete ByteString
tag -> do
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
          Builder -> Maybe NativeResult -> IO NativeResult
go Builder
emptyBuilder (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (Builder -> Maybe ByteString -> NativeResult
commandResult Builder
builder (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
tag)))
        BackendMessage
EmptyQueryResponse -> do
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just ByteString
"")
          Builder -> Maybe NativeResult -> IO NativeResult
go Builder
emptyBuilder (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
EmptyQuery [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
forall k a. Map k a
Map.empty [] ByteString
""))
        ErrorResponse [(Word8, ByteString)]
fs -> do
          let errMap :: Map Word8 ByteString
errMap = [(Word8, ByteString)] -> Map Word8 ByteString
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Word8, ByteString)]
fs
          IORef (Maybe ByteString) -> Maybe ByteString -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef (Maybe ByteString)
lastError Connection
connection) (ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Map Word8 ByteString -> ByteString
formatResultError ByteString
sql Map Word8 ByteString
errMap))
          Builder -> Maybe NativeResult -> IO NativeResult
go Builder
emptyBuilder (NativeResult -> Maybe NativeResult
forall a. a -> Maybe a
Just (ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult ExecStatus
FatalError [] [] Maybe ByteString
forall a. Maybe a
Nothing Map Word8 ByteString
errMap [] ByteString
sql))
        ReadyForQuery Word8
txState -> do
          IORef Word8 -> Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (Connection -> IORef Word8
txStatus Connection
connection) Word8
txState
          NativeResult -> IO NativeResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (NativeResult -> Maybe NativeResult -> NativeResult
forall a. a -> Maybe a -> a
fromMaybe (Builder -> NativeResult
describeResult Builder
builder) Maybe NativeResult
finished)
        BackendMessage
_ -> Builder -> Maybe NativeResult -> IO NativeResult
go Builder
builder Maybe NativeResult
finished

-- | A result terminated by @CommandComplete@\/@PortalSuspended@: 'TuplesOk' if a
-- row description was seen, else 'CommandOk'.
commandResult :: Builder -> Maybe ByteString -> NativeResult
commandResult :: Builder -> Maybe ByteString -> NativeResult
commandResult Builder
builder Maybe ByteString
tag =
  ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult
    (if (Builder -> Bool
accSawRowDescription Builder
builder) then ExecStatus
TuplesOk else ExecStatus
CommandOk)
    (Builder -> [FieldDescription]
accFields Builder
builder)
    ([[Maybe ByteString]] -> [[Maybe ByteString]]
forall a. [a] -> [a]
reverse (Builder -> [[Maybe ByteString]]
accRevRows Builder
builder))
    Maybe ByteString
tag
    Map Word8 ByteString
forall k a. Map k a
Map.empty
    (Builder -> [Word32]
accParamOids Builder
builder)
    ByteString
""

-- | A result with no command completion (a @Describe@\/@Parse@-only flow):
-- 'CommandOk', carrying any column descriptions and parameter OIDs.
describeResult :: Builder -> NativeResult
describeResult :: Builder -> NativeResult
describeResult Builder
builder =
  ExecStatus
-> [FieldDescription]
-> [[Maybe ByteString]]
-> Maybe ByteString
-> Map Word8 ByteString
-> [Word32]
-> ByteString
-> NativeResult
NativeResult
    ExecStatus
CommandOk
    (Builder -> [FieldDescription]
accFields Builder
builder)
    ([[Maybe ByteString]] -> [[Maybe ByteString]]
forall a. [a] -> [a]
reverse (Builder -> [[Maybe ByteString]]
accRevRows Builder
builder))
    Maybe ByteString
forall a. Maybe a
Nothing
    Map Word8 ByteString
forall k a. Map k a
Map.empty
    (Builder -> [Word32]
accParamOids Builder
builder)
    ByteString
""

lastMaybe :: [a] -> Maybe a
lastMaybe :: forall a. [a] -> Maybe a
lastMaybe = (Maybe a -> a -> Maybe a) -> Maybe a -> [a] -> Maybe a
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (\Maybe a
_ a
x -> a -> Maybe a
forall a. a -> Maybe a
Just a
x) Maybe a
forall a. Maybe a
Nothing

-- | Build a synthetic 'FieldDescription' from a COPY format code (0=text,
-- 1=binary).  COPY results have no column names, table OID, type OID, etc.
copyField :: Int16 -> FieldDescription
copyField :: Int16 -> FieldDescription
copyField Int16
fmt =
  FieldDescription
    { name :: ByteString
name = ByteString
"",
      tableOid :: Word32
tableOid = Word32
0,
      columnAttributeNumber :: Int16
columnAttributeNumber = Int16
0,
      typeOid :: Word32
typeOid = Word32
0,
      typeSize :: Int16
typeSize = Int16
0,
      typeModifier :: Int32
typeModifier = Int32
0,
      formatCode :: Int16
formatCode = Int16
fmt
    }