Query.hs

-- | 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 Control.Exception (IOException, catch, mask_)
import qualified Data.Map.Strict as Map
import qualified Data.Sequence as Seq
import qualified Data.Vector.Mutable as MVector
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 sql params resultFormat =
  asyncParamsWrite sql params resultFormat <> syncMessage

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

prepareWrite :: ByteString -> ByteString -> Maybe [Word32] -> Poker.Write
prepareWrite name sql parameterTypes =
  parseMessage name sql (fromMaybe [] parameterTypes) <> 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 sql params resultFormat =
  parseMessage "" sql (fmap paramOid params)
    <> bindMessage "" "" (fmap paramFormat params) (fmap paramValue params) [formatCodeOf resultFormat]
    <> describePortalMessage ""
    <> executeMessage "" 0

asyncPreparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write
asyncPreparedWrite name params resultFormat =
  bindMessage "" name (fmap boundFormat params) (fmap boundValue params) [formatCodeOf resultFormat]
    <> describePortalMessage ""
    <> executeMessage "" 0

-- * Parameter-count limit

-- | The largest parameter count the wire protocol's 16-bit count fields can
-- carry, matching libpq's @PQ_QUERY_PARAM_MAX_LIMIT@. 'Parse' and 'Bind'
-- messages encode their parameter count as an 'Int16'; past this limit that
-- encoding wraps around instead of overflowing, so the count must be
-- rejected locally before it corrupts the message.
maxParamCount :: Int
maxParamCount = 65535

tooManyParams :: [a] -> Bool
tooManyParams = (> maxParamCount) . length

-- * Synchronous flows

-- | Simple query. Returns the last result, mirroring @PQexec@.
exec :: Connection -> ByteString -> IO (Maybe NativeResult)
exec connection sql = withReady connection do
  catch
    (sendMessage connection (queryMessage sql) >> (lastMaybe <$> collectSimple connection sql))
    (fmap Just . connectionLostResult connection sql)

-- | Parameterized query via the extended protocol.
execParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO (Maybe NativeResult)
execParams connection sql params resultFormat
  | tooManyParams params = pure Nothing
  | otherwise = withReady connection do
      catch
        (sendMessage connection (paramsWrite sql params resultFormat) >> (Just <$> collectExtended connection sql))
        (fmap Just . connectionLostResult connection sql)

-- | Prepare a named statement.
prepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO (Maybe NativeResult)
prepare connection name sql parameterTypes
  | maybe False tooManyParams parameterTypes = pure Nothing
  | otherwise = withReady connection do
      catch
        (sendMessage connection (prepareWrite name sql parameterTypes) >> (Just <$> collectExtended connection sql))
        (fmap Just . connectionLostResult connection sql)

-- | Execute a previously prepared statement.
execPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO (Maybe NativeResult)
execPrepared connection name params resultFormat
  | tooManyParams params = pure Nothing
  | otherwise = withReady connection do
      catch
        (sendMessage connection (preparedWrite name params resultFormat) >> (Just <$> collectExtended connection ""))
        (fmap Just . connectionLostResult connection "")

-- * Asynchronous flows

-- | Send a write in async mode, tracking pending commands for pipeline abort.
--
-- A socket death mid-send (e.g. @EPIPE@\/@ECONNRESET@) surfaces here as an
-- escaped 'IOException' from 'Transport.send', which 'sendMessage' does not
-- catch. Left uncaught, it would blow straight through this function - and
-- every caller layered on it ('Pqi.Native.sendQuery' etc., 'Hasql.Comms.Send')
-- - as a raw exception instead of the @False@ that @PQsendQuery@ always
-- returns for a fatal send. Catching it here and marking the connection bad
-- keeps the contract: the caller sees a normal failed send, discoverable via
-- 'Pqi.status', exactly as libpq's own internals never throw and always
-- record the failure on the connection instead.
sendAsync :: Connection -> ByteString -> Poker.Write -> IO Bool
sendAsync connection sql write = do
  status <- readIORef (connStatus connection)
  case status of
    ConnectionOk ->
      catch
        do
          sendMessage connection write
          writeIORef (currentQuery connection) sql
          writeIORef (asyncPending connection) True
          pipeStatus <- readIORef (pipelineStatus connection)
          when (pipeStatus /= PipelineOff) $ modifyIORef' (pendingCommands connection) (+ 1)
          pure True
        (\err -> False <$ markConnectionLost connection err)
    _ -> pure False

-- | Whether the connection is in pipeline mode.
inPipeline :: Connection -> IO Bool
inPipeline connection = (/= PipelineOff) <$> readIORef (pipelineStatus connection)

-- Simple query protocol: no Sync needed (server sends ReadyForQuery on its own).
sendQuery :: Connection -> ByteString -> IO Bool
sendQuery connection sql = sendAsync connection sql (queryMessage 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 sql params resultFormat
  | tooManyParams params = pure False
  | otherwise = do
      pipeline <- inPipeline connection
      ok <-
        sendAsync connection sql
          $ if pipeline
            then asyncParamsWrite sql params resultFormat
            else paramsWrite sql params 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'.
      when (ok && pipeline) $ modifyIORef' (pendingParseOrigins connection) (Seq.|> Just False)
      pure ok

sendPrepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO Bool
sendPrepare connection name sql parameterTypes
  | maybe False tooManyParams parameterTypes = pure False
  | otherwise = do
      pipeline <- inPipeline connection
      ok <-
        sendAsync connection sql
          $ if pipeline
            then parseMessage name sql (fromMaybe [] parameterTypes)
            else prepareWrite name sql parameterTypes
      when (ok && pipeline) $ modifyIORef' (pendingParseOrigins connection) (Seq.|> Just True)
      pure ok

-- | Push a @Nothing@ origin for a pipelined command that sends no @Parse@
-- (so has no @ParseComplete@ to fold or terminate on): it still occupies one
-- slot in the FIFO, popped and discarded at its own terminal message.
sendAsyncNoOrigin :: Connection -> Bool -> ByteString -> Poker.Write -> IO Bool
sendAsyncNoOrigin connection pipeline sql write = do
  ok <- sendAsync connection sql write
  when (ok && pipeline) $ modifyIORef' (pendingParseOrigins connection) (Seq.|> Nothing)
  pure ok

sendQueryPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO Bool
sendQueryPrepared connection name params resultFormat
  | tooManyParams params = pure False
  | otherwise = do
      pipeline <- inPipeline connection
      sendAsyncNoOrigin connection pipeline ""
        $ if pipeline
          then asyncPreparedWrite name params resultFormat
          else preparedWrite name params resultFormat

sendDescribePrepared :: Connection -> ByteString -> IO Bool
sendDescribePrepared connection name = do
  pipeline <- inPipeline connection
  sendAsyncNoOrigin connection pipeline ""
    $ if pipeline
      then describeStatementMessage name
      else describeStatementMessage name <> syncMessage

sendDescribePortal :: Connection -> ByteString -> IO Bool
sendDescribePortal connection name = do
  pipeline <- inPipeline connection
  sendAsyncNoOrigin connection pipeline ""
    $ if pipeline
      then describePortalMessage name
      else describePortalMessage name <> 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.
--
-- Runs 'mask_'ed. The connection's result bookkeeping - the pending-command
-- counter, the separator flag, the @ParseComplete@ origin FIFO, the pipeline
-- status - lives in separate 'IORef's that a single logical transition
-- updates one after another. An async exception landing between two of those
-- updates leaves the pair inconsistent, and an inconsistent pair is not merely
-- wrong: it sends the next 'getNextResult' down the @readAndProcess@ path to
-- wait for a message the backend has already decided not to send, which is a
-- stall with no timer behind it. Masking keeps each transition atomic. The
-- blocking read inside stays interruptible (see 'Transport.receiveFrame'), so
-- this costs no abandonability.
getNextResult :: Connection -> IO (Maybe NativeResult)
getNextResult connection = mask_ do
  pending <- readIORef (asyncPending connection)
  if not pending
    then pure Nothing
    else do
      sepPending <- readIORef (pipelineSeparatorPending connection)
      if sepPending
        then do
          writeIORef (pipelineSeparatorPending connection) False
          pure Nothing
        else do
          singleRow <- readIORef (singleRowMode connection)
          cachedFields <- readIORef (singleRowFields connection)
          builder <- newBuilder
          let initBuilder =
                if singleRow && not (null cachedFields)
                  then builder {accFields = cachedFields, accSawRowDescription = True}
                  else builder
          go singleRow 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 pipeStatus = do
      when (pipeStatus /= PipelineOff) $ do
        modifyIORef' (pendingCommands connection) (subtract 1)
        writeIORef (pipelineSeparatorPending connection) True

    -- Pop and discard this command's 'pendingParseOrigins' entry if it hasn't
    -- already been consumed by its own 'ParseComplete' (see the field's
    -- Haddock on 'Connection'). Called at every terminal message besides a
    -- 'ParseComplete' that itself terminates the command.
    popOriginIfPending pipeStatus builder =
      when (pipeStatus /= PipelineOff && not (accOriginPopped builder))
        $ void (popPendingParseOrigin connection)

    go singleRow builder = do
      pipeStatus <- readIORef (pipelineStatus 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.
      pending <- readIORef (pendingCommands connection)
      if pipeStatus == PipelineAborted && pending > 0
        then do
          -- This command is being synthesized without ever reading a
          -- message for it (the server sends nothing for a discarded
          -- command), so its origin - if it pushed one - is necessarily
          -- still unpopped.
          popOriginIfPending pipeStatus builder
          modifyIORef' (pendingCommands connection) (subtract 1)
          writeIORef (pipelineSeparatorPending connection) True
          pure (Just (NativeResult PipelineAbort [] [] Nothing Map.empty [] ""))
        else readAndProcess singleRow builder pipeStatus

    readAndProcess singleRow builder pipeStatus = do
      message <- nextMessage connection
      case message of
        RowDescription fs ->
          go singleRow builder {accFields = fs, accSawRowDescription = True, accHadResponse = True}
        ParameterDescription oids ->
          go singleRow builder {accParamOids = oids, accHadResponse = True}
        NoData ->
          go singleRow builder {accHadResponse = True}
        DataRow values ->
          if singleRow
            then do
              writeIORef (singleRowFields connection) (accFields builder)
              pure (Just (NativeResult SingleTuple (accFields builder) [values] Nothing Map.empty [] ""))
            else do
              pushRow (accRows builder) values
              go singleRow builder
        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').
          origin <-
            if pipeStatus /= PipelineOff
              then popPendingParseOrigin connection
              else pure Nothing
          case origin of
            Just (Just True) -> do
              finishCommand pipeStatus
              pure (Just (NativeResult CommandOk [] [] Nothing Map.empty [] ""))
            Just _ -> go singleRow builder {accHadResponse = True, accOriginPopped = True}
            Nothing -> go singleRow builder {accHadResponse = True}
        BindComplete ->
          go singleRow builder {accHadResponse = True}
        CloseComplete ->
          go singleRow builder {accHadResponse = True}
        CommandComplete tag -> do
          if singleRow
            then do
              writeIORef (singleRowMode connection) False
              writeIORef (singleRowFields connection) []
              writeIORef (lastError connection) (Just "")
              pure (Just (NativeResult TuplesOk (accFields builder) [] (Just tag) Map.empty [] ""))
            else do
              writeIORef (lastError connection) (Just "")
              popOriginIfPending pipeStatus builder
              finishCommand pipeStatus
              Just <$> commandResult builder (Just tag)
        EmptyQueryResponse -> do
          writeIORef (lastError connection) (Just "")
          popOriginIfPending pipeStatus builder
          finishCommand pipeStatus
          pure (Just (NativeResult EmptyQuery [] [] Nothing Map.empty [] ""))
        ErrorResponse fs -> do
          let errMap = Map.fromList fs
          case pipeStatus of
            PipelineAborted -> do
              -- Should not normally happen (server discards commands in abort
              -- mode) but handle defensively.
              popOriginIfPending pipeStatus builder
              finishCommand pipeStatus
              pure (Just (NativeResult PipelineAbort [] [] Nothing Map.empty [] ""))
            PipelineOn -> do
              writeIORef (pipelineStatus connection) PipelineAborted
              popOriginIfPending pipeStatus builder
              finishCommand PipelineOn
              pure (Just (NativeResult FatalError [] [] Nothing errMap [] ""))
            PipelineOff -> do
              sql <- readIORef (currentQuery connection)
              writeIORef (lastError connection) (Just (formatResultError sql errMap))
              pure (Just (NativeResult FatalError [] [] Nothing errMap [] sql))
        PortalSuspended ->
          Just <$> commandResult builder Nothing
        ReadyForQuery txState -> do
          writeIORef (txStatus connection) txState
          case pipeStatus of
            PipelineOff -> do
              writeIORef (asyncPending connection) False
              if (accHadResponse builder)
                then Just <$> describeResult builder
                else pure Nothing
            _ -> do
              writeIORef (pipelineStatus connection) 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.
              remaining <- atomicModifyIORef' (pendingSyncs connection) (\n -> (n - 1, n - 1))
              when (remaining == 0) $ writeIORef (asyncPending connection) False
              pure (Just (NativeResult PipelineSync [] [] Nothing Map.empty [] ""))
        _ -> go singleRow builder

-- | Pop the next in-flight pipelined command's origin: @Just (Just True)@ if
-- it is the terminal @ParseComplete@ of a 'sendPrepare' (to be materialized
-- as 'CommandOk'), @Just (Just False)@ if it belongs to a 'sendQueryParams'
-- and must fold into that command's accumulating result, @Just Nothing@ for
-- a command with no @Parse@ step. The outer @Nothing@ means the FIFO is
-- empty (defensive: nothing is popped and the caller folds).
popPendingParseOrigin :: Connection -> IO (Maybe (Maybe Bool))
popPendingParseOrigin connection =
  atomicModifyIORef'
    (pendingParseOrigins connection)
    ( \queue -> case Seq.viewl queue of
        origin Seq.:< rest -> (rest, Just origin)
        Seq.EmptyL -> (queue, Nothing)
    )

-- | Describe a prepared statement.
describePrepared :: Connection -> ByteString -> IO (Maybe NativeResult)
describePrepared connection name = withReady connection do
  sendMessage connection (describeStatementMessage name <> syncMessage)
  catch (Just <$> collectExtended connection "") (fmap Just . connectionLostResult connection "")

-- | Describe a portal.
describePortal :: Connection -> ByteString -> IO (Maybe NativeResult)
describePortal connection name = withReady connection do
  sendMessage connection (describePortalMessage name <> syncMessage)
  catch (Just <$> collectExtended connection "") (fmap Just . connectionLostResult connection "")

-- * Parameter projections

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

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

paramValue :: Maybe (Word32, ByteString, Format) -> Maybe ByteString
paramValue = fmap (\(_, value, _) -> value)

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

boundValue :: Maybe (ByteString, Format) -> Maybe ByteString
boundValue = fmap fst

-- | Renamed from @formatCode@ to avoid clashing with 'FieldDescription's
-- @formatCode@ field now that 'DuplicateRecordFields' is no longer enabled.
formatCodeOf :: Format -> Int16
formatCodeOf = \case
  Text -> 0
  Binary -> 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 connection action = do
  status <- readIORef (connStatus connection)
  case status of
    ConnectionOk -> action
    _ -> pure Nothing

-- | Turn a read loop's escaped 'IOException' - e.g. a connection reset while
-- a result is still in flight - into a classified 'FatalError' result,
-- matching @PQexec@: libpq never throws here, it reports the same
-- "server closed the connection unexpectedly" wording it uses for a
-- handshake-time loss (see 'connectionLostMessage'), and marks the
-- connection bad so the caller's next call sees it too.
connectionLostResult :: Connection -> ByteString -> IOException -> IO NativeResult
connectionLostResult connection sql err = do
  message <- markConnectionLost connection err
  pure (NativeResult FatalError [] [] Nothing (Map.singleton 0x4d message) [] sql)

-- | A growable row buffer, replacing a cons-list-plus-reverse accumulation
-- with amortized-O(1) appends into a doubling mutable vector. A statement's
-- rows are pushed one 'DataRow' at a time and can number in the hundreds of
-- thousands; accumulating them as @row : rows@ and reversing at the end (the
-- prior representation) keeps two full-length cons-cell spines alive across
-- the whole statement, which dominates GC copy volume at scale (see
-- 'hasql-pqi-native-decode-gc-bound' memory). This buffer holds only the
-- array (which the RTS moves as one contiguous block, unlike a linked list)
-- until 'freezeRowBuffer' walks it once into the final list.
data RowBuffer
  = RowBuffer
      -- | The mutable vector of rows.
      !(IORef (MVector.IOVector [Maybe ByteString]))
      -- | Length.
      !(IORef Int)

newRowBuffer :: IO RowBuffer
newRowBuffer = do
  v <- MVector.new 64
  RowBuffer <$> newIORef v <*> newIORef 0

pushRow :: RowBuffer -> [Maybe ByteString] -> IO ()
pushRow (RowBuffer vRef lenRef) row = do
  v <- readIORef vRef
  len <- readIORef lenRef
  v' <-
    if len >= MVector.length v
      then do
        grown <- MVector.grow v (MVector.length v)
        writeIORef vRef grown
        pure grown
      else pure v
  MVector.write v' len row
  writeIORef lenRef (len + 1)

-- | Walk the buffer back-to-front once, consing onto an accumulator, which
-- yields the rows in their original order without a separate reverse pass.
freezeRowBuffer :: RowBuffer -> IO [[Maybe ByteString]]
freezeRowBuffer (RowBuffer vRef lenRef) = do
  v <- readIORef vRef
  len <- readIORef lenRef
  let go !i acc
        | i < 0 = pure acc
        | otherwise = do
            x <- MVector.read v i
            go (i - 1) (x : acc)
  go (len - 1) []

-- accumulator for a result under construction
data Builder = Builder
  { accFields :: [FieldDescription],
    accRows :: RowBuffer,
    accParamOids :: [Word32],
    accSawRowDescription :: Bool,
    accHadResponse :: Bool,
    -- | Whether this command's 'pendingParseOrigins' entry has already been
    -- popped (via its own 'ParseComplete'). See 'popOriginIfPending'.
    accOriginPopped :: Bool
  }

newBuilder :: IO Builder
newBuilder = do
  rows <- newRowBuffer
  pure (Builder [] rows [] False False 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 sql = do
  builder0 <- newBuilder
  go builder0 []
  where
    go builder acc = do
      message <- nextMessage connection
      case message of
        RowDescription fs -> go builder {accFields = fs, accSawRowDescription = True} acc
        DataRow values -> do
          pushRow (accRows builder) values
          go builder acc
        CommandComplete tag -> do
          writeIORef (lastError connection) (Just "")
          result <- commandResult builder (Just tag)
          builder' <- newBuilder
          go builder' (result : acc)
        EmptyQueryResponse -> do
          writeIORef (lastError connection) (Just "")
          builder' <- newBuilder
          go builder' (NativeResult EmptyQuery [] [] Nothing Map.empty [] "" : acc)
        ErrorResponse fs -> do
          let errMap = Map.fromList fs
          writeIORef (lastError connection) (Just (formatResultError sql errMap))
          builder' <- newBuilder
          go builder' (NativeResult FatalError [] [] Nothing errMap [] sql : acc)
        CopyInResponse _ formats ->
          let fields = map copyField formats
           in pure (reverse (NativeResult CopyIn fields [] Nothing Map.empty [] "" : acc))
        CopyOutResponse _ formats ->
          let fields = map copyField formats
           in pure (reverse (NativeResult CopyOut fields [] Nothing Map.empty [] "" : acc))
        ReadyForQuery txState -> do
          writeIORef (txStatus connection) txState
          pure (reverse acc)
        _ -> go builder acc

-- | Collect the single result of an extended-protocol command.
collectExtended :: Connection -> ByteString -> IO NativeResult
collectExtended connection sql = do
  builder0 <- newBuilder
  go builder0 Nothing
  where
    go builder finished = do
      message <- nextMessage connection
      case message of
        RowDescription fs -> go builder {accFields = fs, accSawRowDescription = True} finished
        ParameterDescription oids -> go builder {accParamOids = oids} finished
        NoData -> go builder finished
        DataRow values -> do
          pushRow (accRows builder) values
          go builder finished
        ParseComplete -> go builder finished
        BindComplete -> go builder finished
        CloseComplete -> go builder finished
        PortalSuspended -> do
          finished' <- case finished of
            Just _ -> pure finished
            Nothing -> Just <$> commandResult builder Nothing
          builder' <- newBuilder
          go builder' finished'
        CommandComplete tag -> do
          writeIORef (lastError connection) (Just "")
          result <- commandResult builder (Just tag)
          builder' <- newBuilder
          go builder' (Just result)
        EmptyQueryResponse -> do
          writeIORef (lastError connection) (Just "")
          builder' <- newBuilder
          go builder' (Just (NativeResult EmptyQuery [] [] Nothing Map.empty [] ""))
        ErrorResponse fs -> do
          let errMap = Map.fromList fs
          writeIORef (lastError connection) (Just (formatResultError sql errMap))
          builder' <- newBuilder
          go builder' (Just (NativeResult FatalError [] [] Nothing errMap [] sql))
        ReadyForQuery txState -> do
          writeIORef (txStatus connection) txState
          case finished of
            Just result -> pure result
            Nothing -> describeResult builder
        _ -> go builder finished

-- | A result terminated by @CommandComplete@\/@PortalSuspended@: 'TuplesOk' if a
-- row description was seen, else 'CommandOk'.
commandResult :: Builder -> Maybe ByteString -> IO NativeResult
commandResult builder tag = do
  rows <- freezeRowBuffer (accRows builder)
  pure
    $ NativeResult
      (if (accSawRowDescription builder) then TuplesOk else CommandOk)
      (accFields builder)
      rows
      tag
      Map.empty
      (accParamOids builder)
      ""

-- | A result with no command completion (a @Describe@\/@Parse@-only flow):
-- 'CommandOk', carrying any column descriptions and parameter OIDs.
describeResult :: Builder -> IO NativeResult
describeResult builder = do
  rows <- freezeRowBuffer (accRows builder)
  pure
    $ NativeResult
      CommandOk
      (accFields builder)
      rows
      Nothing
      Map.empty
      (accParamOids builder)
      ""

lastMaybe :: [a] -> Maybe a
lastMaybe = foldl (\_ x -> Just x) 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 fmt =
  FieldDescription
    { name = "",
      tableOid = 0,
      columnAttributeNumber = 0,
      typeOid = 0,
      typeSize = 0,
      typeModifier = 0,
      formatCode = fmt
    }