packages feed

pg-wire (empty) → 0.1.0.0

raw patch · 41 files changed

+6772/−0 lines, 41 filesdep +asyncdep +basedep +base64-bytestring

Dependencies added: async, base, base64-bytestring, bytestring, containers, cryptohash-md5, crypton, crypton-x509-store, crypton-x509-system, deepseq, directory, hashable, hedgehog, hspec, hspec-hedgehog, memory, network, nothunks, pg-wire, psqueues, random, stm, text, time, tls, vector

Files

+ CHANGELOG.md view
@@ -0,0 +1,25 @@+# Changelog++The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).++## [0.1.0.0] - 2026-04-25++Initial release. Pure-Haskell implementation of the PostgreSQL v3 wire+protocol.++- Connection setup over TCP and Unix sockets, TLS 1.2/1.3,+  cleartext / MD5 / SCRAM-SHA-256 authentication, multi-host failover+  (`target_session_attrs`, `load_balance_hosts`), startup parameters.+- Connection pool with idle reaping, health checks, jittered max-life,+  queue mode (LIFO/FIFO), recycling modes (always/on-error), lifecycle+  hooks, and `nothunks`-checked invariants.+- Wire frames: extended query (parse / bind / describe / execute),+  simple query, COPY in/out, large objects, LISTEN/NOTIFY, query+  cancellation.+- Binary format type classes (`PgEncode`, `PgDecode`) and OID metadata.+  Concrete codecs for PostgreSQL types live in the+  [`valiant`](https://hackage.haskell.org/package/valiant) runtime.+- Nine structured error types with column-by-column diagnostics.++[0.1.0.0]: https://github.com/joshburgess/valiant/releases/tag/v0.1.0.0
+ LICENSE view
@@ -0,0 +1,27 @@+Copyright (c) 2026 Josh Burgess++Redistribution and use in source and binary forms, with or without+modification, are permitted provided that the following conditions are met:++1. Redistributions of source code must retain the above copyright notice,+   this list of conditions and the following disclaimer.++2. Redistributions in binary form must reproduce the above copyright notice,+   this list of conditions and the following disclaimer in the documentation+   and/or other materials provided with the distribution.++3. Neither the name of the copyright holder nor the names of its+   contributors may be used to endorse or promote products derived from this+   software without specific prior written permission.++THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE+ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE+LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR+CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF+SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS+INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN+CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)+ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE+POSSIBILITY OF SUCH DAMAGE.
+ README.md view
@@ -0,0 +1,51 @@+# pg-wire++Pure-Haskell driver for the PostgreSQL v3 wire protocol. No `libpq`, no+C dependencies.++## What you get++- Connection setup: TCP and Unix sockets, TLS 1.2/1.3, MD5 and+  SCRAM-SHA-256 auth, multi-host failover (`target_session_attrs`,+  `load_balance_hosts`), startup parameters.+- Connection pool: idle reaping, health checks, jittered max-life,+  queue mode (LIFO/FIFO), recycling modes (always/on-error), lifecycle+  hooks, `nothunks` invariants.+- Wire frames: extended query (parse / bind / describe / execute),+  simple query, COPY in/out, large objects, LISTEN/NOTIFY, query+  cancellation.+- Binary format type classes: `PgEncode`, `PgDecode`, OID metadata.+  Codecs for individual types live in the `valiant` runtime library.+- Errors: 8 structured error constructors (`ConnectionError`,+  `AuthError`, `ProtocolError`, `QueryError`, `DecodeError`,+  `PoolTimeout`, `PoolClosed`, `ConnectionDead`).++## When to use this directly++`pg-wire` is the transport layer. Most users want the higher-level+[`valiant`](https://hackage.haskell.org/package/valiant) runtime, which+re-exports everything from `pg-wire` and adds compile-time SQL safety+via the `valiant-plugin` GHC source plugin.++Use `pg-wire` directly only when you want raw wire-protocol access+without the compile-time validation layer (e.g. for building a+different ORM or for testing the protocol itself).++## Installation++```+cabal install pg-wire+```++## Documentation++- [Top-level README](https://github.com/joshburgess/valiant#readme):+  the project overview and how the four packages fit together.+- [`docs/`](https://github.com/joshburgess/valiant/tree/main/docs):+  tutorial, performance notes, async architecture, gap analysis.+- [`docs/benchmark-results/`](https://github.com/joshburgess/valiant/tree/main/docs/benchmark-results):+  archived comparative numbers.++## License++BSD-3-Clause. See `LICENSE`.
+ mockserver/PgWire/MockServer.hs view
@@ -0,0 +1,320 @@+-- | A minimal mock PostgreSQL server for testing.+--+-- Speaks enough of the PG v3 wire protocol to test client behavior+-- without a real database: authentication, simple queries, extended+-- query protocol, error responses, and connection lifecycle.+--+-- @+-- withMockServer defaultMockConfig $ \\port -> do+--   conn <- connectString (\"postgres:\/\/user:pass\@localhost:\" <> BS8.pack (show port) <> \"\/testdb\")+--   (rows, _) <- simpleQuery conn \"SELECT 1\"+--   close conn+-- @+module PgWire.MockServer+  ( -- * Server lifecycle+    withMockServer+  , MockConfig (..)+  , defaultMockConfig+    -- * Query handlers+  , QueryHandler+  , simpleHandler+  , errorHandler+  , failNTimes+    -- * Response builders+  , sendBackendMsg+  , buildBackendMsg+  , buildErrorResponse+  , buildCommandComplete+  , buildRowDescription+  , buildDataRow+  , buildReadyForQuery+  ) where++import Control.Concurrent (forkIO, killThread)+import Control.Exception (SomeException, catch, finally)+import Data.IORef+import Data.Bits (shiftL, (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Builder (Builder)+import Data.ByteString.Builder qualified as B+import Data.ByteString.Char8 qualified as BS8+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int32)+import Data.Word (Word8)+import Network.Socket qualified as NS+import Network.Socket.ByteString qualified as NSB++-- | Configuration for the mock server.+data MockConfig = MockConfig+  { mockUser :: ByteString+  -- ^ Expected username. Default: @\"testuser\"@.+  , mockPassword :: ByteString+  -- ^ Expected password. Default: @\"testpass\"@.+  , mockDatabase :: ByteString+  -- ^ Expected database name. Default: @\"testdb\"@.+  , mockQueryHandler :: QueryHandler+  -- ^ Handler for simple queries. Default: returns empty result.+  }++-- | Handler called for each simple query. Receives the SQL text and+-- a function to send backend messages.+type QueryHandler = ByteString -> (Builder -> IO ()) -> IO ()++-- | Default config: accepts auth, returns empty results for all queries.+defaultMockConfig :: MockConfig+defaultMockConfig = MockConfig+  { mockUser = "testuser"+  , mockPassword = "testpass"+  , mockDatabase = "testdb"+  , mockQueryHandler = \_ send -> do+      send $ buildCommandComplete "SELECT 0"+  }++-- | A simple query handler that returns fixed rows for a specific query.+simpleHandler :: [(ByteString, [[(ByteString, ByteString)]])] -> QueryHandler+simpleHandler table sql send = case lookup sql table of+  Nothing -> send $ buildCommandComplete "SELECT 0"+  Just rows -> do+    case rows of+      [] -> send $ buildCommandComplete "SELECT 0"+      (first : _) -> do+        -- Send RowDescription+        let cols = map fst first+        send $ buildRowDescription cols+        -- Send DataRows+        mapM_ (\row -> send $ buildDataRow (map snd row)) rows+        send $ buildCommandComplete ("SELECT " <> BS8.pack (show (length rows)))++-- | A handler that returns a PG error with the given SQLSTATE and message.+--+-- @+-- errorHandler \"23505\" \"duplicate key value violates unique constraint \\\"users_email_key\\\"\"+-- @+errorHandler :: ByteString -> ByteString -> QueryHandler+errorHandler sqlstate msg _ send =+  send $ buildErrorResponse sqlstate msg++-- | A handler backed by an IORef counter. Returns an error for the first+-- @n@ calls, then delegates to the fallback handler. Useful for testing+-- retry logic.+--+-- @+-- ref <- newIORef 0+-- let handler = failNTimes ref 2 \"40001\" \"serialization failure\" defaultHandler+-- @+failNTimes :: IORef Int -> Int -> ByteString -> ByteString -> QueryHandler -> QueryHandler+failNTimes ref n sqlstate msg fallback sql send = do+  count <- atomicModifyIORef' ref (\c -> (c + 1, c))+  if count < n+    then send $ buildErrorResponse sqlstate msg+    else fallback sql send++-- | Start a mock server on a random port, run an action with the port,+-- then shut down.+withMockServer :: MockConfig -> (NS.PortNumber -> IO a) -> IO a+withMockServer cfg action = do+  -- Create a listening socket+  sock <- NS.socket NS.AF_INET NS.Stream NS.defaultProtocol+  NS.setSocketOption sock NS.ReuseAddr 1+  NS.bind sock (NS.SockAddrInet 0 (NS.tupleToHostAddress (127, 0, 0, 1)))+  NS.listen sock 5+  port <- NS.socketPort sock++  -- Start acceptor thread+  tid <- forkIO $ acceptLoop sock cfg++  -- Run the action, then clean up+  action port `finally` do+    killThread tid+    NS.close sock++-- | Accept loop: handle one connection at a time.+acceptLoop :: NS.Socket -> MockConfig -> IO ()+acceptLoop listenSock cfg = go+  where+    go = do+      (clientSock, _) <- NS.accept listenSock+      _ <- forkIO (handleClient clientSock cfg `catch` \(_ :: SomeException) -> pure ())+      go++-- | Handle a single client connection.+handleClient :: NS.Socket -> MockConfig -> IO ()+handleClient sock cfg = do+  -- Read startup message (no tag byte)+  startupBytes <- recvN sock 4+  let startupLen = decodeInt32BE startupBytes 0+  _payload <- recvN sock (fromIntegral startupLen - 4)++  -- Send AuthOk (cleartext auth skipped for simplicity)+  sendBuilder sock $ buildAuth 0 -- AuthOk++  -- Send server parameters+  sendBuilder sock $ buildParameterStatus "server_version" "16.0"+  sendBuilder sock $ buildParameterStatus "server_encoding" "UTF8"+  sendBuilder sock $ buildParameterStatus "client_encoding" "UTF8"+  sendBuilder sock $ buildParameterStatus "is_superuser" "off"++  -- Send BackendKeyData+  sendBuilder sock $ buildBackendKeyData 1234 5678++  -- Send ReadyForQuery (Idle)+  sendBuilder sock $ buildReadyForQuery 'I'++  -- Message loop+  let loop = do+        mTag <- recvMaybe sock 1+        case mTag of+          Nothing -> pure () -- client disconnected+          Just tagBs -> do+            let tag = BS.index tagBs 0+            lenBs <- recvN sock 4+            let len = decodeInt32BE lenBs 0+            payload <- if len > 4 then recvN sock (fromIntegral len - 4) else pure BS.empty+            handleMessage sock cfg tag payload+            loop++  loop `catch` \(_ :: SomeException) -> pure ()+  NS.close sock++handleMessage :: NS.Socket -> MockConfig -> Word8 -> ByteString -> IO ()+handleMessage sock cfg tag payload = case tag of+  -- 'Q' = Simple Query+  0x51 -> do+    let sql = BS.takeWhile (/= 0) payload -- strip NUL terminator+    mockQueryHandler cfg sql (sendBuilder sock)+    sendBuilder sock $ buildReadyForQuery 'I'++  -- 'P' = Parse+  0x50 -> do+    sendBuilder sock $ buildSimpleTag '1' -- ParseComplete++  -- 'B' = Bind+  0x42 -> do+    sendBuilder sock $ buildSimpleTag '2' -- BindComplete++  -- 'D' = Describe+  0x44 -> do+    sendBuilder sock $ buildNoData++  -- 'E' = Execute+  0x45 -> do+    sendBuilder sock $ buildCommandComplete "SELECT 0"++  -- 'S' = Sync+  0x53 -> do+    sendBuilder sock $ buildReadyForQuery 'I'++  -- 'H' = Flush+  0x48 -> pure ()++  -- 'C' = Close+  0x43 -> do+    sendBuilder sock $ buildSimpleTag '3' -- CloseComplete++  -- 'X' = Terminate+  0x58 -> pure () -- connection will close++  -- Unknown+  _ -> do+    sendBuilder sock $ buildErrorResponse "08P01" ("Unknown message tag: " <> BS8.pack (show tag))+    sendBuilder sock $ buildReadyForQuery 'E'++------------------------------------------------------------------------+-- Backend message builders+------------------------------------------------------------------------++-- | Build a complete backend message as a Builder.+buildBackendMsg :: Word8 -> Builder -> Builder+buildBackendMsg tag payload =+  let payloadBs = LBS.toStrict (B.toLazyByteString payload)+      len = fromIntegral (BS.length payloadBs + 4) :: Int32+   in B.word8 tag <> B.int32BE len <> B.byteString payloadBs++-- | Send a backend message builder to a socket.+sendBackendMsg :: NS.Socket -> Word8 -> Builder -> IO ()+sendBackendMsg sock tag payload = sendBuilder sock (buildBackendMsg tag payload)++buildAuth :: Int32 -> Builder+buildAuth authType = buildBackendMsg 0x52 (B.int32BE authType)++buildParameterStatus :: ByteString -> ByteString -> Builder+buildParameterStatus name value =+  buildBackendMsg 0x53 (B.byteString name <> B.word8 0 <> B.byteString value <> B.word8 0)++buildBackendKeyData :: Int32 -> Int32 -> Builder+buildBackendKeyData pid key =+  buildBackendMsg 0x4B (B.int32BE pid <> B.int32BE key)++buildReadyForQuery :: Char -> Builder+buildReadyForQuery status =+  buildBackendMsg 0x5A (B.word8 (fromIntegral (fromEnum status)))++buildSimpleTag :: Char -> Builder+buildSimpleTag c = buildBackendMsg (fromIntegral (fromEnum c)) mempty++buildNoData :: Builder+buildNoData = buildSimpleTag 'n'++buildCommandComplete :: ByteString -> Builder+buildCommandComplete tag =+  buildBackendMsg 0x43 (B.byteString tag <> B.word8 0)++buildRowDescription :: [ByteString] -> Builder+buildRowDescription cols =+  buildBackendMsg 0x54 $ do+    B.int16BE (fromIntegral (length cols))+    <> foldMap buildFieldInfo cols+  where+    buildFieldInfo name =+      B.byteString name <> B.word8 0  -- name + NUL+        <> B.word32BE 0               -- table OID+        <> B.int16BE 0                -- column number+        <> B.word32BE 25              -- type OID (text)+        <> B.int16BE (-1)             -- type size+        <> B.int32BE (-1)             -- type modifier+        <> B.int16BE 0                -- format (text)++buildDataRow :: [ByteString] -> Builder+buildDataRow cols =+  buildBackendMsg 0x44 $+    B.int16BE (fromIntegral (length cols))+    <> foldMap (\v -> B.int32BE (fromIntegral (BS.length v)) <> B.byteString v) cols++buildErrorResponse :: ByteString -> ByteString -> Builder+buildErrorResponse sqlState msg =+  buildBackendMsg 0x45 $+    B.word8 (fromIntegral (fromEnum 'S')) <> B.byteString "ERROR" <> B.word8 0+    <> B.word8 (fromIntegral (fromEnum 'C')) <> B.byteString sqlState <> B.word8 0+    <> B.word8 (fromIntegral (fromEnum 'M')) <> B.byteString msg <> B.word8 0+    <> B.word8 0 -- terminator++------------------------------------------------------------------------+-- Socket helpers+------------------------------------------------------------------------++sendBuilder :: NS.Socket -> Builder -> IO ()+sendBuilder sock b = NSB.sendAll sock (LBS.toStrict (B.toLazyByteString b))++recvN :: NS.Socket -> Int -> IO ByteString+recvN sock n = go n []+  where+    go 0 acc = pure (BS.concat (reverse acc))+    go remaining acc = do+      chunk <- NSB.recv sock (min 65536 remaining)+      if BS.null chunk+        then pure (BS.concat (reverse acc))+        else go (remaining - BS.length chunk) (chunk : acc)++recvMaybe :: NS.Socket -> Int -> IO (Maybe ByteString)+recvMaybe sock n = do+  bs <- NSB.recv sock n+  if BS.null bs then pure Nothing else pure (Just bs)++decodeInt32BE :: ByteString -> Int -> Int32+decodeInt32BE bs off =+  let b0 = fromIntegral (BS.index bs off) :: Int32+      b1 = fromIntegral (BS.index bs (off + 1)) :: Int32+      b2 = fromIntegral (BS.index bs (off + 2)) :: Int32+      b3 = fromIntegral (BS.index bs (off + 3)) :: Int32+   in b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3
+ pg-wire.cabal view
@@ -0,0 +1,168 @@+cabal-version:   3.0+name:            pg-wire+version:         0.1.0.0+synopsis:        Pure Haskell PostgreSQL wire protocol driver+description:+  A complete implementation of the PostgreSQL v3 wire protocol in+  pure Haskell. Includes connection management, authentication+  (SCRAM-SHA-256, MD5, cleartext), TLS, connection pooling, and+  binary format type classes. No dependency on libpq or any C library.++license:         BSD-3-Clause+license-file:    LICENSE+author:          Josh Burgess+maintainer:      joshburgess.webdev@gmail.com+category:        Database+homepage:        https://github.com/joshburgess/valiant+bug-reports:     https://github.com/joshburgess/valiant/issues+build-type:      Simple+extra-doc-files:+  README.md+  CHANGELOG.md+tested-with:     GHC ==9.10.3++source-repository head+  type:     git+  location: https://github.com/joshburgess/valiant+  subdir:   wire++flag werror+  description: Enable -Werror for development builds.+  default:     False+  manual:      True++common warnings+  ghc-options: -Wall -Wcompat -Wno-unticked-promoted-constructors -funbox-strict-fields -fspecialise-aggressively+  if flag(werror)+    ghc-options: -Werror++library+  import:           warnings+  hs-source-dirs:   src+  default-language: GHC2021+  default-extensions:+    DerivingStrategies+    LambdaCase+    OverloadedStrings+    RecordWildCards+    StrictData++  exposed-modules:+    PgWire.Async+    PgWire.Auth.Cleartext+    PgWire.Auth.MD5+    PgWire.Auth.ScramSHA256+    PgWire.Binary.Types+    PgWire.Cancel+    PgWire.Connection+    PgWire.Connection.Config+    PgWire.Error+    PgWire.Pool+    PgWire.Pool.Config+    PgWire.Pool.Observation+    PgWire.Protocol.Backend+    PgWire.Protocol.Builders+    PgWire.Protocol.Frontend+    PgWire.Protocol.Oid+    PgWire.Protocol.Parsers+    PgWire.TypeCache+    PgWire.Wire++  build-depends:+    -- Core+    , base                  >=4.17    && <5+    , bytestring            >=0.11    && <0.13+    , text                  >=2.0     && <2.2+    , containers            >=0.6     && <0.8+    , vector                >=0.13    && <0.14+    , time                  >=1.12    && <1.15+    , deepseq               >=1.4     && <1.6++    -- Networking+    , network               >=3.1     && <3.3++    -- TLS+    , tls                   >=1.7     && <2.2+    , crypton-x509-store    >=1.6     && <1.7+    , crypton-x509-system   >=1.6     && <1.7++    -- Crypto (auth)+    , crypton               >=0.34    && <1.1+    , cryptohash-md5        >=0.11    && <0.12+    , memory                >=0.18    && <0.19+    , base64-bytestring     >=1.2     && <1.3++    -- Misc+    , directory             >=1.3     && <1.4+    , hashable              >=1.4     && <1.6+    , nothunks              >=0.1     && <0.3+    , psqueues              >=0.2.8   && <0.3+    , random                >=1.2     && <1.4++    -- Concurrency / pool+    , stm                   >=2.5     && <2.6+    , async                 >=2.2     && <2.3++-- Mock PostgreSQL server for testing client behaviour without a real+-- database. Sub-library so this isn't part of the public pg-wire API.+library mockserver+  import:           warnings+  visibility:       public+  hs-source-dirs:   mockserver+  default-language: GHC2021+  default-extensions:+    DerivingStrategies+    LambdaCase+    OverloadedStrings+    RecordWildCards+    StrictData+  exposed-modules:+    PgWire.MockServer+  build-depends:+    , base                  >=4.17    && <5+    , bytestring            >=0.11    && <0.13+    , network               >=3.1     && <3.3+    , pg-wire++test-suite pg-wire-test+  import:           warnings+  type:             exitcode-stdio-1.0+  hs-source-dirs:   test+  main-is:          Main.hs+  default-language: GHC2021+  default-extensions:+    OverloadedStrings++  other-modules:+    PgWire.Auth.MD5Spec+    PgWire.Auth.ScramFieldsSpec+    PgWire.CancelSpec+    PgWire.Connection.ConfigSpec+    PgWire.Connection.EscapingSpec+    PgWire.Connection.FeaturesSpec+    PgWire.ErrorSpec+    PgWire.MockServerSpec+    PgWire.Pool.ConfigSpec+    PgWire.Pool.NoThunksSpec+    PgWire.Pool.PropertySpec+    PgWire.Pool.StateMachineSpec+    PgWire.Protocol.BuildersSpec+    PgWire.Protocol.FuzzSpec+    PgWire.Protocol.OidSpec+    PgWire.Protocol.ParsersSpec++  ghc-options: -rtsopts "-with-rtsopts=-K8K" -finfo-table-map -fdistinct-constructor-tables+  build-depends:+    , async        >=2.2  && <2.3+    , base         >=4.17 && <5+    , bytestring+    , crypton+    , deepseq+    , hedgehog     >=1.2  && <1.6+    , hspec        >=2.11 && <2.13+    , hspec-hedgehog >=0.1 && <0.3+    , memory+    , nothunks+    , pg-wire+    , pg-wire:mockserver+    , vector
+ src/PgWire/Async.hs view
@@ -0,0 +1,675 @@+{-# OPTIONS_GHC -fno-full-laziness #-}++-- | Asynchronous sender/receiver split for PostgreSQL connections.+--+-- __Audience:__ this module is exposed for downstream library authors+-- (the @valiant@ runtime, custom adapters). Most application code+-- should reach for "PgWire.Connection" instead. The types and+-- functions here are public and stable, but they sit below the+-- request/response abstraction and require knowledge of the v3 wire+-- protocol to use safely.+--+-- Each connection spawns two green threads after startup:+--+-- * __Writer thread__ — drains the send queue, batches messages via+--   vectored I/O ('sendMany'), and enqueues response MVars.+-- * __Reader thread__ — parses backend messages from the socket and+--   fills MVars in FIFO order.+--+-- Application threads submit 'Request's and block on the returned 'MVar'.+-- PostgreSQL processes messages in strict FIFO order, so the i-th response+-- always corresponds to the i-th pending request — no correlation IDs needed.+--+-- This architecture enables automatic pipelining: when N threads query+-- concurrently, their messages are coalesced into a single 'sendMany'+-- syscall by the writer, and responses are demultiplexed by the reader.+module PgWire.Async+  ( -- * Types+    AsyncWireConn (..)+  , Request (..)+  , Response (..)+  , ResponseCollector (..)+    -- * Lifecycle+  , spawnAsyncWireConn+  , shutdownAsyncWireConn+    -- * Submitting requests+  , submitRequest+  , submitExclusive+  ) where++import Control.Concurrent (forkIO)+import Control.Concurrent.Async (Async, async, cancel, link2)+import Control.Concurrent.MVar+import Control.Concurrent.STM+import Control.Exception (SomeException, catch, mask, onException, throwIO, try)+import Control.Monad (void)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Int (Int32, Int64)+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Vector.Mutable qualified as VM+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend+import PgWire.Protocol.Frontend (FrontendMsg (..))+import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsgs)++------------------------------------------------------------------------+-- Types+------------------------------------------------------------------------++-- | What application threads submit to the async connection.+data Request+  = -- | Extended query: pre-built messages (Bind+Execute+Sync) with a collector.+    ReqExtendedQuery ![FrontendMsg] !ResponseCollector+  | -- | Simple text query.+    ReqSimpleQuery !ByteString+  | -- | Parse a prepared statement (Parse+Sync).+    ReqPrepare !FrontendMsg+  | -- | Close a prepared statement (Close+Sync).+    ReqClose !FrontendMsg++-- | What comes back to the caller via MVar.+data Response+  = -- | Rows collected from DataRow messages.+    RespRows ![Vector (Maybe ByteString)]+  | -- | Command tag from CommandComplete.+    RespCommand !CommandTag+  | -- | Simple query result: rows + optional tag.+    RespSimple ![[Maybe ByteString]] !(Maybe CommandTag)+  | -- | ParseComplete acknowledged.+    RespParsed+  | -- | CloseComplete acknowledged.+    RespClosed+  | -- | Multiple sub-results (one per Bind+Execute in a batch).+    RespBatchRows ![[Vector (Maybe ByteString)]]+  | -- | Batch command result: total rows affected.+    RespBatchCommand !Int64+  | -- | Rows collected from DataRow messages AND the command tag.+    RespRowsAndCommand ![Vector (Maybe ByteString)] !CommandTag+  | -- | Rows collected as a Vector (not a list). Used by fetchAllVec.+    RespRowsVec !(Vector (Vector (Maybe ByteString)))+  | -- | At most one row. Used by fetchOne/fetchScalar/fetchExists.+    RespFirstRow !(Maybe (Vector (Maybe ByteString)))++-- | Tells the reader thread how to interpret backend messages for a request.+data ResponseCollector+  = -- | Collect DataRow until CommandComplete, then ReadyForQuery.+    CollectRows+  | -- | Expect CommandComplete + ReadyForQuery, extract row count.+    CollectCommand+  | -- | Expect ParseComplete + ReadyForQuery.+    CollectParseComplete+  | -- | Expect CloseComplete + ReadyForQuery.+    CollectCloseComplete+  | -- | Collect N sub-results (Bind+Execute repeated N times), one ReadyForQuery at end.+    CollectBatch !Int+  | -- | Collect N command sub-results, one ReadyForQuery at end.+    CollectBatchCommand !Int+  | -- | Collect DataRow until CommandComplete, returning both rows and the command tag.+    CollectRowsAndCommand+  | -- | Like CollectRows but collects into a Vector instead of a list.+    CollectRowsVec+  | -- | Collect at most the first row. Discards remaining DataRows.+    CollectFirstRow+  | -- | Simple query protocol: text rows + optional tag + ReadyForQuery.+    CollectSimple++-- | Enqueued by the writer, dequeued by the reader.+data PendingResponse = PendingResponse+  { prCollector :: !ResponseCollector+  , prMVar :: !(MVar (Either PgWireError Response))+  }++-- | The async connection state. Owns the writer and reader threads.+data AsyncWireConn = AsyncWireConn+  { awcWire :: !WireConn+  , awcSendQueue :: !(TBQueue (Request, MVar (Either PgWireError Response)))+  , awcPending :: !(TQueue PendingResponse)+  , awcExclusive :: !(TMVar (MVar (Either PgWireError Response)))+  -- ^ When set, the writer pauses the pipeline for exclusive access.+  , awcAlive :: !(TVar Bool)+  , awcTxStatus :: !(IORef TxStatus)+  , awcNotifyHandler :: !(IORef (Int32 -> ByteString -> ByteString -> IO ()))+  , awcNoticeHandler :: !(IORef (PgNotice -> IO ()))+  , awcParamStatus :: !(IORef (Map ByteString ByteString))+  , awcWriterThread :: !(Async ())+  , awcReaderThread :: !(Async ())+  , awcSendLock :: !(TMVar ())+  -- ^ Fast-path send lock. When the writer is idle, a caller can take+  -- this lock and send directly on the wire, bypassing the TBQueue.+  -- The reader still collects the response normally via the pending queue.+  }++-- | Shared state between writer/reader threads. Doesn't include the+-- thread handles themselves (avoids circular dependency with StrictData).+data AsyncCore = AsyncCore+  { acWire :: !WireConn+  , acSendQueue :: !(TBQueue (Request, MVar (Either PgWireError Response)))+  , acPending :: !(TQueue PendingResponse)+  , acExclusive :: !(TMVar (MVar (Either PgWireError Response)))+  , acAlive :: !(TVar Bool)+  , acTxStatus :: !(IORef TxStatus)+  , acNotifyHandler :: !(IORef (Int32 -> ByteString -> ByteString -> IO ()))+  , acNoticeHandler :: !(IORef (PgNotice -> IO ()))+  , acParamStatus :: !(IORef (Map ByteString ByteString))+  , acSendLock :: !(TMVar ())+  }++------------------------------------------------------------------------+-- Lifecycle+------------------------------------------------------------------------++-- | Spawn writer and reader threads for an established connection.+-- Call this after the serial startup handshake completes.+spawnAsyncWireConn+  :: WireConn+  -> IORef TxStatus+  -> IORef (Map ByteString ByteString)+  -> IO AsyncWireConn+spawnAsyncWireConn wc txRef paramRef = do+  sendQueue <- newTBQueueIO 64+  pending <- newTQueueIO+  exclusive <- newEmptyTMVarIO+  alive <- newTVarIO True+  sendLock <- newTMVarIO ()  -- initially available+  notifyHandler <- newIORef (\_ _ _ -> pure ())+  noticeHandler <- newIORef (\_ -> pure ())++  -- The writer and reader threads need access to the shared STM/IORef state+  -- but NOT to each other's Async handles. We build a "core" with just the+  -- shared state, spawn threads from that, then assemble the final record.+  let core = AsyncCore+        { acWire = wc+        , acSendQueue = sendQueue+        , acPending = pending+        , acExclusive = exclusive+        , acAlive = alive+        , acTxStatus = txRef+        , acNotifyHandler = notifyHandler+        , acNoticeHandler = noticeHandler+        , acParamStatus = paramRef+        , acSendLock = sendLock+        }++  writer <- async (writerThread core)+  reader <- async (readerThread core)++  -- Link threads: if either dies, both die.+  link2 writer reader++  pure AsyncWireConn+    { awcWire = wc+    , awcSendQueue = sendQueue+    , awcPending = pending+    , awcExclusive = exclusive+    , awcAlive = alive+    , awcTxStatus = txRef+    , awcNotifyHandler = notifyHandler+    , awcNoticeHandler = noticeHandler+    , awcParamStatus = paramRef+    , awcWriterThread = writer+    , awcReaderThread = reader+    , awcSendLock = sendLock+    }++-- | Shut down the async connection. Signals threads to stop and cancels them.+shutdownAsyncWireConn :: AsyncWireConn -> IO ()+shutdownAsyncWireConn awc = do+  atomically $ writeTVar (awcAlive awc) False+  cancel (awcWriterThread awc) `catch` \(_ :: SomeException) -> pure ()+  cancel (awcReaderThread awc) `catch` \(_ :: SomeException) -> pure ()+  drainPendingWithError awc ConnectionDead++------------------------------------------------------------------------+-- Submitting requests+------------------------------------------------------------------------++-- | Submit a request and block until the response arrives.+-- Throws 'ConnectionDead' if the async threads have died.+-- | Submit a request and block until the response arrives.+-- Fast path: if the send lock is available (no contention), send directly+-- on the wire without going through the writer thread's TBQueue. The reader+-- still collects the response normally. This eliminates ~2-4μs of thread+-- coordination overhead for single-threaded workloads.+-- Slow path: enqueue to the writer thread (enables automatic pipelining+-- under concurrency).+submitRequest :: AsyncWireConn -> Request -> IO Response+submitRequest awc req = do+  respVar <- newEmptyMVar+  -- Try fast path: take the send lock and send directly.+  -- If the writer is busy (lock unavailable), fall back to the queue.+  path <- atomically $ do+    a <- readTVar (awcAlive awc)+    if not a+      then pure Nothing  -- dead+      else do+        mLock <- tryTakeTMVar (awcSendLock awc)+        case mLock of+          Just () -> pure (Just True)   -- fast path: we have the lock+          Nothing -> do+            writeTBQueue (awcSendQueue awc) (req, respVar)+            pure (Just False)           -- slow path: queued+  case path of+    Nothing -> throwPgWire ConnectionDead+    Just False -> do+      -- Slow path: wait for response from reader via writer+      result <- takeMVar respVar+      case result of+        Left err -> throwIO err+        Right resp -> pure resp+    Just True -> do+      -- Fast path: send directly, enqueue pending for reader+      let (msgs, collector) = requestToMsgs req+          releaseLock = atomically $ putTMVar (awcSendLock awc) ()+      sendFrontendMsgs (awcWire awc) msgs+        `onException` releaseLock+      -- Enqueue pending and release lock atomically so another fast-path+      -- caller cannot send and enqueue before our pending is visible to+      -- the reader — preserving the FIFO invariant.+      atomically $ do+        writeTQueue (awcPending awc) (PendingResponse collector respVar)+        putTMVar (awcSendLock awc) ()+      -- Wait for response from reader+      result <- takeMVar respVar+      case result of+        Left err -> throwIO err+        Right resp -> pure resp++-- | Submit an exclusive request. The writer thread will:+-- 1. Stop accepting new requests+-- 2. Wait for all pending responses to be collected by the reader+-- 3. Hand control to the caller+-- 4. Resume normal operation when the caller finishes+--+-- Used for COPY, cursors, folds — operations that need direct socket access.+submitExclusive :: AsyncWireConn -> (WireConn -> IORef TxStatus -> IO a) -> IO a+submitExclusive awc action = mask $ \restore -> do+  respVar <- newEmptyMVar+  -- Combine alive check with exclusive signal in one STM transaction.+  signaled <- atomically $ do+    a <- readTVar (awcAlive awc)+    if a+      then putTMVar (awcExclusive awc) respVar >> pure True+      else pure False+  if not signaled+    then throwPgWire ConnectionDead+    else do++      -- Wait for writer to drain pending and hand us control+      result <- takeMVar respVar+      case result of+        Left err -> throwIO err+        Right _ -> do+          -- We now have exclusive access to the WireConn.+          -- Reader is blocked on empty awcPending.+          -- Writer is blocked waiting for us to signal done.+          let signalDone = do+                doneVar <- newEmptyMVar+                atomically $ putTMVar (awcExclusive awc) doneVar+                putMVar doneVar (Right RespParsed)+          a <- restore (action (awcWire awc) (awcTxStatus awc))+                 `onException` signalDone+          signalDone+          pure a++-- | Extract messages and collector from a request (used by fast path).+requestToMsgs :: Request -> ([FrontendMsg], ResponseCollector)+requestToMsgs (ReqExtendedQuery msgs collector) = (msgs, collector)+requestToMsgs (ReqSimpleQuery sql) = ([Query sql], CollectSimple)+requestToMsgs (ReqPrepare parseMsg) = ([parseMsg, Flush], CollectParseComplete)+requestToMsgs (ReqClose closeMsg) = ([closeMsg, Sync], CollectCloseComplete)++------------------------------------------------------------------------+-- Writer thread+------------------------------------------------------------------------++data WriterAction+  = WriterExclusive !(MVar (Either PgWireError Response))+  | WriterBatch ![(Request, MVar (Either PgWireError Response))]++writerThread :: AsyncCore -> IO ()+writerThread ac = go `catch` onDeath+  where+    go :: IO ()+    go = do+      action <- atomically $ tryExclusive `orElse` drainQueue+      case action of+        WriterExclusive respVar -> handleExclusive respVar+        WriterBatch items -> handleBatch items+      go++    tryExclusive :: STM WriterAction+    tryExclusive = do+      respVar <- takeTMVar (acExclusive ac)+      pure (WriterExclusive respVar)++    drainQueue :: STM WriterAction+    drainQueue = do+      first <- readTBQueue (acSendQueue ac)+      -- Take the send lock so fast-path callers can't send while we're sending.+      takeTMVar (acSendLock ac)+      rest <- drainTBQueue (acSendQueue ac)+      pure (WriterBatch (first : rest))++    handleBatch :: [(Request, MVar (Either PgWireError Response))] -> IO ()+    handleBatch items = do+      let (allMsgs, pendings) = unzip (map buildItem items)+      atomically $ mapM_ (writeTQueue (acPending ac)) pendings+      sendFrontendMsgs (acWire ac) (concat allMsgs)+      -- Release the send lock so fast-path callers can send again.+      atomically $ putTMVar (acSendLock ac) ()++    buildItem :: (Request, MVar (Either PgWireError Response)) -> ([FrontendMsg], PendingResponse)+    buildItem (req, respVar) = case req of+      ReqExtendedQuery msgs collector ->+        (msgs, PendingResponse collector respVar)+      ReqSimpleQuery sql ->+        ([Query sql], PendingResponse CollectSimple respVar)+      ReqPrepare parseMsg ->+        ([parseMsg, Flush], PendingResponse CollectParseComplete respVar)+      ReqClose closeMsg ->+        ([closeMsg, Sync], PendingResponse CollectCloseComplete respVar)++    handleExclusive :: MVar (Either PgWireError Response) -> IO ()+    handleExclusive respVar = do+      atomically $ do+        empty <- isEmptyTQueue (acPending ac)+        check empty+      putMVar respVar (Right RespParsed)+      doneVar <- atomically $ takeTMVar (acExclusive ac)+      _ <- takeMVar doneVar+      pure ()++    onDeath :: SomeException -> IO ()+    onDeath _ = do+      atomically $ writeTVar (acAlive ac) False+      drainPendingWithErrorCore ac ConnectionDead++------------------------------------------------------------------------+-- Reader thread+------------------------------------------------------------------------++readerThread :: AsyncCore -> IO ()+readerThread ac = go `catch` onDeath+  where+    go :: IO ()+    go = do+      pr <- atomically $ readTQueue (acPending ac)+      -- Catch query errors so they don't kill the reader thread.+      -- A QueryError is per-request, not connection-fatal.+      result <- try (collectResponse (prCollector pr))+      case result of+        Right resp -> putMVar (prMVar pr) (Right resp)+        Left (err :: PgWireError) -> do+          -- After ErrorResponse with Sync, Postgres sends ReadyForQuery.+          -- After ErrorResponse with Flush (preparation), it does NOT.+          case prCollector pr of+            CollectParseComplete -> pure ()+            _ -> drainUntilReady+          putMVar (prMVar pr) (Left err)+      go++    -- Drain messages until ReadyForQuery (error recovery).+    drainUntilReady :: IO ()+    drainUntilReady = do+      msg <- recvAndDispatch+      case msg of+        ReadyForQuery status -> writeIORef (acTxStatus ac) status+        _ -> drainUntilReady++    collectResponse :: ResponseCollector -> IO Response+    collectResponse CollectRows = do+      rows <- collectRowsLoop+      waitReadyForQuery+      pure (RespRows rows)+    collectResponse CollectCommand = do+      tag <- collectCommandLoop+      waitReadyForQuery+      pure (RespCommand tag)+    collectResponse CollectParseComplete = do+      waitParseCompleteLoop+      -- No waitReadyForQuery: Flush doesn't trigger ReadyForQuery.+      pure RespParsed+    collectResponse CollectCloseComplete = do+      waitCloseCompleteLoop+      waitReadyForQuery+      pure RespClosed+    collectResponse CollectRowsAndCommand = do+      (rows, tag) <- collectRowsAndCommandLoop+      waitReadyForQuery+      pure (RespRowsAndCommand rows tag)+    collectResponse CollectRowsVec = do+      vec <- collectRowsVecLoop+      waitReadyForQuery+      pure (RespRowsVec vec)+    collectResponse CollectFirstRow = do+      mRow <- collectFirstRowLoop+      waitReadyForQuery+      pure (RespFirstRow mRow)+    collectResponse (CollectBatch n) = do+      results <- collectBatchLoop n+      waitReadyForQuery+      pure (RespBatchRows results)+    collectResponse (CollectBatchCommand n) = do+      total <- collectBatchCommandLoop n 0+      waitReadyForQuery+      pure (RespBatchCommand total)+    collectResponse CollectSimple = do+      (rows, tag) <- collectSimpleLoop [] Nothing+      pure (RespSimple rows tag)++    -- Collectors skip both ParseComplete and BindComplete so that+    -- Parse+Bind+Execute can be coalesced into a single request.+    collectRowsLoop :: IO [Vector (Maybe ByteString)]+    collectRowsLoop = loop id+      where+        loop !acc = do+          msg <- recvAndDispatch+          case msg of+            ParseComplete -> loop acc+            BindComplete -> loop acc+            DataRow vals -> loop (acc . (vals :))+            CommandComplete _ -> pure (acc [])+            EmptyQueryResponse -> pure (acc [])+            ErrorResponse err -> throwIO (QueryError err)+            other -> throwIO (ProtocolError ("Unexpected in rows: " <> BS8.pack (show other)))++    -- | Collect rows into a growable mutable vector. Starts at capacity 64,+    -- doubles when full. Final freeze+slice produces an exact-size immutable Vector.+    collectRowsVecLoop :: IO (Vector (Vector (Maybe ByteString)))+    collectRowsVecLoop = do+      mv <- VM.new 64+      (finalMv, !n) <- loop mv 0+      V.unsafeFreeze (VM.slice 0 n finalMv)+      where+        loop !mv !i = do+          msg <- recvAndDispatch+          case msg of+            ParseComplete -> loop mv i+            BindComplete -> loop mv i+            DataRow vals -> do+              mv' <- if i >= VM.length mv+                then VM.grow mv (VM.length mv) -- double capacity+                else pure mv+              VM.write mv' i vals+              loop mv' (i + 1)+            CommandComplete _ -> pure (mv, i)+            EmptyQueryResponse -> pure (mv, i)+            ErrorResponse err -> throwIO (QueryError err)+            other -> throwIO (ProtocolError ("Unexpected in rows vec: " <> BS8.pack (show other)))++    -- | Collect at most the first row, discard the rest.+    collectFirstRowLoop :: IO (Maybe (Vector (Maybe ByteString)))+    collectFirstRowLoop = loop+      where+        loop = do+          msg <- recvAndDispatch+          case msg of+            ParseComplete -> loop+            BindComplete -> loop+            DataRow vals -> do+              -- Got first row, now drain remaining rows+              drainRows+              pure (Just vals)+            CommandComplete _ -> pure Nothing+            EmptyQueryResponse -> pure Nothing+            ErrorResponse err -> throwIO (QueryError err)+            other -> throwIO (ProtocolError ("Unexpected in first row: " <> BS8.pack (show other)))+        drainRows = do+          msg <- recvAndDispatch+          case msg of+            DataRow _ -> drainRows -- discard+            CommandComplete _ -> pure ()+            EmptyQueryResponse -> pure ()+            NoticeResponse _ -> drainRows+            ErrorResponse err -> throwIO (QueryError err)+            other -> throwIO (ProtocolError ("Unexpected draining rows: " <> BS8.pack (show other)))++    collectRowsAndCommandLoop :: IO ([Vector (Maybe ByteString)], CommandTag)+    collectRowsAndCommandLoop = loop id+      where+        loop !acc = do+          msg <- recvAndDispatch+          case msg of+            ParseComplete -> loop acc+            BindComplete -> loop acc+            DataRow vals -> loop (acc . (vals :))+            CommandComplete tag -> pure (acc [], tag)+            EmptyQueryResponse -> pure (acc [], OtherTag "EMPTY")+            ErrorResponse err -> throwIO (QueryError err)+            other -> throwIO (ProtocolError ("Unexpected in rows+command: " <> BS8.pack (show other)))++    collectCommandLoop :: IO CommandTag+    collectCommandLoop = do+      msg <- recvAndDispatch+      case msg of+        ParseComplete -> collectCommandLoop+        BindComplete -> collectCommandLoop+        CommandComplete tag -> pure tag+        EmptyQueryResponse -> pure (OtherTag "EMPTY")+        ErrorResponse err -> throwIO (QueryError err)+        other -> throwIO (ProtocolError ("Unexpected in command: " <> BS8.pack (show other)))++    waitParseCompleteLoop :: IO ()+    waitParseCompleteLoop = do+      msg <- recvAndDispatch+      case msg of+        ParseComplete -> pure ()+        ErrorResponse err -> throwIO (QueryError err)+        other -> throwIO (ProtocolError ("Expected ParseComplete, got: " <> BS8.pack (show other)))++    waitCloseCompleteLoop :: IO ()+    waitCloseCompleteLoop = do+      msg <- recvAndDispatch+      case msg of+        CloseComplete -> pure ()+        ErrorResponse err -> throwIO (QueryError err)+        other -> throwIO (ProtocolError ("Expected CloseComplete, got: " <> BS8.pack (show other)))++    collectBatchLoop :: Int -> IO [[Vector (Maybe ByteString)]]+    collectBatchLoop 0 = pure []+    collectBatchLoop n = do+      rows <- collectRowsLoop+      rest <- collectBatchLoop (n - 1)+      pure (rows : rest)++    collectBatchCommandLoop :: Int -> Int64 -> IO Int64+    collectBatchCommandLoop 0 !total = pure total+    collectBatchCommandLoop n !total = do+      msg <- recvAndDispatch+      case msg of+        ParseComplete -> collectBatchCommandLoop n total+        BindComplete -> collectBatchCommandLoop n total+        CommandComplete tag -> collectBatchCommandLoop (n - 1) (total + tagRows tag)+        ErrorResponse err -> throwIO (QueryError err)+        other -> throwIO (ProtocolError ("Unexpected in batch cmd: " <> BS8.pack (show other)))++    collectSimpleLoop :: [[Maybe ByteString]] -> Maybe CommandTag -> IO ([[Maybe ByteString]], Maybe CommandTag)+    collectSimpleLoop !rows !tag = do+      msg <- recvAndDispatch+      case msg of+        RowDescription _ -> collectSimpleLoop rows tag+        DataRow vals -> collectSimpleLoop (V.toList vals : rows) tag+        CommandComplete ct -> collectSimpleLoop rows (Just ct)+        EmptyQueryResponse -> collectSimpleLoop rows tag+        ReadyForQuery status -> do+          writeIORef (acTxStatus ac) status+          pure (reverse rows, tag)+        ErrorResponse err -> throwIO (QueryError err)+        other -> throwIO (ProtocolError ("Unexpected in simple query: " <> BS8.pack (show other)))++    waitReadyForQuery :: IO ()+    waitReadyForQuery = do+      msg <- recvAndDispatch+      case msg of+        ReadyForQuery status -> writeIORef (acTxStatus ac) status+        _ -> waitReadyForQuery++    recvAndDispatch :: IO BackendMsg+    recvAndDispatch = do+      msg <- recvBackendMsg (acWire ac)+      case msg of+        NotificationResponse pid channel payload -> do+          handler <- readIORef (acNotifyHandler ac)+          _ <- forkIO (handler pid channel payload `catch` \(_ :: SomeException) -> pure ())+          recvAndDispatch+        NoticeResponse notice -> do+          handler <- readIORef (acNoticeHandler ac)+          _ <- forkIO (handler notice `catch` \(_ :: SomeException) -> pure ())+          recvAndDispatch+        ParameterStatus key value -> do+          modifyIORef' (acParamStatus ac) (Map.insert key value)+          recvAndDispatch+        _ -> pure msg++    onDeath :: SomeException -> IO ()+    onDeath _ = do+      atomically $ writeTVar (acAlive ac) False+      drainPendingWithErrorCore ac ConnectionDead++------------------------------------------------------------------------+-- Helpers+------------------------------------------------------------------------++drainTBQueue :: TBQueue a -> STM [a]+drainTBQueue q = loop id+  where+    loop acc = do+      mItem <- tryReadTBQueue q+      case mItem of+        Just item -> loop (acc . (item :))+        Nothing -> pure (acc [])++drainPendingWithError :: AsyncWireConn -> PgWireError -> IO ()+drainPendingWithError awc err = do+  pendings <- atomically $ flushTQueue (awcPending awc)+  mapM_ (\pr -> void (tryPutMVar (prMVar pr) (Left err))) pendings+  queued <- atomically $ drainTBQueue (awcSendQueue awc)+  mapM_ (\(_, mv) -> void (tryPutMVar mv (Left err))) queued+  mExcl <- atomically $ tryTakeTMVar (awcExclusive awc)+  case mExcl of+    Just mv -> void (tryPutMVar mv (Left err))+    Nothing -> pure ()++-- | Like 'drainPendingWithError' but operates on 'AsyncCore' (used by threads).+drainPendingWithErrorCore :: AsyncCore -> PgWireError -> IO ()+drainPendingWithErrorCore ac err = do+  pendings <- atomically $ flushTQueue (acPending ac)+  mapM_ (\pr -> void (tryPutMVar (prMVar pr) (Left err))) pendings+  queued <- atomically $ drainTBQueue (acSendQueue ac)+  mapM_ (\(_, mv) -> void (tryPutMVar mv (Left err))) queued+  mExcl <- atomically $ tryTakeTMVar (acExclusive ac)+  case mExcl of+    Just mv -> void (tryPutMVar mv (Left err))+    Nothing -> pure ()++tagRows :: CommandTag -> Int64+tagRows (InsertTag n) = n+tagRows (UpdateTag n) = n+tagRows (DeleteTag n) = n+tagRows (SelectTag n) = n+tagRows (OtherTag _) = 0
+ src/PgWire/Auth/Cleartext.hs view
@@ -0,0 +1,17 @@+-- | Cleartext password authentication.+--+-- __Audience:__ exposed for downstream library authors. The standard+-- authentication flow is driven from "PgWire.Connection"; this module+-- is only useful if you are reimplementing handshake.+module PgWire.Auth.Cleartext+  ( cleartextAuth+  ) where++import Data.ByteString (ByteString)+import PgWire.Protocol.Frontend (FrontendMsg (..))+import PgWire.Wire (WireConn, sendFrontendMsg)++-- | Handle cleartext password authentication.+cleartextAuth :: WireConn -> ByteString -> IO ()+cleartextAuth wc password =+  sendFrontendMsg wc (PasswordMessage password)
+ src/PgWire/Auth/MD5.hs view
@@ -0,0 +1,39 @@+-- | MD5 password authentication (legacy PostgreSQL auth method).+--+-- __Audience:__ exposed for downstream library authors. The standard+-- authentication flow is driven from "PgWire.Connection"; this module+-- is only useful if you are reimplementing handshake.+module PgWire.Auth.MD5+  ( md5Auth+  , md5Password+  ) where++import Crypto.Hash.MD5 qualified as MD5+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Word (Word8)+import PgWire.Protocol.Frontend (FrontendMsg (..))+import PgWire.Wire (WireConn, sendFrontendMsg)+import Numeric (showHex)++-- | Handle MD5 password authentication.+-- PG expects: @"md5" ++ md5(md5(password ++ user) ++ salt)@+md5Auth :: WireConn -> ByteString -> ByteString -> ByteString -> IO ()+md5Auth wc user password salt =+  sendFrontendMsg wc (PasswordMessage (md5Password user password salt))++-- | Compute the MD5 password hash.+md5Password :: ByteString -> ByteString -> ByteString -> ByteString+md5Password user password salt =+  let inner = hexMd5 (password <> user)+      outer = hexMd5 (inner <> salt)+   in "md5" <> outer++hexMd5 :: ByteString -> ByteString+hexMd5 = BS.pack . concatMap toHexBytes . BS.unpack . MD5.hash++toHexBytes :: Word8 -> [Word8]+toHexBytes w =+  let s = showHex w ""+      padded = if length s == 1 then '0' : s else s+   in map (fromIntegral . fromEnum) padded
+ src/PgWire/Auth/ScramSHA256.hs view
@@ -0,0 +1,163 @@+-- | SCRAM-SHA-256 authentication (RFC 5802 / RFC 7677), with optional+-- @tls-server-end-point@ channel binding.+--+-- __Audience:__ exposed for downstream library authors. The standard+-- authentication flow is driven from "PgWire.Connection"; this module+-- is only useful if you are reimplementing handshake.+module PgWire.Auth.ScramSHA256+  ( scramAuth+  , scramAuthWithChannelBinding+  ) where++import Control.Monad (unless)+import Text.Read (readMaybe)+import Crypto.Hash (SHA256 (..), hashWith)+import Crypto.KDF.PBKDF2 qualified as PBKDF2+import Crypto.MAC.HMAC (HMAC, hmac)+import Crypto.Random (getRandomBytes)+import Data.ByteArray qualified as BA+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Base64 qualified as B64+import Data.ByteString.Char8 qualified as BS8+import Data.Bits (xor)+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend (AuthType (..), BackendMsg (..), PgError (..))+import PgWire.Protocol.Frontend (FrontendMsg (..))+import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsg)++-- | Perform SCRAM-SHA-256 authentication without channel binding.+scramAuth :: WireConn -> ByteString -> ByteString -> IO ()+scramAuth wc user password = scramAuthInternal wc user password Nothing++-- | Perform SCRAM-SHA-256-PLUS authentication with tls-server-end-point+-- channel binding. The ByteString argument is the SHA-256 hash of the+-- server's TLS certificate (DER encoded).+scramAuthWithChannelBinding :: WireConn -> ByteString -> ByteString -> ByteString -> IO ()+scramAuthWithChannelBinding wc user password certHash =+  scramAuthInternal wc user password (Just certHash)++scramAuthInternal :: WireConn -> ByteString -> ByteString -> Maybe ByteString -> IO ()+scramAuthInternal wc user password mCertHash = do+  nonceBytes <- getRandomBytes 18 :: IO ByteString+  let clientNonce = B64.encode nonceBytes++      -- Channel binding header:+      --   "n,," for no channel binding+      --   "p=tls-server-end-point,," for tls-server-end-point+      (gs2Header, cbData, mechName) = case mCertHash of+        Nothing ->+          ("n,,", "n,,", "SCRAM-SHA-256")+        Just certHash ->+          let hdr = "p=tls-server-end-point,,"+           in (hdr, hdr <> certHash, "SCRAM-SHA-256-PLUS")++      clientFirstBare = "n=" <> user <> ",r=" <> clientNonce+      clientFirstMsg = gs2Header <> clientFirstBare++  -- Send SASL initial response+  sendFrontendMsg wc (SASLInitialResponse mechName clientFirstMsg)++  -- Receive server first message+  msg1 <- recvBackendMsg wc+  serverFirstMsg <- case msg1 of+    Authentication (AuthSASLContinue serverData) -> pure serverData+    Authentication AuthOk -> pure ""+    ErrorResponse err -> throwPgWire (AuthError (pgMessage err))+    other -> throwPgWire (AuthError ("Unexpected message during SCRAM: " <> BS8.pack (show other)))++  -- Parse server first message: r=<nonce>,s=<salt>,i=<iterations>+  let serverFields = parseScramFields serverFirstMsg+  serverNonce <- lookupField "r" serverFields+  saltB64 <- lookupField "s" serverFields+  iterStr <- lookupField "i" serverFields++  iterations <- case readMaybe (BS8.unpack iterStr) :: Maybe Int of+    Just n | n > 0 -> pure n+    _ -> throwPgWire (AuthError ("Bad SCRAM iteration count: " <> iterStr))+  salt <- case B64.decode saltB64 of+    Left err -> throwPgWire (AuthError ("Bad salt base64: " <> BS8.pack err))+    Right s -> pure s++  -- Verify server nonce starts with our client nonce+  unless (clientNonce `BS.isPrefixOf` serverNonce) $+    throwPgWire (AuthError "Server nonce doesn't start with client nonce")++  -- Compute proofs+  let saltedPassword = hi password salt iterations+      clientKey = hmacSHA256 saltedPassword "Client Key"+      storedKey = hashSHA256 clientKey+      serverKey = hmacSHA256 saltedPassword "Server Key"++      channelBinding = B64.encode cbData+      clientFinalWithoutProof = "c=" <> channelBinding <> ",r=" <> serverNonce+      authMessage = clientFirstBare <> "," <> serverFirstMsg <> "," <> clientFinalWithoutProof++      clientSignature = hmacSHA256 storedKey authMessage+      clientProof = BS.pack (BS.zipWith xor clientKey clientSignature)+      serverSignature = hmacSHA256 serverKey authMessage++      clientFinalMsg = clientFinalWithoutProof <> ",p=" <> B64.encode clientProof++  -- Send client final message+  sendFrontendMsg wc (SASLResponse clientFinalMsg)++  -- Receive server final message+  msg2 <- recvBackendMsg wc+  case msg2 of+    Authentication (AuthSASLFinal serverData) -> do+      let sFields = parseScramFields serverData+      case lookup "v" sFields of+        Nothing -> throwPgWire (AuthError "Missing server signature in SASL final")+        Just sigB64 -> case B64.decode sigB64 of+          Left err -> throwPgWire (AuthError ("Bad server sig base64: " <> BS8.pack err))+          Right sig+            | not (BA.constEq sig serverSignature) ->+                throwPgWire (AuthError "Server signature mismatch")+            | otherwise -> pure ()+    Authentication AuthOk -> pure ()+    ErrorResponse err -> throwPgWire (AuthError (pgMessage err))+    other -> throwPgWire (AuthError ("Unexpected message during SCRAM final: " <> BS8.pack (show other)))++  -- Wait for AuthOk+  msg3 <- recvBackendMsg wc+  case msg3 of+    Authentication AuthOk -> pure ()+    ErrorResponse err -> throwPgWire (AuthError (pgMessage err))+    _ -> throwPgWire (AuthError "Expected AuthOk after SCRAM")++-- Crypto helpers ----------------------------------------------------------++hmacSHA256 :: ByteString -> ByteString -> ByteString+hmacSHA256 key msg = BA.convert (hmac key msg :: HMAC SHA256)+{-# INLINE hmacSHA256 #-}++hashSHA256 :: ByteString -> ByteString+hashSHA256 bs = BA.convert (hashWith SHA256 bs)+{-# INLINE hashSHA256 #-}++-- | PBKDF2-SHA256 (called "Hi" in SCRAM spec).+hi :: ByteString -> ByteString -> Int -> ByteString+hi password salt iterations =+  PBKDF2.generate+    (PBKDF2.prfHMAC SHA256)+    (PBKDF2.Parameters iterations 32)+    password+    salt++-- SCRAM field parsing -----------------------------------------------------++parseScramFields :: ByteString -> [(ByteString, ByteString)]+parseScramFields bs =+  [ (key, val)+  | field <- BS8.split ',' bs+  , let (key, rest) = BS.break (== 61) field -- '='+  , not (BS.null rest)+  , let val = BS.drop 1 rest+  ]++lookupField :: ByteString -> [(ByteString, ByteString)] -> IO ByteString+lookupField key fields =+  case lookup key fields of+    Nothing -> throwPgWire (AuthError ("Missing SCRAM field: " <> key))+    Just v -> pure v
+ src/PgWire/Binary/Types.hs view
@@ -0,0 +1,25 @@+-- | Type classes for encoding and decoding Haskell values to and from+-- PostgreSQL binary format.+--+-- Instances are provided for common types ('Int32', 'Text', 'Bool',+-- 'UTCTime', etc.) in "Valiant.Binary.Encode" and "Valiant.Binary.Decode".+-- Custom types can provide their own instances.+module PgWire.Binary.Types+  ( PgEncode (..)+  , PgDecode (..)+  ) where++import Data.ByteString (ByteString)+import Data.Proxy (Proxy)+import PgWire.Protocol.Oid (Oid)++-- | Encode a Haskell value to PostgreSQL binary format.+class PgEncode a where+  pgEncode :: a -> ByteString+  pgOid :: Proxy a -> Oid+  {-# MINIMAL pgEncode, pgOid #-}++-- | Decode a Haskell value from PostgreSQL binary format.+class PgDecode a where+  pgDecode :: ByteString -> Either String a+  {-# MINIMAL pgDecode #-}
+ src/PgWire/Cancel.hs view
@@ -0,0 +1,88 @@+-- | Query cancellation via the CancelRequest protocol.+--+-- PostgreSQL's cancel mechanism works by opening a separate TCP connection+-- and sending a CancelRequest message containing the backend PID and+-- secret key from the original connection's BackendKeyData. The server+-- then signals the backend process to cancel its current query.+--+-- @+-- -- Cancel a long-running query from another thread:+-- 'cancelQuery' conn+--+-- -- Or use a timeout wrapper:+-- 'withQueryTimeout' conn 5.0 $ do+--   fetchAll conn expensiveQuery ()+-- @+module PgWire.Cancel+  ( cancelQuery+  , withQueryTimeout+  ) where++import Control.Concurrent.Async (race)+import Control.Concurrent (threadDelay)+import Control.Exception (SomeException, try)+import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int32)+import Data.Time (NominalDiffTime)+import Network.Socket qualified as NS+import Network.Socket.ByteString qualified as NSB+import PgWire.Connection (Connection (..))+import PgWire.Connection.Config (ConnConfig (..))+import PgWire.Error (PgWireError (..), throwPgWire)++-- | Cancel the currently executing query on a connection.+--+-- Opens a separate TCP connection to the same Postgres server and sends+-- a CancelRequest message. The original query will receive an ErrorResponse+-- with SQLSTATE 57014 (query_canceled).+--+-- This is safe to call from any thread.+cancelQuery :: Connection -> IO ()+cancelQuery conn = do+  let cfg = connConfig conn+      pid = connBackendPid conn+      key = connBackendKey conn+  result <- try @SomeException $ sendCancelRequest (ccHost cfg) (ccPort cfg) pid key+  case result of+    Right () -> pure ()+    Left _ -> pure () -- Best-effort: cancel may fail if server is unreachable++-- | Run an IO action with a timeout. If the timeout expires, the query+-- is cancelled and a 'PoolTimeout' error is thrown.+--+-- > withQueryTimeout conn 5.0 $ do+-- >   fetchAll conn someExpensiveQuery ()+withQueryTimeout :: Connection -> NominalDiffTime -> IO a -> IO a+withQueryTimeout conn seconds action = do+  let micros = round (seconds * 1000000) :: Int+  result <- race (threadDelay micros) action+  case result of+    Left () -> do+      cancelQuery conn+      throwPgWire (ConnectionError "Query timed out")+    Right a -> pure a++-- | Send a CancelRequest to the given host/port.+--+-- CancelRequest format (not a normal frontend message):+--   [length :: Int32 = 16]+--   [cancel code :: Int32 = 80877102]+--   [pid :: Int32]+--   [key :: Int32]+sendCancelRequest :: NS.HostName -> NS.PortNumber -> Int32 -> Int32 -> IO ()+sendCancelRequest host port pid key = do+  let hints = NS.defaultHints {NS.addrSocketType = NS.Stream}+  addrs <- NS.getAddrInfo (Just hints) (Just host) (Just (show port))+  case addrs of+    [] -> throwPgWire (ConnectionError "Cancel: no address found")+    (addr : _) -> do+      sock <- NS.socket (NS.addrFamily addr) NS.Stream NS.defaultProtocol+      NS.connect sock (NS.addrAddress addr)+      let msg = LBS.toStrict . B.toLazyByteString $+            B.int32BE 16+              <> B.int32BE 80877102+              <> B.int32BE pid+              <> B.int32BE key+      NSB.sendAll sock msg+      NS.close sock
+ src/PgWire/Connection.hs view
@@ -0,0 +1,598 @@+-- | PostgreSQL connection management.+--+-- A 'Connection' represents an active TCP (or TLS) session with a+-- PostgreSQL server. It handles authentication, parameter negotiation,+-- and prepared statement caching.+--+-- After startup, each connection spawns dedicated writer and reader+-- threads (see 'PgWire.Async') that enable automatic pipelining when+-- multiple green threads share a connection.+--+-- For production use, prefer 'PgWire.Pool' over direct connections.+--+-- @+-- conn <- 'connectString' \"postgres:\/\/user:pass\@localhost:5432\/mydb\"+-- (rows, _) <- 'simpleQuery' conn \"SELECT 1\"+-- 'close' conn+-- @+module PgWire.Connection+  ( -- * Connection type+    Connection (..)+  , connWire+    -- * Connecting+  , connect+  , connectString+  , reset+    -- * Closing+  , close+  , withConnection+    -- * Simple queries+  , simpleQuery+    -- * Connection status+  , connectionStatus+  , transactionStatus+  , parameterStatus+  , serverVersion+  , backendPid+  , isSslInUse+    -- * Notifications (non-blocking)+  , checkNotification+    -- * Notice handling+  , setNoticeHandler+    -- * SQL escaping+  , escapeLiteral+  , escapeIdentifier+    -- * Statement introspection+  , describePrepared+  , ParamDescription (..)+  , ColumnDescription (..)+    -- * Server ping+  , ping+    -- * Password utilities+  , lookupPgpass+  , encryptPassword+    -- * Protocol tracing+  , setTraceHandler+  ) where++import Control.Exception (SomeException, bracket, catch, try)+import Crypto.Hash (MD5, hash, Digest)+import Data.ByteArray qualified as BA+import Data.ByteString qualified as BS+import System.Directory (doesFileExist)+import System.Environment (lookupEnv)+import System.Random (randomRIO)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.HashPSQ (HashPSQ)+import Data.HashPSQ qualified as PSQ+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Int (Int32)+import Data.Vector qualified as V+import Data.Vector (Vector)+import Data.Word (Word32, Word64)+import PgWire.Async (AsyncWireConn (..), Request (..), Response (..), spawnAsyncWireConn, shutdownAsyncWireConn, submitRequest, submitExclusive)+import PgWire.Auth.Cleartext (cleartextAuth)+import PgWire.Auth.MD5 (md5Auth)+import PgWire.Auth.ScramSHA256 (scramAuth)+import PgWire.Connection.Config (ConnConfig (..), TargetSessionAttrs (..), TlsMode (..), parseConnString, parseHosts)+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend+import PgWire.Protocol.Builders (buildStartup)+import PgWire.Protocol.Frontend (DescribeTarget (..), FrontendMsg (..), StartupParams (..))+import PgWire.Wire (TlsConfig (..), TraceDirection (..), WireConn (..), connectTcpTimeout, recvBackendMsg, sendFrontendMsg, sendRawBytes, upgradeTls)++-- | A connection to a PostgreSQL database.+--+-- After the initial handshake, dedicated writer and reader threads handle+-- all wire I/O. Use 'submitRequest' or the higher-level functions in+-- "Valiant.Execute" to interact with the connection.+data Connection = Connection+  { connAsync :: !AsyncWireConn+  -- ^ The async wire connection (writer + reader threads).+  , connConfig :: !ConnConfig+  , connBackendPid :: {-# UNPACK #-} !Int32+  , connBackendKey :: {-# UNPACK #-} !Int32+  , connStmtCache :: !(IORef (HashPSQ ByteString Word64 ByteString))+  -- ^ LRU prepared statement cache. Keys are SQL text, priorities are+  -- monotonic access ticks (lower = least recently used), values are+  -- server-side statement names.+  , connStmtCounter :: !(IORef Word64)+  , connStmtTick :: !(IORef Word64)+  -- ^ Monotonic tick counter for LRU cache priorities.+  , connSslActive :: !Bool+  , connPreparedStatements :: !Bool+  -- ^ Whether to use named prepared statements (default: True).+  -- Set to False for compatibility with PgBouncer in transaction mode.+  -- When False, all queries use the unnamed statement (parsed + discarded+  -- each time, never cached).+  , connMaxPreparedStatements :: !Int+  -- ^ Maximum number of prepared statements cached on this connection.+  }++-- | Access the underlying 'WireConn' for direct wire operations.+-- Only use this inside exclusive mode ('submitExclusive').+connWire :: Connection -> WireConn+connWire = awcWire . connAsync++-- | Access the transaction status IORef.+connTxStatus :: Connection -> IORef TxStatus+connTxStatus = awcTxStatus . connAsync++-- | Access the server parameters IORef.+connParams :: Connection -> IORef (Map ByteString ByteString)+connParams = awcParamStatus . connAsync++-- | Access the notice handler IORef.+connNoticeHandler :: Connection -> IORef (PgNotice -> IO ())+connNoticeHandler = awcNoticeHandler . connAsync++-- | Parameter type information returned by 'describePrepared'.+data ParamDescription = ParamDescription+  { pdOids :: Vector Word32+  -- ^ OIDs of each parameter's type (one per @$N@ placeholder).+  }+  deriving stock (Show, Eq)++-- | Column metadata returned by 'describePrepared'.+data ColumnDescription = ColumnDescription+  { cdFields :: Vector FieldInfo+  -- ^ Name, type OID, table OID, column number, and format for each result column.+  }+  deriving stock (Show, Eq)++-- | Connect using a 'ConnConfig'. Supports multi-host failover:+-- if @ccHost@ contains comma-separated hosts (@\"host1,host2\"@),+-- tries each in order until one succeeds. If @ccTargetSessionAttrs@+-- is set, verifies the server matches (e.g., primary vs standby).+connect :: ConnConfig -> IO Connection+connect cfg0 = do+  -- Auto-lookup password from ~/.pgpass if not provided+  cfg <- if BS8.null (ccPassword cfg0)+    then do+      mPw <- lookupPgpass cfg0+      pure $ case mPw of+        Just pw -> cfg0 { ccPassword = pw }+        Nothing -> cfg0+    else pure cfg0+  let hosts = parseHosts (ccHost cfg)+  orderedHosts <- if ccLoadBalanceHosts cfg+    then shuffleHosts hosts+    else pure hosts+  tryHosts cfg orderedHosts++tryHosts :: ConnConfig -> [String] -> IO Connection+tryHosts _ [] = throwPgWire (ConnectionError "All hosts failed")+tryHosts cfg [h] = connectSingleHost cfg { ccHost = h }+tryHosts cfg (h : hs) = do+  result <- try @SomeException (connectSingleHost cfg { ccHost = h })+  case result of+    Right conn -> do+      ok <- checkSessionAttrs conn (ccTargetSessionAttrs cfg)+      if ok+        then pure conn+        else do+          close conn+          tryHosts cfg hs+    Left _ -> tryHosts cfg hs++checkSessionAttrs :: Connection -> TargetSessionAttrs -> IO Bool+checkSessionAttrs _ SessionAny = pure True+checkSessionAttrs conn attr = do+  (rows, _) <- simpleQuery conn "SHOW transaction_read_only"+  case rows of+    [[Just val]] ->+      let readOnly = val == "on"+       in pure $ case attr of+            SessionReadWrite -> not readOnly+            SessionReadOnly -> readOnly+            SessionPrimary -> not readOnly+            SessionStandby -> readOnly+            SessionPreferStandby -> True+    _ -> pure True++-- | Fisher-Yates shuffle for load balancing+shuffleHosts :: [a] -> IO [a]+shuffleHosts [] = pure []+shuffleHosts [x] = pure [x]+shuffleHosts xs = do+  let arr = zip [0 :: Int ..] xs+      n = length xs+  go (reverse arr) (n - 1)+  where+    go [] _ = pure []+    go ((_, x) : _rest) 0 = pure [x]+    go items i = do+      j <- randomRIO (0, i)+      let (before, rest) = splitAt j items+      case rest of+        []                   -> pure (map snd before)  -- unreachable: j <= length items - 1+        (_, picked) : after  -> (picked :) <$> go (before ++ after) (i - 1)++connectSingleHost :: ConnConfig -> IO Connection+connectSingleHost cfg = do+  wc0 <- connectTcpTimeout (ccConnectTimeout cfg) (ccHost cfg) (ccPort cfg)++  let tlsCfg = TlsConfig+        { tlsHostname = ccHost cfg+        , tlsVerify = ccTls cfg `elem` [TlsVerifyCa, TlsVerifyFull]+        , tlsVerifyHostname = ccTls cfg == TlsVerifyFull+        , tlsClientCert = ccSslCert cfg+        , tlsClientKey = ccSslKey cfg+        , tlsCaCert = ccSslRootCert cfg+        }++  -- Optionally upgrade to TLS+  wc <- case ccTls cfg of+    TlsDisable -> pure wc0+    TlsRequire -> upgradeTls wc0 tlsCfg+    TlsVerifyCa -> upgradeTls wc0 tlsCfg+    TlsVerifyFull -> upgradeTls wc0 tlsCfg+    TlsPrefer -> do+      result <- try @SomeException (upgradeTls wc0 tlsCfg)+      case result of+        Right tlsWc -> pure tlsWc+        Left _ -> do+          wcClose wc0+          connectTcpTimeout (ccConnectTimeout cfg) (ccHost cfg) (ccPort cfg)++  -- Send startup message+  let extraParams =+        [ ("client_encoding", ccClientEncoding cfg) | ccClientEncoding cfg /= "UTF8" ]+        ++ [ ("options", ccOptions cfg) | not (BS8.null (ccOptions cfg)) ]+      startup =+        StartupParams+          { spUser = ccUser cfg+          , spDatabase = ccDatabase cfg+          , spAppName = ccAppName cfg+          , spExtraParams = extraParams+          }+  sendRawBytes wc (buildStartup startup)++  -- Handle authentication+  handleAuth wc cfg++  -- Collect server parameters until ReadyForQuery+  paramsRef <- newIORef Map.empty+  pidRef <- newIORef 0+  keyRef <- newIORef 0+  txRef <- newIORef TxIdle++  let loop = do+        msg <- recvBackendMsg wc+        case msg of+          ParameterStatus name value -> do+            modifyIORef' paramsRef (Map.insert name value)+            loop+          BackendKeyData pid key -> do+            writeIORef pidRef pid+            writeIORef keyRef key+            loop+          ReadyForQuery status -> do+            writeIORef txRef status+            pure ()+          NoticeResponse _ -> loop+          other ->+            throwPgWire (ProtocolError ("Unexpected message during startup: " <> BS8.pack (show other)))++  loop++  pid <- readIORef pidRef+  key <- readIORef keyRef+  stmtCache <- newIORef PSQ.empty+  stmtCounter <- newIORef 0+  stmtTick <- newIORef 0+  let sslActive = case ccTls cfg of+        TlsDisable -> False+        _ -> True++  -- Spawn async writer + reader threads+  asyncConn <- spawnAsyncWireConn wc txRef paramsRef++  pure+    Connection+      { connAsync = asyncConn+      , connConfig = cfg+      , connBackendPid = pid+      , connBackendKey = key+      , connStmtCache = stmtCache+      , connStmtCounter = stmtCounter+      , connStmtTick = stmtTick+      , connSslActive = sslActive+      , connPreparedStatements = ccPreparedStatements cfg+      , connMaxPreparedStatements = max 1 (ccMaxPreparedStatements cfg)+      }++-- | Connect using a connection string.+connectString :: ByteString -> IO Connection+connectString bs = case parseConnString bs of+  Left err -> throwPgWire (ConnectionError (BS8.pack err))+  Right cfg -> connect cfg++-- | Close a connection.+close :: Connection -> IO ()+close conn = do+  shutdownAsyncWireConn (connAsync conn)+  -- Send Terminate and close the socket (best effort, threads are already dead)+  sendFrontendMsg (connWire conn) Terminate `catch_` pure ()+  wcClose (connWire conn) `catch_` pure ()+  where+    catch_ :: IO a -> IO a -> IO a+    catch_ action fallback = action `catch` \(_ :: SomeException) -> fallback++-- | Bracket-style connection management.+withConnection :: ConnConfig -> (Connection -> IO a) -> IO a+withConnection cfg = bracket (connect cfg) close++-- | Execute a simple text query (no parameters). Returns all result rows+-- as lists of nullable bytestrings, plus the command tag.+simpleQuery :: Connection -> ByteString -> IO ([[Maybe ByteString]], Maybe CommandTag)+simpleQuery conn sql = do+  resp <- submitRequest (connAsync conn) (ReqSimpleQuery sql)+  case resp of+    RespSimple rows tag -> pure (rows, tag)+    _ -> throwPgWire (ProtocolError "simpleQuery: unexpected response type")++-- | Reset the connection: close and reconnect using the same config.+reset :: Connection -> IO Connection+reset conn = do+  close conn+  connect (connConfig conn)++-- Connection status -------------------------------------------------------++-- | Check if the connection is alive by sending an empty query.+connectionStatus :: Connection -> IO Bool+connectionStatus conn = do+  result <- try @SomeException (simpleQuery conn "")+  pure $ case result of+    Right _ -> True+    Left _ -> False++-- | Get the current transaction status.+transactionStatus :: Connection -> IO TxStatus+transactionStatus conn = readIORef (connTxStatus conn)++-- | Look up a server parameter (e.g., @\"server_version\"@, @\"server_encoding\"@).+parameterStatus :: Connection -> ByteString -> IO (Maybe ByteString)+parameterStatus conn key = do+  params <- readIORef (connParams conn)+  pure (Map.lookup key params)++-- | Get the server version as an integer (e.g., 160004 for 16.4).+-- Parses from the @server_version@ parameter.+serverVersion :: Connection -> IO (Maybe Int)+serverVersion conn = do+  mVersion <- parameterStatus conn "server_version"+  pure (mVersion >>= parseServerVersion)+  where+    parseServerVersion bs =+      -- Strip distribution suffix: "16.6 (Ubuntu 16.6-0ubuntu0.24.04.1)" → "16.6"+      let versionOnly = BS8.takeWhile (\c -> c == '.' || (c >= '0' && c <= '9')) bs+          parts = BS8.split '.' versionOnly+       in case parts of+            [major, minor] -> do+              maj <- readInt major+              mn <- readInt minor+              pure (maj * 10000 + mn)+            [major, minor, patch] -> do+              maj <- readInt major+              mn <- readInt minor+              p <- readInt patch+              pure (maj * 10000 + mn * 100 + p)+            _ -> Nothing+    readInt s = case BS8.readInt s of+      Just (n, _) -> Just n+      Nothing -> Nothing++-- | Get the backend process ID.+backendPid :: Connection -> Int32+backendPid = connBackendPid++-- | Check if SSL/TLS is in use on this connection.+isSslInUse :: Connection -> Bool+isSslInUse = connSslActive++-- Notifications (non-blocking) --------------------------------------------++-- | Check for a pending notification without blocking.+-- Returns 'Nothing' if no notification is available.+checkNotification :: Connection -> IO (Maybe (Int32, ByteString, ByteString))+checkNotification conn = do+  -- Send empty query to flush pending notifications from the server+  _ <- simpleQuery conn ""+  -- Notifications are dispatched by the reader thread to the handler.+  -- For non-blocking check, we'd need a different mechanism.+  -- For now, return Nothing — notifications arrive via the handler.+  pure Nothing++-- Notice handling ---------------------------------------------------------++-- | Set a callback for server notice messages (warnings, info, etc.).+-- The default handler discards all notices.+setNoticeHandler :: Connection -> (PgNotice -> IO ()) -> IO ()+setNoticeHandler conn handler = writeIORef (connNoticeHandler conn) handler++-- SQL escaping ------------------------------------------------------------++-- | Escape a string for use as a SQL literal. Returns a properly quoted+-- and escaped string including the surrounding single quotes.+--+-- @+-- escapeLiteral conn "O'Brien"  ==  "'O''Brien'"+-- @+escapeLiteral :: Connection -> ByteString -> ByteString+escapeLiteral _conn bs =+  "'" <> BS8.concatMap escapeChar bs <> "'"+  where+    escapeChar '\'' = "''"+    escapeChar '\\' = "\\\\"+    escapeChar c = BS8.singleton c++-- | Escape a string for use as a SQL identifier (table, column, function name).+-- Returns a properly quoted identifier with surrounding double quotes.+--+-- @+-- escapeIdentifier conn "user table"  ==  "\"user table\""+-- @+escapeIdentifier :: Connection -> ByteString -> ByteString+escapeIdentifier _conn bs =+  "\"" <> BS8.concatMap escapeChar bs <> "\""+  where+    escapeChar '"' = "\"\""+    escapeChar c = BS8.singleton c++-- Statement introspection -------------------------------------------------++-- | Describe a prepared statement. Returns parameter OIDs and column metadata.+-- The statement must already be prepared on this connection.+describePrepared :: Connection -> ByteString -> IO (ParamDescription, ColumnDescription)+describePrepared conn stmtName = do+  -- describePrepared needs direct wire access (Describe + Sync protocol)+  submitExclusive (connAsync conn) $ \wc txRef -> do+    sendFrontendMsg wc (Describe DescribeStatement stmtName)+    sendFrontendMsg wc Sync+    (params, cols) <- collectDescribe wc txRef+    pure (ParamDescription params, ColumnDescription cols)+  where+    collectDescribe wc txRef = do+      pds <- collectParams wc+      cds <- collectColumns wc+      waitDescribeReady wc txRef+      pure (pds, cds)++    collectParams wc = do+      msg <- recvBackendMsg wc+      case msg of+        ParameterDescription oids -> pure oids+        ErrorResponse err -> throwPgWire (QueryError err)+        _ -> collectParams wc++    collectColumns wc = do+      msg <- recvBackendMsg wc+      case msg of+        RowDescription fields -> pure fields+        NoData -> pure V.empty+        ErrorResponse err -> throwPgWire (QueryError err)+        _ -> collectColumns wc++    waitDescribeReady wc txRef = do+      msg <- recvBackendMsg wc+      case msg of+        ReadyForQuery status -> writeIORef txRef status+        _ -> waitDescribeReady wc txRef++-- Server ping -------------------------------------------------------------++-- | Check if a PostgreSQL server is accepting connections, without+-- fully authenticating. Attempts a TCP connection and checks if the+-- server responds to the startup sequence.+ping :: ConnConfig -> IO Bool+ping cfg = do+  result <- try @SomeException (connect cfg >>= close)+  pure $ case result of+    Right () -> True+    Left _ -> False++-- Password file -----------------------------------------------------------++-- | Look up a password from @~/.pgpass@ file.+-- Format: @hostname:port:database:username:password@ (one per line).+-- @*@ matches any value in a field.+lookupPgpass :: ConnConfig -> IO (Maybe ByteString)+lookupPgpass cfg = do+  home <- lookupEnv "HOME"+  case home of+    Nothing -> pure Nothing+    Just h -> do+      let path = h <> "/.pgpass"+      exists <- doesFileExist path+      if not exists+        then pure Nothing+        else do+          contents <- BS8.readFile path+          let host = BS8.pack (ccHost cfg)+              port = BS8.pack (show (ccPort cfg))+              db = ccDatabase cfg+              user = ccUser cfg+          pure (findMatch host port db user (BS8.lines contents))+  where+    findMatch _ _ _ _ [] = Nothing+    findMatch host port db user (line : rest)+      | BS8.null line || BS8.isPrefixOf "#" line = findMatch host port db user rest+      | otherwise = case BS8.split ':' line of+          [h, p, d, u, pw]+            | matches h host && matches p port && matches d db && matches u user ->+                Just pw+          _ -> findMatch host port db user rest+    matches "*" _ = True+    matches pattern value = pattern == value++-- Password encryption -----------------------------------------------------++-- | Encrypt a password for storage, using the same algorithm as+-- @PQencryptPasswordConn@. Supports MD5 format.+encryptPassword :: ByteString -> ByteString -> ByteString+encryptPassword user password =+  -- MD5 format: "md5" + hex(md5(password + user))+  let digest = BA.convert (hash (password <> user) :: Digest MD5) :: ByteString+   in "md5" <> toHex digest++toHex :: ByteString -> ByteString+toHex bs = BS8.pack (concatMap byteToHex (BS.unpack bs))+  where+    byteToHex w =+      let (hi, lo) = w `divMod` 16+       in [hexDigit hi, hexDigit lo]+    hexDigit n+      | n < 10 = toEnum (fromEnum '0' + fromIntegral n)+      | otherwise = toEnum (fromEnum 'a' + fromIntegral n - 10)++-- Protocol tracing --------------------------------------------------------++-- | Enable protocol tracing on this connection.+-- The callback receives 'True' for send, 'False' for recv, plus the raw bytes.+-- Set to @Nothing@ to disable tracing.+--+-- @+-- setTraceHandler conn $ \\isSend bytes ->+--   BS8.putStrLn $ (if isSend then \">>> \" else \"<<< \") <> BS8.take 40 bytes+-- @+setTraceHandler :: Connection -> (Bool -> ByteString -> IO ()) -> IO ()+setTraceHandler conn handler =+  writeIORef (wcTrace (connWire conn)) (Just wrapper)+  where+    wrapper TraceSend bs = handler True bs+    wrapper TraceRecv bs = handler False bs++-- Authentication ----------------------------------------------------------++handleAuth :: WireConn -> ConnConfig -> IO ()+handleAuth wc cfg = do+  msg <- recvBackendMsg wc+  case msg of+    Authentication AuthOk -> pure ()+    Authentication AuthCleartextPassword -> do+      cleartextAuth wc (ccPassword cfg)+      waitForAuthOk wc+    Authentication (AuthMD5Password salt) -> do+      md5Auth wc (ccUser cfg) (ccPassword cfg) salt+      waitForAuthOk wc+    Authentication (AuthSASL mechs)+      | "SCRAM-SHA-256" `elem` mechs -> do+          scramAuth wc (ccUser cfg) (ccPassword cfg)+          -- scramAuth handles waiting for AuthOk internally+      | otherwise ->+          throwPgWire (AuthError ("Unsupported SASL mechanisms: " <> BS8.pack (show mechs)))+    ErrorResponse err -> throwPgWire (AuthError (pgMessage err))+    other -> throwPgWire (AuthError ("Unexpected auth message: " <> BS8.pack (show other)))++waitForAuthOk :: WireConn -> IO ()+waitForAuthOk wc = do+  msg <- recvBackendMsg wc+  case msg of+    Authentication AuthOk -> pure ()+    ErrorResponse err -> throwPgWire (AuthError (pgMessage err))+    other -> throwPgWire (AuthError ("Expected AuthOk, got: " <> BS8.pack (show other)))
+ src/PgWire/Connection/Config.hs view
@@ -0,0 +1,298 @@+-- | Configuration types for PostgreSQL connections.+--+-- Supports both URI format (@postgres:\/\/user:pass\@host:port\/db@) and+-- key=value format (@host=localhost port=5432 dbname=mydb@). Multi-host+-- failover, TLS modes, and session attribute targeting are configurable.+module PgWire.Connection.Config+  ( ConnConfig (..)+  , TlsMode (..)+  , TargetSessionAttrs (..)+  , defaultConnConfig+  , parseConnString+  , parseHosts+  ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.Char (isDigit)+import Data.Maybe (fromMaybe)+import Data.Time (NominalDiffTime)+import Data.Word (Word16)+import Network.Socket (HostName, PortNumber)++-- | TLS connection mode, matching PostgreSQL's @sslmode@ parameter.+data TlsMode+  = TlsDisable       -- ^ No SSL+  | TlsPrefer        -- ^ Try SSL, fall back to plain+  | TlsRequire       -- ^ Require SSL, don't verify server cert+  | TlsVerifyCa      -- ^ Require SSL, verify server cert against CA+  | TlsVerifyFull    -- ^ Require SSL, verify cert + hostname match+  deriving stock (Show, Eq)++-- | Session attribute requirements for multi-host failover.+data TargetSessionAttrs+  = SessionAny            -- ^ Connect to any server+  | SessionReadWrite      -- ^ Must be a read-write server (primary)+  | SessionReadOnly       -- ^ Must be a read-only server (standby)+  | SessionPrimary        -- ^ Must be the primary+  | SessionStandby        -- ^ Must be a standby+  | SessionPreferStandby  -- ^ Prefer standby, fall back to primary+  deriving stock (Show, Eq)++-- | Connection configuration for a single PostgreSQL server.+data ConnConfig = ConnConfig+  { ccHost :: HostName+  -- ^ Primary host. For multi-host, use comma-separated: @\"host1,host2\"@.+  , ccPort :: PortNumber+  , ccDatabase :: ByteString+  , ccUser :: ByteString+  , ccPassword :: ByteString+  , ccTls :: TlsMode+  , ccAppName :: ByteString+  , ccConnectTimeout :: NominalDiffTime+  -- ^ Timeout for establishing the TCP connection (seconds). 0 = no timeout.+  , ccQueryTimeout :: NominalDiffTime+  -- ^ Default timeout for query execution (seconds). 0 = no timeout.+  -- Can be overridden per-query with 'withQueryTimeout'.+  , ccSslCert :: Maybe FilePath+  -- ^ Path to client SSL certificate file (@sslcert@).+  , ccSslKey :: Maybe FilePath+  -- ^ Path to client SSL private key file (@sslkey@).+  , ccSslRootCert :: Maybe FilePath+  -- ^ Path to SSL certificate authority (CA) file (@sslrootcert@).+  -- If @\"system\"@, uses the system certificate store.+  , ccKeepalives :: Bool+  -- ^ Enable TCP keepalives (default: True).+  , ccKeepalivesIdle :: Int+  -- ^ Seconds before first keepalive probe (default: 0 = system default).+  , ccKeepalivesInterval :: Int+  -- ^ Seconds between keepalive probes (default: 0 = system default).+  , ccKeepalivesCount :: Int+  -- ^ Number of failed probes before disconnect (default: 0 = system default).+  , ccTargetSessionAttrs :: TargetSessionAttrs+  -- ^ Required session attributes for multi-host failover (default: 'SessionAny').+  , ccClientEncoding :: ByteString+  -- ^ Client encoding (default: @\"UTF8\"@). Set to @\"auto\"@ for auto-detection.+  , ccOptions :: ByteString+  -- ^ Extra command-line options to send to the server at startup.+  , ccLoadBalanceHosts :: Bool+  -- ^ Randomize host order for multi-host connections (default: False).+  , ccPreparedStatements :: Bool+  -- ^ Use named prepared statements (default: True). Set to False for+  -- compatibility with PgBouncer in transaction mode or pgpool.+  -- When False, all queries use the unnamed statement which is+  -- parsed and discarded each time.+  , ccMaxPreparedStatements :: !Int+  -- ^ Maximum number of prepared statements cached per connection+  -- (default: 256). When the cache is full, the least recently used+  -- statement is evicted. Higher values reduce re-parses for workloads+  -- with many distinct queries, at the cost of server-side memory+  -- (each cached statement holds a parse tree in PostgreSQL).+  -- Ignored when 'ccPreparedStatements' is False.+  }+  deriving stock (Show)++-- | Default connection configuration: @localhost:5432@, no TLS, 10s connect timeout.+defaultConnConfig :: ConnConfig+defaultConnConfig =+  ConnConfig+    { ccHost = "localhost"+    , ccPort = 5432+    , ccDatabase = ""+    , ccUser = ""+    , ccPassword = ""+    , ccTls = TlsDisable+    , ccAppName = "pg-wire"+    , ccConnectTimeout = 10+    , ccQueryTimeout = 0+    , ccSslCert = Nothing+    , ccSslKey = Nothing+    , ccSslRootCert = Nothing+    , ccKeepalives = True+    , ccKeepalivesIdle = 0+    , ccKeepalivesInterval = 0+    , ccKeepalivesCount = 0+    , ccTargetSessionAttrs = SessionAny+    , ccClientEncoding = "UTF8"+    , ccOptions = ""+    , ccLoadBalanceHosts = False+    , ccPreparedStatements = True+    , ccMaxPreparedStatements = 256+    }++-- | Parse a PostgreSQL connection string.+-- Supports both URI format (@postgres://user:pass\@host:port/db@)+-- and key=value format (@host=localhost port=5432 dbname=mydb@).+parseConnString :: ByteString -> Either String ConnConfig+parseConnString bs+  | "postgres://" `BS8.isPrefixOf` bs || "postgresql://" `BS8.isPrefixOf` bs =+      parseUri bs+  | otherwise =+      parseKeyValue bs++parseUri :: ByteString -> Either String ConnConfig+parseUri bs = do+  -- Strip scheme: "postgres://..." -> drop up to first '/', then drop "//"+  let afterScheme = BS8.drop 2 (BS8.dropWhile (/= '/') bs) -- drop "scheme://"+      afterSlash = afterScheme++  -- Split user:pass@host:port/db?params+  let (authHost, pathQuery) = case BS8.break (== '/') afterSlash of+        (ah, pq) -> (ah, BS8.drop 1 pq)+      dbAndParams = pathQuery+      (db, queryStr) = BS8.break (== '?') dbAndParams++  -- Split auth@host+  let (auth, hostPort) = case BS8.breakEnd (== '@') authHost of+        (a, _) | BS8.null a -> ("", authHost)+        (a, hp) -> (BS8.init a, hp) -- drop trailing '@'++  -- Split user:pass+  let (user, pass) = case BS8.break (== ':') auth of+        (u, p)+          | BS8.null p -> (u, "")+          | otherwise -> (u, BS8.drop 1 p)++  -- Split host:port+  let (host, portStr) = case BS8.break (== ':') hostPort of+        (h, p)+          | BS8.null p -> (h, "5432")+          | otherwise -> (h, BS8.drop 1 p)+      port = maybe 5432 fromIntegral (readPort portStr)++  -- Parse query params for sslmode+  let params = parseQueryParams (BS8.drop 1 queryStr) -- drop '?'+      tlsMode = case lookup "sslmode" params of+        Just "require"     -> TlsRequire+        Just "prefer"      -> TlsPrefer+        Just "disable"     -> TlsDisable+        Just "verify-ca"   -> TlsVerifyCa+        Just "verify-full" -> TlsVerifyFull+        _ -> TlsDisable+      sslCert = BS8.unpack <$> lookup "sslcert" params+      sslKey = BS8.unpack <$> lookup "sslkey" params+      sslRootCert = BS8.unpack <$> lookup "sslrootcert" params++  Right+    ConnConfig+      { ccHost = BS8.unpack host+      , ccPort = port+      , ccDatabase = db+      , ccUser = user+      , ccPassword = pass+      , ccTls = tlsMode+      , ccAppName = "pg-wire"+      , ccConnectTimeout = readTimeout (lookup "connect_timeout" params)+      , ccQueryTimeout = 0+      , ccSslCert = sslCert+      , ccSslKey = sslKey+      , ccSslRootCert = sslRootCert+      , ccKeepalives = True+      , ccKeepalivesIdle = 0+      , ccKeepalivesInterval = 0+      , ccKeepalivesCount = 0+      , ccTargetSessionAttrs = maybe SessionAny parseSessionAttrs+          (lookup "target_session_attrs" params)+      , ccClientEncoding = fromMaybe "UTF8" (lookup "client_encoding" params)+      , ccOptions = fromMaybe "" (lookup "options" params)+      , ccLoadBalanceHosts = lookup "load_balance_hosts" params == Just "random"+      , ccPreparedStatements = lookup "prepared_statements" params /= Just "false"+      , ccMaxPreparedStatements = maybe 256 (\v -> readIntDef 256 v) (lookup "statement_cache_size" params)+      }++readTimeout :: Maybe ByteString -> NominalDiffTime+readTimeout Nothing = 10+readTimeout (Just bs) = case BS8.readInt bs of+  Just (n, _) | n > 0 -> fromIntegral n+  _ -> 10++parseKeyValue :: ByteString -> Either String ConnConfig+parseKeyValue bs =+  let pairs = map parsePair (BS8.words bs)+      get key def = fromMaybe def (lookup key pairs)+      portStr = get "port" "5432"+      port = maybe 5432 fromIntegral (readPort portStr)+      tlsMode = case get "sslmode" "disable" of+        "require"     -> TlsRequire+        "prefer"      -> TlsPrefer+        "verify-ca"   -> TlsVerifyCa+        "verify-full" -> TlsVerifyFull+        _ -> TlsDisable+      getMaybe key = let v = get key "" in if BS8.null v then Nothing else Just (BS8.unpack v)+   in Right+        ConnConfig+          { ccHost = BS8.unpack (get "host" "localhost")+          , ccPort = port+          , ccDatabase = get "dbname" (get "database" "")+          , ccUser = get "user" ""+          , ccPassword = get "password" ""+          , ccTls = tlsMode+          , ccAppName = get "application_name" "pg-wire"+          , ccConnectTimeout = case BS8.readInt (get "connect_timeout" "10") of+              Just (n, _) | n > 0 -> fromIntegral n+              _ -> 10+          , ccQueryTimeout = case BS8.readInt (get "query_timeout" "0") of+              Just (n, _) | n > 0 -> fromIntegral n+              _ -> 0+          , ccSslCert = getMaybe "sslcert"+          , ccSslKey = getMaybe "sslkey"+          , ccSslRootCert = getMaybe "sslrootcert"+          , ccKeepalives = get "keepalives" "1" /= "0"+          , ccKeepalivesIdle = readIntDef 0 (get "keepalives_idle" "0")+          , ccKeepalivesInterval = readIntDef 0 (get "keepalives_interval" "0")+          , ccKeepalivesCount = readIntDef 0 (get "keepalives_count" "0")+          , ccTargetSessionAttrs = parseSessionAttrs (get "target_session_attrs" "any")+          , ccClientEncoding = get "client_encoding" "UTF8"+          , ccOptions = get "options" ""+          , ccLoadBalanceHosts = get "load_balance_hosts" "disable" == "random"+          , ccPreparedStatements = get "prepared_statements" "true" /= "false"+          , ccMaxPreparedStatements = readIntDef 256 (get "statement_cache_size" "256")+          }+  where+    parsePair p =+      let (k, v) = BS8.break (== '=') p+       in (k, BS8.drop 1 v)++readIntDef :: Int -> ByteString -> Int+readIntDef d bs = case BS8.readInt bs of+  Just (n, _) -> n+  Nothing -> d++readPort :: ByteString -> Maybe Word16+readPort bs+  | BS8.null bs = Nothing+  | BS8.all (\c -> isDigit c) bs =+      let n = BS8.foldl' (\acc c -> acc * 10 + fromIntegral (fromEnum c - 48)) 0 bs :: Int+       in if n > 0 && n <= 65535 then Just (fromIntegral n) else Nothing+  | otherwise = Nothing++parseQueryParams :: ByteString -> [(ByteString, ByteString)]+parseQueryParams bs+  | BS8.null bs = []+  | otherwise =+      [ (k, BS8.drop 1 v)+      | param <- BS8.split '&' bs+      , let (k, v) = BS8.break (== '=') param+      , not (BS8.null k)+      ]++parseSessionAttrs :: ByteString -> TargetSessionAttrs+parseSessionAttrs "any"              = SessionAny+parseSessionAttrs "read-write"       = SessionReadWrite+parseSessionAttrs "read-only"        = SessionReadOnly+parseSessionAttrs "primary"          = SessionPrimary+parseSessionAttrs "standby"          = SessionStandby+parseSessionAttrs "prefer-standby"   = SessionPreferStandby+parseSessionAttrs _                  = SessionAny++-- | Parse a comma-separated host string into individual hosts.+-- @\"host1,host2,host3\"@ → @[\"host1\", \"host2\", \"host3\"]@+parseHosts :: String -> [String]+parseHosts [] = ["localhost"]+parseHosts s = go s+  where+    go [] = []+    go str = let (h, rest) = break (== ',') str+              in h : case rest of+                       [] -> []+                       (_ : rs) -> go rs
+ src/PgWire/Error.hs view
@@ -0,0 +1,44 @@+-- | Error types for pg-wire runtime exceptions.+--+-- All errors are thrown as 'PgWireError' via 'throwIO' and can be caught+-- with the standard @Control.Exception@ machinery.+module PgWire.Error+  ( PgWireError (..)+  , throwPgWire+  ) where++import Control.Exception (Exception, throwIO)+import Data.ByteString (ByteString)+import GHC.Generics (Generic)+import NoThunks.Class (NoThunks)+import PgWire.Protocol.Backend (PgError)++-- | All runtime errors thrown by pg-wire.+data PgWireError+  = ConnectionError ByteString+  -- ^ Failed to establish a TCP\/TLS connection to the server.+  | AuthError ByteString+  -- ^ Authentication failed (wrong password, unsupported mechanism, etc.).+  | ProtocolError ByteString+  -- ^ Unexpected message in the wire protocol (indicates a driver bug or+  -- incompatible server).+  | QueryError PgError+  -- ^ The server returned an error in response to a query. The 'PgError'+  -- contains SQLSTATE, message, detail, hint, and position fields.+  | DecodeError ByteString+  -- ^ Failed to decode a result row (type mismatch, unexpected NULL, or+  -- no rows returned to 'fetchOneOrThrow').+  | PoolTimeout+  -- ^ Timed out waiting to acquire a connection from the pool.+  | PoolClosed+  -- ^ Attempted to use a pool after 'PgWire.Pool.closePool' was called.+  | ConnectionDead+  -- ^ The connection's reader thread has died (network failure, server crash).+  deriving stock (Show, Eq, Generic)++instance NoThunks PgWireError+instance Exception PgWireError++-- | Throw a 'PgWireError' as an exception.+throwPgWire :: PgWireError -> IO a+throwPgWire = throwIO
+ src/PgWire/Pool.hs view
@@ -0,0 +1,796 @@+{-# LANGUAGE DataKinds #-}++-- | Thread-safe connection pool for PostgreSQL.+--+-- Manages a set of reusable 'Connection's with configurable pool size,+-- idle timeout, max lifetime, health checking, and lifecycle hooks.+-- Inspired by deadpool-postgres (Rust).+--+-- @+-- pool <- 'newPool' 'PgWire.Pool.Config.defaultPoolConfig'+--   { poolConnString = \"postgres:\/\/...\"+--   , poolSize = 10+--   }+-- 'withResource' pool $ \\conn -> ...+-- 'closePool' pool+-- @+module PgWire.Pool+  ( Pool (..)+  , PoolStats (..)+  , newPool+  , closePool+  , drainPool+  , withResource+  , poolStats+  , poolIsAlive+  , resize+  , retain+  , setPostCreateHook+  , setOnAcquireHook+  , setPreReleaseHook+  , withResourceTimeout+  ) where++import GHC.Generics (Generic)+import NoThunks.Class (NoThunks (..), OnlyCheckWhnfNamed (..), allNoThunks)+import Control.Concurrent (threadDelay)+import Control.Concurrent.Async (Async, async, cancel, race)+import Control.Concurrent.STM+import Control.Exception (AsyncException, SomeException, catch, mask, onException, throwIO, try)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Int (Int32)+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Sequence (Seq (..))+import Data.Sequence qualified as Seq+import Data.Time (NominalDiffTime, UTCTime, addUTCTime, diffUTCTime, getCurrentTime)+import Data.HashPSQ qualified as PSQ+import PgWire.Async (AsyncWireConn (..))+import PgWire.Connection (Connection (..), close, connectString, simpleQuery)+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Pool.Config (PoolConfig (..), QueueMode (..), RecyclingMethod (..))+import PgWire.Pool.Observation (PoolEvent (ConnectionAcquired, ConnectionCreated, ConnectionDestroyed, ConnectionRecycled, ConnectionReleased, HealthCheckFailed, AcquireTimeout, ReaperSwept, WarmerCreated, PoolResized, PoolShutdown))+import PgWire.TypeCache (TypeCache, newTypeCache)+import System.Random (randomRIO)++-- | A thread-safe connection pool for PostgreSQL.+--+-- Manages a bounded set of reusable 'Connection's with idle timeout, max lifetime+-- (with jitter to avoid thundering-herd reconnects), health checking, and lifecycle hooks.+-- The pool is safe to share across threads. Use 'newPool' to create and 'closePool' to shut down.+data Pool = Pool+  { pConfig :: PoolConfig+  , pIdle :: TVar (Seq PoolEntry)+  , pActive :: TVar Int+  , pWaiters :: TVar (Seq (TMVar (Either PgWireError Connection)))+  , pClosed :: TVar Bool+  , pReaper :: IORef (Async ())+  , pWarmer :: IORef (Maybe (Async ()))+  , pEffectiveSize :: TVar Int+  -- ^ Mutable pool size for runtime 'resize'.+  , pTotalCreated :: TVar Int+  , pTotalDestroyed :: TVar Int+  , pTotalTimeouts :: TVar Int+  , pConnMeta :: TVar (Map Int32 ConnMeta)+  -- ^ Per-connection metadata keyed by backend PID.+  -- Tracks creation time and jittered max-life for each connection.+  , pOnCreate :: IORef (Connection -> IO ())+  , pOnAcquire :: IORef (Connection -> IO ())+  , pOnRelease :: IORef (Connection -> IO ())+  , pTypeCache :: TypeCache+  -- ^ Pool-level cache for resolved PG type metadata (OID → TypeInfo).+  -- Avoids redundant @pg_type@ round-trips across connections.+  }++-- | Per-connection metadata, stored by backend PID.+data ConnMeta = ConnMeta+  { cmCreatedAt :: !UTCTime+  , cmDeadline :: !UTCTime+  -- ^ Creation time + jittered max-life. Connection is expired after this.+  }+  deriving stock (Generic)++-- | 'UTCTime' is treated as opaque (its internal 'Day'/'DiffTime' are+-- lazy in the @time@ package). Bounded size, not a leak target.+instance NoThunks ConnMeta where+  showTypeOf _ = "ConnMeta"+  wNoThunks ctx (ConnMeta createdAt deadline) = allNoThunks+    [ noThunks ctx (OnlyCheckWhnfNamed @"UTCTime" createdAt)+    , noThunks ctx (OnlyCheckWhnfNamed @"UTCTime" deadline)+    ]++data PoolEntry = PoolEntry+  { peConn :: Connection+  , peCreatedAt :: UTCTime+  , peLastUsed :: IORef UTCTime+  }++-- | Treats 'Connection', the @IORef UTCTime@, and 'UTCTime' itself as+-- opaque references (their internal state is not pool-bookkeeping and+-- 'UTCTime' has lazy fields in the @time@ package).+instance NoThunks PoolEntry where+  showTypeOf _ = "PoolEntry"+  wNoThunks ctx (PoolEntry conn createdAt lastUsed) = allNoThunks+    [ noThunks ctx (OnlyCheckWhnfNamed @"Connection" conn)+    , noThunks ctx (OnlyCheckWhnfNamed @"UTCTime" createdAt)+    , noThunks ctx (OnlyCheckWhnfNamed @"IORef UTCTime" lastUsed)+    ]++-- | Snapshot of pool statistics at a point in time, obtained via 'poolStats'.+data PoolStats = PoolStats+  { psIdle :: !Int+  -- ^ Number of connections currently idle in the pool.+  , psInUse :: !Int+  -- ^ Number of connections currently checked out by callers.+  , psWaiters :: !Int+  -- ^ Number of threads blocked waiting to acquire a connection.+  , psMaxSize :: !Int+  -- ^ Current maximum pool size (may differ from initial config after 'resize').+  , psTotalCreated :: !Int+  -- ^ Cumulative count of connections created since pool creation.+  , psTotalDestroyed :: !Int+  -- ^ Cumulative count of connections destroyed since pool creation.+  , psTotalTimeouts :: !Int+  -- ^ Cumulative count of acquire attempts that timed out.+  }+  deriving stock (Show, Eq, Generic)++instance NoThunks PoolStats++-- | Create a new connection pool from the given configuration.+--+-- Starts background threads for reaping idle\/expired connections and (if+-- @poolMinIdle > 0@) warming the pool to maintain a minimum number of idle+-- connections. No connections are created eagerly; the first 'withResource'+-- call triggers the initial connection.+--+-- Call 'closePool' when the pool is no longer needed to release all resources.+newPool :: PoolConfig -> IO Pool+newPool cfg = do+  idle <- newTVarIO Seq.empty+  active <- newTVarIO 0+  waiters <- newTVarIO Seq.empty+  closed <- newTVarIO False+  reaperRef <- newIORef (error "reaper not started")+  warmerRef <- newIORef Nothing+  effectiveSize <- newTVarIO (poolSize cfg)+  totalCreated <- newTVarIO 0+  totalDestroyed <- newTVarIO 0+  totalTimeouts <- newTVarIO 0+  connMeta <- newTVarIO Map.empty+  onCreate <- newIORef (\_ -> pure ())+  onAcquire <- newIORef (\_ -> pure ())+  onRelease <- newIORef (\_ -> pure ())+  typeCache <- newTypeCache+  let pool =+        Pool+          { pConfig = cfg+          , pIdle = idle+          , pActive = active+          , pWaiters = waiters+          , pClosed = closed+          , pReaper = reaperRef+          , pWarmer = warmerRef+          , pEffectiveSize = effectiveSize+          , pTotalCreated = totalCreated+          , pTotalDestroyed = totalDestroyed+          , pTotalTimeouts = totalTimeouts+          , pConnMeta = connMeta+          , pOnCreate = onCreate+          , pOnAcquire = onAcquire+          , pOnRelease = onRelease+          , pTypeCache = typeCache+          }+  reaper <- async (reaperThread pool)+  writeIORef reaperRef reaper+  -- Spawn warmer if poolMinIdle > 0+  when (poolMinIdle cfg > 0) $ do+    w <- async (warmerThread pool)+    writeIORef warmerRef (Just w)+  pure pool++-- | Close the pool, releasing all resources.+--+-- Closes every idle connection, cancels background threads (reaper and warmer),+-- and wakes all blocked waiters with a 'PoolClosed' error. After this call,+-- any subsequent 'withResource' will throw 'PoolClosed'. Idempotent: calling+-- 'closePool' on an already-closed pool is a no-op for the connection teardown.+closePool :: Pool -> IO ()+closePool pool = do+  -- Atomically: mark closed, drain idle, drain waiters+  (entries, blockedWaiters) <- atomically $ do+    writeTVar (pClosed pool) True+    idle <- readTVar (pIdle pool)+    writeTVar (pIdle pool) Seq.empty+    ws <- readTVar (pWaiters pool)+    writeTVar (pWaiters pool) Seq.empty+    pure (idle, ws)+  -- Cancel background threads+  reaper <- readIORef (pReaper pool)+  cancel reaper+  mWarmer <- readIORef (pWarmer pool)+  mapM_ cancel mWarmer+  -- Wake all blocked waiters with PoolClosed+  mapM_ (\w -> atomically $ tryPutTMVar w (Left PoolClosed)) blockedWaiters+  -- Close all idle connections (properly tracked)+  mapM_ (destroyEntry pool "pool closed") entries+  observe pool PoolShutdown++-- | Gracefully drain the pool for rolling deployments.+--+-- Marks the pool as closed (no new acquires), then polls until all+-- in-use connections are returned, then closes everything. Blocks+-- until the drain is complete or the timeout expires.+--+-- @+-- -- During rolling deployment:+-- drainPool pool 30  -- wait up to 30 seconds for in-flight queries+-- @+drainPool :: Pool -> NominalDiffTime -> IO ()+drainPool pool timeout = do+  -- Mark as closed — new withResource calls will get PoolClosed+  atomically $ writeTVar (pClosed pool) True+  -- Wake blocked waiters immediately+  ws <- atomically $ do+    waiters <- readTVar (pWaiters pool)+    writeTVar (pWaiters pool) Seq.empty+    pure waiters+  mapM_ (\w -> atomically $ tryPutTMVar w (Left PoolClosed)) ws+  -- Poll until all active connections are returned or timeout+  let timeoutMicros = round (timeout * 1000000) :: Int+      pollInterval = 100000 -- 100ms+      pollLoop !remaining+        | remaining <= 0 = pure () -- timeout, force close+        | otherwise = do+            active <- readTVarIO (pActive pool)+            idle <- Seq.length <$> readTVarIO (pIdle pool)+            if active <= idle+              then pure () -- all connections returned to idle+              else do+                threadDelay pollInterval+                pollLoop (remaining - pollInterval)+  pollLoop timeoutMicros+  -- Now close everything (reuse closePool's cleanup logic)+  entries <- atomically $ do+    idle <- readTVar (pIdle pool)+    writeTVar (pIdle pool) Seq.empty+    pure idle+  reaper <- readIORef (pReaper pool)+  cancel reaper+  mWarmer <- readIORef (pWarmer pool)+  mapM_ cancel mWarmer+  mapM_ (destroyEntry pool "drain") entries+  observe pool PoolShutdown++-- | Acquire a connection from the pool, run an action, and return the connection.+--+-- The connection is guaranteed to be returned to the pool (or destroyed) even+-- if the action throws an exception. Uses 'mask' internally so the acquire and+-- release cannot be interrupted by async exceptions.+--+-- @+-- withResource pool $ \\conn ->+--   simpleQuery conn \"SELECT 1\"+-- @+withResource :: Pool -> (Connection -> IO a) -> IO a+withResource pool action = mask $ \restore -> do+  t0 <- getCurrentTime+  conn <- acquire pool+  t1 <- getCurrentTime+  observe pool (ConnectionAcquired (diffUTCTime t1 t0))+  result <- restore (action conn) `onException` destroyConn pool conn "exception"+  release pool conn+  observe pool ConnectionReleased+  pure result++-- | Like 'withResource' but with a custom acquire timeout.+--+-- Overrides 'poolAcquireTimeout' for this single acquisition. Useful when+-- certain operations can tolerate longer (or shorter) waits than the pool+-- default.+--+-- @+-- -- Wait up to 30 seconds for this particular query+-- withResourceTimeout pool 30 $ \\conn ->+--   simpleQuery conn \"SELECT expensive_function()\"+-- @+withResourceTimeout :: Pool -> NominalDiffTime -> (Connection -> IO a) -> IO a+withResourceTimeout pool timeout action = mask $ \restore -> do+  t0 <- getCurrentTime+  conn <- acquireWithTimeout pool timeout+  t1 <- getCurrentTime+  observe pool (ConnectionAcquired (diffUTCTime t1 t0))+  result <- restore (action conn) `onException` destroyConn pool conn "exception"+  release pool conn+  observe pool ConnectionReleased+  pure result++-- | Get an atomic snapshot of the pool's current statistics.+--+-- The returned 'PoolStats' is consistent (all fields read in a single STM+-- transaction) but represents a point-in-time view that may be stale by+-- the time you inspect it.+poolStats :: Pool -> IO PoolStats+poolStats pool = atomically $ do+  idle <- Seq.length <$> readTVar (pIdle pool)+  active <- readTVar (pActive pool)+  waiters <- Seq.length <$> readTVar (pWaiters pool)+  maxSz <- readTVar (pEffectiveSize pool)+  created <- readTVar (pTotalCreated pool)+  destroyed <- readTVar (pTotalDestroyed pool)+  timeouts <- readTVar (pTotalTimeouts pool)+  pure PoolStats+    { psIdle = idle+    , psInUse = active - idle+    , psWaiters = waiters+    , psMaxSize = maxSz+    , psTotalCreated = created+    , psTotalDestroyed = destroyed+    , psTotalTimeouts = timeouts+    }++-- | Check if the pool is still open (not closed).+--+-- Returns 'True' if the pool is alive and accepting connections,+-- 'False' if 'closePool' has been called.+poolIsAlive :: Pool -> IO Bool+poolIsAlive pool = not <$> readTVarIO (pClosed pool)++-- | Resize the pool at runtime to a new maximum connection count.+--+-- If shrinking, excess idle connections are closed immediately. In-use+-- connections are not interrupted and will finish naturally; they are simply+-- not returned to the pool once the count exceeds the new limit. If growing,+-- new connections are created on demand by subsequent 'withResource' calls.+resize :: Pool -> Int -> IO ()+resize pool newSize = do+  (oldSize, excess) <- atomically $ do+    old <- readTVar (pEffectiveSize pool)+    writeTVar (pEffectiveSize pool) newSize+    idle <- readTVar (pIdle pool)+    let idleCount = Seq.length idle+        excessCount = max 0 (idleCount - newSize)+    if excessCount > 0+      then do+        let (keep, toDrop) = Seq.splitAt (idleCount - excessCount) idle+        writeTVar (pIdle pool) keep+        pure (old, seqToList toDrop)+      else pure (old, [])+  logPool pool "info" ("resized from " <> BS8.pack (show oldSize) <> " to " <> BS8.pack (show newSize))+  observe pool (PoolResized oldSize newSize)+  mapM_ (destroyEntry pool "resize") excess++-- | Filter idle connections, keeping only those that satisfy the predicate.+--+-- Atomically drains the idle queue, applies the predicate to each connection+-- in IO, returns matching ones to the pool, and destroys the rest. Useful for+-- invalidating connections after a credential rotation or schema migration.+--+-- @+-- -- Drop all connections that were created before the migration+-- retain pool (\\conn -> isPostMigration conn)+-- @+retain :: Pool -> (Connection -> IO Bool) -> IO ()+retain pool predicate = do+  entries <- atomically $ do+    idle <- readTVar (pIdle pool)+    writeTVar (pIdle pool) Seq.empty+    pure (seqToList idle)+  (kept, dropped) <- partitionM (\e -> predicate (peConn e)) entries+  atomically $ modifyTVar' (pIdle pool) (Seq.fromList kept Seq.><)+  mapM_ (destroyEntry pool "retain") dropped++-- | Set a hook called immediately after a new connection is created.+--+-- Useful for running session-level setup such as @SET search_path@ or+-- @SET statement_timeout@. If the hook throws, the exception is silently+-- caught and the connection is still placed in the pool.+setPostCreateHook :: Pool -> (Connection -> IO ()) -> IO ()+setPostCreateHook pool hook = writeIORef (pOnCreate pool) $! hook++-- | Set a hook called each time a connection is handed out to a caller via 'withResource'.+--+-- Runs after any recycling\/health checks have passed. If the hook throws, the+-- exception is silently caught and the connection is still provided.+setOnAcquireHook :: Pool -> (Connection -> IO ()) -> IO ()+setOnAcquireHook pool hook = writeIORef (pOnAcquire pool) $! hook++-- | Set a hook called before a connection is returned to the idle pool.+--+-- Runs after the caller's action completes and before the connection is placed+-- back in the idle queue or delivered to a waiting thread. Useful for cleanup+-- such as @RESET ALL@ or @DISCARD TEMP@. If the hook throws, the exception is+-- silently caught.+setPreReleaseHook :: Pool -> (Connection -> IO ()) -> IO ()+setPreReleaseHook pool hook = writeIORef (pOnRelease pool) $! hook++-- Internal ----------------------------------------------------------------++-- | What the single STM transaction decided we should do.+data AcquireAction+  = GotIdle !PoolEntry+  | CreateNew+  | MustWait !(TMVar (Either PgWireError Connection))+  | PoolIsClosed++-- | Atomically decide what to do: take idle, create new, or enqueue waiter.+-- This is the core fix for the acquire race condition — a single STM+-- transaction replaces the old two-step check.+acquireAction :: Pool -> STM AcquireAction+acquireAction pool = do+  closed <- readTVar (pClosed pool)+  if closed+    then pure PoolIsClosed+    else do+      idle <- readTVar (pIdle pool)+      case takeIdle (poolQueueMode (pConfig pool)) idle of+        Just (entry, rest) -> do+          writeTVar (pIdle pool) rest+          pure (GotIdle entry)+        Nothing -> do+          active <- readTVar (pActive pool)+          maxSz <- readTVar (pEffectiveSize pool)+          if active < maxSz+            then do+              writeTVar (pActive pool) (active + 1)+              pure CreateNew+            else do+              waiter <- newEmptyTMVar+              modifyTVar' (pWaiters pool) (Seq.|> waiter)+              pure (MustWait waiter)++-- | Take from the appropriate end based on QueueMode.+takeIdle :: QueueMode -> Seq PoolEntry -> Maybe (PoolEntry, Seq PoolEntry)+takeIdle QueueLIFO s = case Seq.viewr s of+  Seq.EmptyR -> Nothing+  rest Seq.:> e -> Just (e, rest)+takeIdle QueueFIFO s = case Seq.viewl s of+  Seq.EmptyL -> Nothing+  e Seq.:< rest -> Just (e, rest)++acquire :: Pool -> IO Connection+acquire pool = acquireWithTimeout pool (poolAcquireTimeout (pConfig pool))++acquireWithTimeout :: Pool -> NominalDiffTime -> IO Connection+acquireWithTimeout pool timeout = do+  act <- atomically (acquireAction pool)+  case act of+    PoolIsClosed -> throwPgWire PoolClosed+    GotIdle entry -> tryRecycle pool entry+    CreateNew -> createConnection pool+    MustWait waiter -> waitForConnectionWith pool waiter timeout++-- | Try to recycle an idle connection. If it's expired or unhealthy,+-- destroy it and retry.+tryRecycle :: Pool -> PoolEntry -> IO Connection+tryRecycle pool entry = do+  now <- getCurrentTime+  expired <- isExpiredIO pool entry now+  if expired+    then do+      destroyEntry pool "expired" entry+      acquire pool+    else do+      ok <- recycleCheck pool entry now+      if ok+        then do+          logPool pool "debug" "recycled connection"+          observe pool ConnectionRecycled+          writeIORef (peLastUsed entry) now+          runHook (pOnAcquire pool) (peConn entry)+          pure (peConn entry)+        else do+          observe pool HealthCheckFailed+          destroyEntry pool "unhealthy" entry+          acquire pool++-- | Recycling check based on the configured method.+recycleCheck :: Pool -> PoolEntry -> UTCTime -> IO Bool+recycleCheck pool entry now = case poolRecyclingMethod (pConfig pool) of+  RecycleFast -> do+    readTVarIO (awcAlive (connAsync (peConn entry)))+  RecycleVerified -> do+    alive <- readTVarIO (awcAlive (connAsync (peConn entry)))+    if not alive+      then pure False+      else do+        lastUsed <- readIORef (peLastUsed entry)+        if diffUTCTime now lastUsed <= poolHealthCheckAge (pConfig pool)+          then pure True+          else checkHealthWithTimeout (poolAcquireTimeout (pConfig pool)) (peConn entry)+  RecycleClean -> do+    alive <- readTVarIO (awcAlive (connAsync (peConn entry)))+    if not alive+      then pure False+      else do+        let timeoutSecs = poolAcquireTimeout (pConfig pool)+            timeoutMicros = round (timeoutSecs * 1000000) :: Int+        result <- race (threadDelay timeoutMicros)+                       (try @SomeException (simpleQuery (peConn entry) "DISCARD ALL"))+        case result of+          Right (Right _) -> do+            -- DISCARD ALL deallocates all prepared statements on the server,+            -- so the client-side cache must be cleared to stay in sync.+            writeIORef (connStmtCache (peConn entry)) PSQ.empty+            pure True+          _ -> pure False++createConnection :: Pool -> IO Connection+createConnection pool = do+  conn <- connectWithRetry (poolConnectionRetries (pConfig pool)) 100000+    `onException` atomically (modifyTVar' (pActive pool) (subtract 1))+  now <- getCurrentTime+  deadline <- jitteredDeadline (pConfig pool) now+  atomically $ do+    modifyTVar' (pTotalCreated pool) (+ 1)+    modifyTVar' (pConnMeta pool) (Map.insert (connBackendPid conn) (ConnMeta now deadline))+  logPool pool "debug" "created connection"+  observe pool ConnectionCreated+  runHook (pOnCreate pool) conn+  runHook (pOnAcquire pool) conn+  pure conn+  where+    -- Exponential backoff: 100ms, 200ms, 400ms, ...+    connectWithRetry :: Int -> Int -> IO Connection+    connectWithRetry 0 _ = connectString (poolConnString (pConfig pool))+    connectWithRetry !retriesLeft !delayMicros = do+      result <- try @SomeException (connectString (poolConnString (pConfig pool)))+      case result of+        Right c -> pure c+        Left _ -> do+          logPool pool "warn" ("connection failed, retrying in " <> BS8.pack (show (delayMicros `div` 1000)) <> "ms")+          threadDelay delayMicros+          connectWithRetry (retriesLeft - 1) (delayMicros * 2)++-- | Compute a jittered deadline for a connection.+-- deadline = now + maxLife + uniform(-jitter, +jitter)+jitteredDeadline :: PoolConfig -> UTCTime -> IO UTCTime+jitteredDeadline cfg now = do+  let jitter = poolMaxLifeJitter cfg+  offset <- if jitter > 0+    then do+      let jitterMicros = round (jitter * 1000000) :: Int+      r <- randomRIO (negate jitterMicros, jitterMicros)+      pure (fromIntegral r / 1000000 :: NominalDiffTime)+    else pure 0+  pure (addUTCTime (poolMaxLife cfg + offset) now)++waitForConnectionWith :: Pool -> TMVar (Either PgWireError Connection) -> NominalDiffTime -> IO Connection+waitForConnectionWith pool waiter timeout = do+  let timeoutMicros = round (timeout * 1000000) :: Int+  result <- race+    (threadDelay timeoutMicros)+    (atomically $ takeTMVar waiter)+  case result of+    Left () -> do+      atomically $ modifyTVar' (pTotalTimeouts pool) (+ 1)+      logPool pool "warn" "acquire timeout"+      observe pool AcquireTimeout+      throwPgWire PoolTimeout+    Right (Left err) -> throwPgWire err+    Right (Right conn) -> do+      runHook (pOnAcquire pool) conn+      pure conn++release :: Pool -> Connection -> IO ()+release pool conn = do+  runHook (pOnRelease pool) conn+  now <- getCurrentTime+  -- Try to deliver to a waiter first. Use tryPutTMVar and loop+  -- to handle dead waiters (waiter died between enqueue and delivery).+  delivered <- deliverToWaiter pool conn+  if delivered+    then pure ()+    else do+      -- Look up creation time from the tracking map (fixes maxLife bug)+      meta <- Map.lookup (connBackendPid conn) <$> readTVarIO (pConnMeta pool)+      let createdAt = maybe now cmCreatedAt meta+      lastUsed <- newIORef now+      let !entry = PoolEntry conn createdAt lastUsed+      atomically $ modifyTVar' (pIdle pool) (Seq.|> entry)++-- | Try to deliver a connection to a waiting thread. Loops through waiters,+-- skipping any that have died (timed out). Returns True if delivered.+deliverToWaiter :: Pool -> Connection -> IO Bool+deliverToWaiter pool conn = atomically $ do+  ws <- readTVar (pWaiters pool)+  case Seq.viewl ws of+    Seq.EmptyL -> pure False+    w Seq.:< rest -> do+      delivered <- tryPutTMVar w (Right conn)+      if delivered+        then do+          writeTVar (pWaiters pool) rest+          pure True+        else do+          -- Waiter died (timeout). Skip it and try next.+          writeTVar (pWaiters pool) rest+          -- Retry with remaining waiters (re-enter STM)+          deliverLoop pool rest conn++-- | STM loop to deliver to next available waiter.+deliverLoop :: Pool -> Seq (TMVar (Either PgWireError Connection)) -> Connection -> STM Bool+deliverLoop _ Empty _ = pure False+deliverLoop pool (w :<| rest) conn = do+  delivered <- tryPutTMVar w (Right conn)+  if delivered+    then do+      writeTVar (pWaiters pool) rest+      pure True+    else do+      writeTVar (pWaiters pool) rest+      deliverLoop pool rest conn++destroyConn :: Pool -> Connection -> BS8.ByteString -> IO ()+destroyConn pool conn reason = do+  safeClose conn+  atomically $ do+    modifyTVar' (pActive pool) (subtract 1)+    modifyTVar' (pTotalDestroyed pool) (+ 1)+    modifyTVar' (pConnMeta pool) (Map.delete (connBackendPid conn))+  logPool pool "debug" ("destroyed connection: " <> reason)+  observe pool (ConnectionDestroyed reason)++destroyEntry :: Pool -> BS8.ByteString -> PoolEntry -> IO ()+destroyEntry pool reason entry = destroyConn pool (peConn entry) reason++-- | Check if a connection has exceeded its (jittered) max-life.+isExpiredIO :: Pool -> PoolEntry -> UTCTime -> IO Bool+isExpiredIO pool entry now = do+  meta <- Map.lookup (connBackendPid (peConn entry)) <$> readTVarIO (pConnMeta pool)+  pure $ case meta of+    Just cm -> now >= cmDeadline cm+    Nothing -> diffUTCTime now (peCreatedAt entry) > poolMaxLife (pConfig pool)++safeClose :: Connection -> IO ()+safeClose conn = close conn `catch` \(_ :: SomeException) -> pure ()++-- | Lightweight health check: send an empty query and see if we get a response.+-- Times out after the given duration to prevent a hung backend from blocking acquire.+checkHealthWithTimeout :: NominalDiffTime -> Connection -> IO Bool+checkHealthWithTimeout timeout conn = do+  let timeoutMicros = round (timeout * 1000000) :: Int+  result <- race (threadDelay timeoutMicros) (try @SomeException (simpleQuery conn ""))+  pure $ case result of+    Left () -> False    -- timed out+    Right (Right _) -> True+    Right (Left _) -> False++-- | Run a lifecycle hook, catching and ignoring exceptions so a hook+-- failure doesn't break the pool.+runHook :: IORef (Connection -> IO ()) -> Connection -> IO ()+runHook ref conn = do+  hook <- readIORef ref+  hook conn `catch` \(_ :: SomeException) -> pure ()++-- | Log a pool event using the configured logger.+logPool :: Pool -> BS8.ByteString -> BS8.ByteString -> IO ()+logPool pool = poolLogger (pConfig pool)++-- | Emit a structured observation event, catching and ignoring exceptions.+observe :: Pool -> PoolEvent -> IO ()+observe pool event =+  poolObserver (pConfig pool) event `catch` \(_ :: SomeException) -> pure ()++-- | Background thread that periodically reaps idle/expired connections.+--+-- Atomically drains the idle queue, filters in IO (reading IORefs),+-- then atomically merges kept entries back. This avoids the race where+-- acquire/release modify the idle queue between read and write.+reaperThread :: Pool -> IO ()+reaperThread pool = go+  where+    intervalMicros = round (poolReaperInterval (pConfig pool) * 1000000) :: Int++    go = do+      threadDelay intervalMicros+      reap+        `catch` \(e :: AsyncException) -> throwIO e+        `catch` \(_ :: SomeException) ->+          logPool pool "warn" "reaper: exception in sweep cycle (continuing)"+      go++    reap = do+      now <- getCurrentTime+      -- Atomically drain the idle queue+      entries <- atomically $ do+        idle <- readTVar (pIdle pool)+        writeTVar (pIdle pool) Seq.empty+        pure (seqToList idle)+      (keep, toReap) <- partitionM (shouldKeep now) entries+      -- Merge kept entries back (prepend; entries added by release+      -- during filtering are already at the tail)+      atomically $ modifyTVar' (pIdle pool) (Seq.fromList keep Seq.><)+      let reapCount = length toReap+      mapM_ (destroyEntry pool "reaped") toReap+      when (reapCount > 0) $ do+        logPool pool "info" ("reaper swept " <> BS8.pack (show reapCount) <> " connections")+        observe pool (ReaperSwept reapCount)++    shouldKeep now entry = do+      expired <- isExpiredIO pool entry now+      if expired+        then pure False+        else do+          lastUsed <- readIORef (peLastUsed entry)+          pure (diffUTCTime now lastUsed <= poolIdleTime (pConfig pool))++-- | Background thread that maintains 'poolMinIdle' connections.+warmerThread :: Pool -> IO ()+warmerThread pool = go+  where+    intervalMicros = round (poolReaperInterval (pConfig pool) * 1000000) :: Int++    go = do+      threadDelay intervalMicros+      warm+        `catch` \(e :: AsyncException) -> throwIO e+        `catch` \(_ :: SomeException) ->+          logPool pool "warn" "warmer: exception in warm cycle (continuing)"+      go++    warm = do+      deficit <- atomically $ do+        closed <- readTVar (pClosed pool)+        if closed+          then pure 0+          else do+            idle <- Seq.length <$> readTVar (pIdle pool)+            active <- readTVar (pActive pool)+            maxSz <- readTVar (pEffectiveSize pool)+            let minIdle = poolMinIdle (pConfig pool)+                canCreate = maxSz - active+                need = max 0 (minIdle - idle)+            pure (min need canCreate)+      created <- warmN 0 deficit+      when (created > 0) $ do+        logPool pool "info" ("warmer created " <> BS8.pack (show created) <> " connections")+        observe pool (WarmerCreated created)++    warmN :: Int -> Int -> IO Int+    warmN !acc 0 = pure acc+    warmN !acc n = do+      result <- try @SomeException $ do+        -- Reserve a slot+        reserved <- atomically $ do+          active <- readTVar (pActive pool)+          maxSz <- readTVar (pEffectiveSize pool)+          if active < maxSz+            then do+              writeTVar (pActive pool) (active + 1)+              pure True+            else pure False+        if reserved+          then do+            conn <- connectString (poolConnString (pConfig pool))+              `onException` atomically (modifyTVar' (pActive pool) (subtract 1))+            now <- getCurrentTime+            deadline <- jitteredDeadline (pConfig pool) now+            atomically $ do+              modifyTVar' (pTotalCreated pool) (+ 1)+              modifyTVar' (pConnMeta pool) (Map.insert (connBackendPid conn) (ConnMeta now deadline))+            runHook (pOnCreate pool) conn+            lastUsed <- newIORef now+            let !entry = PoolEntry conn now lastUsed+            atomically $ modifyTVar' (pIdle pool) (Seq.|> entry)+            pure True+          else pure False+      case result of+        Left _ -> pure acc -- Connection failed, stop warming this round+        Right True -> warmN (acc + 1) (n - 1)+        Right False -> pure acc++when :: Bool -> IO () -> IO ()+when True f = f+when False _ = pure ()++seqToList :: Seq a -> [a]+seqToList = foldr (:) []++partitionM :: (a -> IO Bool) -> [a] -> IO ([a], [a])+partitionM _ [] = pure ([], [])+partitionM p (x : xs) = do+  b <- p x+  (!ys, !zs) <- partitionM p xs+  pure (if b then (x : ys, zs) else (ys, x : zs))
+ src/PgWire/Pool/Config.hs view
@@ -0,0 +1,141 @@+-- | Configuration types for the connection pool.+--+-- Use 'defaultPoolConfig' as a starting point and override fields:+--+-- @+-- cfg = 'defaultPoolConfig'+--   { poolConnString = \"postgres:\/\/user:pass\@localhost\/mydb\"+--   , poolSize = 20+--   , poolRecyclingMethod = 'RecycleVerified'+--   }+-- @+module PgWire.Pool.Config+  ( PoolConfig (..)+  , defaultPoolConfig+  , RecyclingMethod (..)+  , QueueMode (..)+  , PoolLogger+  , nullPoolLogger+  ) where++import Data.ByteString (ByteString)+import Data.Time (NominalDiffTime)+import PgWire.Pool.Observation (PoolObserver, nullObserver)++-- | How to validate a connection before handing it to the caller.+data RecyclingMethod+  = RecycleFast+  -- ^ Check @awcAlive@ TVar only — free, no I/O.+  | RecycleVerified+  -- ^ Send an empty query if the connection has been idle longer than+  -- 'poolHealthCheckAge'. Falls back to 'RecycleFast' for recently used+  -- connections.+  | RecycleClean+  -- ^ Send @DISCARD ALL@ before every reuse (resets all session state).+  deriving stock (Show, Eq)++-- | Which end of the idle queue to draw from.+data QueueMode+  = QueueFIFO+  -- ^ Oldest idle connection first — spreads load across backends.+  | QueueLIFO+  -- ^ Newest idle connection first — keeps fewer total connections alive.+  deriving stock (Show, Eq)++-- | A pool logging callback. The pool calls this with a severity string+-- (@\"debug\"@, @\"info\"@, @\"warn\"@) and a message.+type PoolLogger = ByteString -> ByteString -> IO ()++-- | Logger that discards all events.+nullPoolLogger :: PoolLogger+nullPoolLogger _ _ = pure ()++-- | Configuration for a connection pool. Use 'defaultPoolConfig' and+-- override the fields you need:+--+-- @+-- pool <- 'PgWire.Pool.newPool' defaultPoolConfig+--   { poolConnString = \"postgres:\/\/user:pass\@localhost:5432\/mydb\"+--   , poolSize = 20+--   }+-- @+data PoolConfig = PoolConfig+  { poolConnString :: ByteString+  -- ^ PostgreSQL connection string. Required.+  , poolSize :: !Int+  -- ^ Maximum number of connections. Default: 10.+  , poolIdleTime :: !NominalDiffTime+  -- ^ Close idle connections after this duration. Default: 600s (10 min).+  , poolMaxLife :: !NominalDiffTime+  -- ^ Close connections after this total lifetime. Default: 3600s (1 hour).+  , poolAcquireTimeout :: !NominalDiffTime+  -- ^ Maximum time to wait for a connection. Throws 'PgWire.Error.PoolTimeout'+  -- if exceeded. Default: 10s.+  , poolMinIdle :: !Int+  -- ^ Minimum idle connections to maintain. A background warmer thread+  -- creates connections to fill the deficit. Default: 0 (no warming).+  , poolReaperInterval :: !NominalDiffTime+  -- ^ How often the reaper sweeps idle/expired connections. Default: 30s.+  , poolRecyclingMethod :: !RecyclingMethod+  -- ^ How to validate connections before reuse. Default: 'RecycleFast'.+  , poolQueueMode :: !QueueMode+  -- ^ Which end of the idle queue to draw from. Default: 'QueueLIFO'.+  , poolHealthCheckAge :: !NominalDiffTime+  -- ^ Only verify (empty query) if a connection has been idle longer than+  -- this. Only relevant when 'poolRecyclingMethod' is 'RecycleVerified'.+  -- Default: 5s.+  , poolMaxLifeJitter :: !NominalDiffTime+  -- ^ Randomize 'poolMaxLife' ± this amount to avoid thundering herd+  -- expiration. Default: 60s.+  , poolConnectionRetries :: !Int+  -- ^ Number of retries when creating a connection fails. Retries use+  -- exponential backoff (100ms, 200ms, 400ms, ...). Default: 3.+  , poolLogger :: PoolLogger+  -- ^ Logging callback. Default: 'nullPoolLogger' (discards all events).+  , poolObserver :: PoolObserver+  -- ^ Structured event callback for metrics and monitoring.+  -- Called for every significant pool state transition.+  -- Default: 'nullObserver' (discards all events).+  -- See "PgWire.Pool.Observation" for the event types.+  }++-- Manual Show instance that omits the function field.+instance Show PoolConfig where+  show cfg = unwords+    [ "PoolConfig"+    , "{poolConnString=" <> show (poolConnString cfg)+    , ",poolSize=" <> show (poolSize cfg)+    , ",poolIdleTime=" <> show (poolIdleTime cfg)+    , ",poolMaxLife=" <> show (poolMaxLife cfg)+    , ",poolAcquireTimeout=" <> show (poolAcquireTimeout cfg)+    , ",poolMinIdle=" <> show (poolMinIdle cfg)+    , ",poolReaperInterval=" <> show (poolReaperInterval cfg)+    , ",poolRecyclingMethod=" <> show (poolRecyclingMethod cfg)+    , ",poolQueueMode=" <> show (poolQueueMode cfg)+    , ",poolHealthCheckAge=" <> show (poolHealthCheckAge cfg)+    , ",poolMaxLifeJitter=" <> show (poolMaxLifeJitter cfg)+    , ",poolConnectionRetries=" <> show (poolConnectionRetries cfg)+    , ",poolLogger=<function>"+    , ",poolObserver=<function>}"+    ]++-- | Sensible defaults: 10 connections, 10-minute idle timeout, 1-hour max life,+-- 10-second acquire timeout, 'RecycleFast', 'QueueLIFO', no warming, no logging.+defaultPoolConfig :: PoolConfig+defaultPoolConfig =+  PoolConfig+    { poolConnString = ""+    , poolSize = 10+    , poolIdleTime = 600+    , poolMaxLife = 3600+    , poolAcquireTimeout = 10+    , poolMinIdle = 0+    , poolReaperInterval = 30+    , poolRecyclingMethod = RecycleFast+    , poolQueueMode = QueueLIFO+    , poolHealthCheckAge = 5+    , poolMaxLifeJitter = 60+    , poolConnectionRetries = 3+    , poolLogger = nullPoolLogger+    , poolObserver = nullObserver+    }
+ src/PgWire/Pool/Observation.hs view
@@ -0,0 +1,60 @@+-- | Structured pool lifecycle events for observability.+--+-- Set an observation handler via 'PgWire.Pool.Config.poolObserver' to+-- receive events for every significant pool state transition. Use these+-- for metrics, structured logging, alerting, or debugging.+--+-- @+-- pool <- newPool defaultPoolConfig+--   { poolObserver = \\event -> do+--       case event of+--         ConnectionCreated -> incrementCounter \"pool.created\"+--         ConnectionDestroyed reason -> logInfo $ \"destroyed: \" <> reason+--         AcquireTimeout -> incrementCounter \"pool.timeouts\"+--         _ -> pure ()+--   }+-- @+module PgWire.Pool.Observation+  ( PoolEvent (..)+  , PoolObserver+  , nullObserver+  ) where++import Data.ByteString (ByteString)+import Data.Time (NominalDiffTime)++-- | A structured event emitted by the connection pool.+data PoolEvent+  = ConnectionCreated+  -- ^ A new connection was established.+  | ConnectionDestroyed !ByteString+  -- ^ A connection was closed. The 'ByteString' is the reason+  -- (e.g., @\"expired\"@, @\"unhealthy\"@, @\"reaped\"@, @\"pool closed\"@,+  -- @\"resize\"@, @\"exception\"@, @\"retain\"@).+  | ConnectionAcquired !NominalDiffTime+  -- ^ A connection was handed to a caller. The 'NominalDiffTime' is+  -- how long the caller waited (0 if a connection was immediately available).+  | ConnectionReleased+  -- ^ A connection was returned to the pool.+  | ConnectionRecycled+  -- ^ An idle connection passed its recycling check and was reused.+  | HealthCheckFailed+  -- ^ A connection failed its health check and was destroyed.+  | AcquireTimeout+  -- ^ A caller timed out waiting for a connection.+  | ReaperSwept !Int+  -- ^ The background reaper closed N idle/expired connections.+  | WarmerCreated !Int+  -- ^ The background warmer created N connections to fill min-idle.+  | PoolResized !Int !Int+  -- ^ Pool was resized from old size to new size.+  | PoolShutdown+  -- ^ 'closePool' was called.+  deriving stock (Show, Eq)++-- | Callback for pool events. Set in 'PgWire.Pool.Config.PoolConfig'.+type PoolObserver = PoolEvent -> IO ()++-- | Observer that discards all events.+nullObserver :: PoolObserver+nullObserver _ = pure ()
+ src/PgWire/Protocol/Backend.hs view
@@ -0,0 +1,121 @@+-- | Decoded backend (server-to-client) message types from the+-- PostgreSQL v3 wire protocol.+--+-- __Audience:__ exposed for downstream library authors that want to+-- decode raw protocol responses or build alternative pipeline+-- abstractions. Most application code only sees these types+-- indirectly via "PgWire.Connection" results and "PgWire.Error".+module PgWire.Protocol.Backend+  ( BackendMsg (..)+  , FieldInfo (..)+  , PgError (..)+  , PgNotice (..)+  , TxStatus (..)+  , CommandTag (..)+  , AuthType (..)+  ) where++import Data.ByteString (ByteString)+import Data.Int (Int16, Int32, Int64)+import Data.Vector (Vector)+import Data.Word (Word32)+import GHC.Generics (Generic)+import NoThunks.Class (NoThunks)++-- | Transaction status indicator.+data TxStatus+  = TxIdle+  | TxInTransaction+  | TxFailed+  deriving stock (Show, Eq)++-- | Parsed command completion tag.+data CommandTag+  = SelectTag Int64+  | InsertTag Int64+  | UpdateTag Int64+  | DeleteTag Int64+  | OtherTag ByteString+  deriving stock (Show, Eq)++-- | Authentication sub-types.+data AuthType+  = AuthOk+  | AuthCleartextPassword+  | AuthMD5Password ByteString -- 4-byte salt+  | AuthSASL [ByteString] -- mechanism names+  | AuthSASLContinue ByteString -- server data+  | AuthSASLFinal ByteString -- server signature+  deriving stock (Show, Eq)++-- | Column metadata from RowDescription.+data FieldInfo = FieldInfo+  { fiName :: ByteString+  , fiTableOid :: {-# UNPACK #-} !Word32+  , fiColumnNum :: {-# UNPACK #-} !Int16+  , fiTypeOid :: {-# UNPACK #-} !Word32+  , fiTypeSize :: {-# UNPACK #-} !Int16+  , fiTypeMod :: {-# UNPACK #-} !Int32+  , fiFormatCode :: {-# UNPACK #-} !Int16+  }+  deriving stock (Show, Eq)++-- | Structured error from ErrorResponse.+-- See https://www.postgresql.org/docs/current/protocol-error-fields.html+data PgError = PgError+  { pgSeverity :: ByteString+  , pgCode :: ByteString+  , pgMessage :: ByteString+  , pgDetail :: Maybe ByteString+  , pgHint :: Maybe ByteString+  , pgPosition :: !(Maybe Int32)+  , pgInternalPosition :: !(Maybe Int32)+  , pgInternalQuery :: Maybe ByteString+  , pgWhere :: Maybe ByteString+  , pgSchema :: Maybe ByteString+  , pgTable :: Maybe ByteString+  , pgColumn :: Maybe ByteString+  , pgDataType :: Maybe ByteString+  , pgConstraint :: Maybe ByteString+  , pgFile :: Maybe ByteString+  , pgLine :: !(Maybe Int32)+  , pgRoutine :: Maybe ByteString+  }+  deriving stock (Show, Eq, Generic)++instance NoThunks PgError++-- | Structured notice from NoticeResponse.+newtype PgNotice = PgNotice {unPgNotice :: PgError}+  deriving stock (Show, Eq)++-- | Server-to-client (backend) wire protocol messages.+data BackendMsg+  = Authentication AuthType+  | ParameterStatus ByteString ByteString+  | BackendKeyData Int32 Int32+  | ReadyForQuery TxStatus+  | ParseComplete+  | BindComplete+  | CloseComplete+  | NoData+  | EmptyQueryResponse+  | ParameterDescription (Vector Word32)+  | RowDescription (Vector FieldInfo)+  | DataRow (Vector (Maybe ByteString))+  | CommandComplete CommandTag+  | ErrorResponse PgError+  | NoticeResponse PgNotice+  | NotificationResponse Int32 ByteString ByteString+  | -- | CopyInResponse: server is ready to receive COPY data+    CopyInResponse Int8Format (Vector Int16)+  | -- | CopyOutResponse: server will send COPY data+    CopyOutResponse Int8Format (Vector Int16)+  | -- | CopyData: a chunk of COPY data from the server+    CopyDataMsg ByteString+  | -- | CopyDone: end of COPY data from server+    CopyDoneMsg+  deriving stock (Show, Eq)++-- | COPY format code (0 = text, 1 = binary)+type Int8Format = Int16
+ src/PgWire/Protocol/Builders.hs view
@@ -0,0 +1,177 @@+-- | Encoders for frontend (client-to-server) messages.+--+-- __Audience:__ exposed for downstream library authors that need to+-- emit raw protocol bytes (e.g. custom pipelining or proxies). Most+-- application code drives these encoders indirectly through+-- "PgWire.Connection".+module PgWire.Protocol.Builders+  ( buildFrontendMsg+  , buildFrontendMsgsConcat+  , buildStartup+  ) where++import Data.Bits (shiftL)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Builder (Builder)+import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int16, Int32)+import Data.Vector (Vector)+import Data.Vector qualified as V+import PgWire.Protocol.Frontend++-- | Build a complete frontend message as a strict 'ByteString'.+buildFrontendMsg :: FrontendMsg -> ByteString+buildFrontendMsg = LBS.toStrict . B.toLazyByteString . encodeFrontendMsg+{-# INLINE buildFrontendMsg #-}++-- | Fuse multiple frontend messages into a single strict 'ByteString'.+-- Uses one 'Builder' pass and one 'toStrict' call, eliminating N-1+-- intermediate allocations compared to mapping 'buildFrontendMsg'.+buildFrontendMsgsConcat :: [FrontendMsg] -> ByteString+buildFrontendMsgsConcat = LBS.toStrict . B.toLazyByteString . foldMap encodeFrontendMsg+{-# INLINE buildFrontendMsgsConcat #-}++-- | Build the startup message (special: no tag byte).+buildStartup :: StartupParams -> ByteString+buildStartup = LBS.toStrict . B.toLazyByteString . encodeStartup++-- Encoding ----------------------------------------------------------------+--+-- Each message is [tag :: Word8] [length :: Int32] [payload].+-- The length field includes itself (4 bytes) but not the tag byte.+--+-- We pre-compute the payload size from the message fields, then write+-- tag + length + payload in a single Builder pass. This avoids the+-- previous approach of materializing the payload to a ByteString just+-- to compute its length (which caused a double-copy).++encodeFrontendMsg :: FrontendMsg -> Builder+encodeFrontendMsg = \case+  Startup params -> encodeStartup params+  Parse name sql oids ->+    let !sz = cstringSize name + cstringSize sql + 2 + V.length oids * 4+     in tag 'P' sz <> cstring name <> cstring sql+          <> B.int16BE (fromIntegral (V.length oids) :: Int16)+          <> V.foldl' (\b oid -> b <> B.word32BE oid) mempty oids+  Bind portal stmt pfmts vals rfmts ->+    let !sz = cstringSize portal + cstringSize stmt+            + formatCodesSize pfmts+            + 2 + paramValuesSize vals+            + formatCodesSize rfmts+     in tag 'B' sz <> cstring portal <> cstring stmt+          <> encodeFormatCodes pfmts+          <> B.int16BE (fromIntegral (V.length vals) :: Int16)+          <> V.foldl' (\b mv -> b <> encodeParamValue mv) mempty vals+          <> encodeFormatCodes rfmts+  Describe target name ->+    let !sz = 1 + cstringSize name+     in tag 'D' sz <> B.char8 (targetChar target) <> cstring name+  Execute portal maxRows ->+    let !sz = cstringSize portal + 4+     in tag 'E' sz <> cstring portal <> B.int32BE maxRows+  Close target name ->+    let !sz = 1 + cstringSize name+     in tag 'C' sz <> B.char8 (targetChar target) <> cstring name+  Sync -> tag 'S' 0+  Flush -> tag 'H' 0+  Terminate -> tag 'X' 0+  Query sql ->+    let !sz = cstringSize sql+     in tag 'Q' sz <> cstring sql+  PasswordMessage pw ->+    let !sz = cstringSize pw+     in tag 'p' sz <> cstring pw+  SASLInitialResponse mech clientFirst ->+    let !sz = cstringSize mech + 4 + BS.length clientFirst+     in tag 'p' sz <> cstring mech+          <> B.int32BE (fromIntegral (BS.length clientFirst))+          <> B.byteString clientFirst+  SASLResponse msg ->+    let !sz = BS.length msg+     in tag 'p' sz <> B.byteString msg+  CopyData dat ->+    let !sz = BS.length dat+     in tag 'd' sz <> B.byteString dat+  CopyDone -> tag 'c' 0+  CopyFail msg ->+    let !sz = cstringSize msg+     in tag 'f' sz <> cstring msg+{-# INLINE encodeFrontendMsg #-}++-- | Write tag byte + length (payload size + 4 for the length field itself).+tag :: Char -> Int -> Builder+tag c payloadSize =+  B.char8 c <> B.int32BE (fromIntegral (payloadSize + 4) :: Int32)+{-# INLINE tag #-}++-- Size computation helpers ------------------------------------------------++-- | Size of a NUL-terminated C string: data bytes + 1 NUL byte.+cstringSize :: ByteString -> Int+cstringSize bs = BS.length bs + 1+{-# INLINE cstringSize #-}++-- | Size of format codes: 2 (count) + 2 per code.+formatCodesSize :: Vector FormatCode -> Int+formatCodesSize fmts = 2 + V.length fmts * 2+{-# INLINE formatCodesSize #-}++-- | Size of parameter values: sum of (4 + data) per non-NULL, 4 per NULL.+paramValuesSize :: Vector (Maybe ByteString) -> Int+paramValuesSize = V.foldl' (\acc mv -> acc + paramValueSize mv) 0+{-# INLINE paramValuesSize #-}++paramValueSize :: Maybe ByteString -> Int+paramValueSize Nothing = 4+paramValueSize (Just bs) = 4 + BS.length bs+{-# INLINE paramValueSize #-}++-- Payload encoding helpers ------------------------------------------------++encodeFormatCodes :: Vector FormatCode -> Builder+encodeFormatCodes fmts =+  B.int16BE (fromIntegral (V.length fmts) :: Int16)+    <> V.foldl' (\b fc -> b <> B.int16BE (formatCodeToInt16 fc)) mempty fmts+{-# INLINE encodeFormatCodes #-}++formatCodeToInt16 :: FormatCode -> Int16+formatCodeToInt16 TextFormat = 0+formatCodeToInt16 BinaryFormat = 1+{-# INLINE formatCodeToInt16 #-}++encodeParamValue :: Maybe ByteString -> Builder+encodeParamValue Nothing = B.int32BE (-1)+encodeParamValue (Just bs) =+  B.int32BE (fromIntegral (BS.length bs))+    <> B.byteString bs+{-# INLINE encodeParamValue #-}++targetChar :: DescribeTarget -> Char+targetChar DescribeStatement = 'S'+targetChar DescribePortal = 'P'+{-# INLINE targetChar #-}++-- | NUL-terminated C string.+cstring :: ByteString -> Builder+cstring bs = B.byteString bs <> B.word8 0+{-# INLINE cstring #-}++-- Startup message ---------------------------------------------------------+-- Startup has a different format: [length :: Int32] [protocol] [params] [NUL]+-- No tag byte. We still need to materialize to compute the length here,+-- but startup is called once per connection, not on the hot path.++encodeStartup :: StartupParams -> Builder+encodeStartup StartupParams {..} =+  let params =+        cstring "user" <> cstring spUser+          <> cstring "database" <> cstring spDatabase+          <> (if BS.null spAppName then mempty else cstring "application_name" <> cstring spAppName)+          <> mconcat [cstring k <> cstring v | (k, v) <- spExtraParams]+          <> B.word8 0 -- terminator+      paramsBs = LBS.toStrict (B.toLazyByteString params)+      len = fromIntegral (4 + 4 + BS.length paramsBs) :: Int32+      protocolVersion = (3 :: Int32) `shiftL` 16 -- 3.0+   in B.int32BE len <> B.int32BE protocolVersion <> B.byteString paramsBs
+ src/PgWire/Protocol/Frontend.hs view
@@ -0,0 +1,94 @@+-- | Frontend (client-to-server) message types from the PostgreSQL v3+-- wire protocol.+--+-- __Audience:__ exposed for downstream library authors that need to+-- assemble protocol messages directly. Most application code uses+-- "PgWire.Connection" instead.+module PgWire.Protocol.Frontend+  ( FrontendMsg (..)+  , StartupParams (..)+  , DescribeTarget (..)+  , FormatCode (..)+  ) where++import Data.ByteString (ByteString)+import Data.Int (Int32)+import Data.Vector (Vector)+import Data.Word (Word32)++-- | Format codes for parameter/result encoding.+data FormatCode = TextFormat | BinaryFormat+  deriving stock (Show, Eq)++-- | Target for a Describe message.+data DescribeTarget = DescribeStatement | DescribePortal+  deriving stock (Show, Eq)++-- | Parameters for the startup message.+data StartupParams = StartupParams+  { spUser :: ByteString+  , spDatabase :: ByteString+  , spAppName :: ByteString+  , spExtraParams :: [(ByteString, ByteString)]+  }+  deriving stock (Show)++-- | Client-to-server (frontend) wire protocol messages.+data FrontendMsg+  = -- | Initial startup (no tag byte)+    Startup StartupParams+  | -- | Parse a prepared statement+    Parse+      ByteString+      -- ^ Statement name (empty = unnamed)+      ByteString+      -- ^ SQL query+      (Vector Word32)+      -- ^ Parameter type OIDs (0 = let server infer)+  | -- | Bind parameters to a prepared statement+    Bind+      ByteString+      -- ^ Portal name (empty = unnamed)+      ByteString+      -- ^ Statement name+      (Vector FormatCode)+      -- ^ Parameter format codes+      (Vector (Maybe ByteString))+      -- ^ Parameter values (Nothing = NULL)+      (Vector FormatCode)+      -- ^ Result format codes+  | -- | Describe a prepared statement or portal+    Describe DescribeTarget ByteString+  | -- | Execute a portal+    Execute+      ByteString+      -- ^ Portal name+      Int32+      -- ^ Max rows (0 = unlimited)+  | -- | Close a prepared statement or portal+    Close DescribeTarget ByteString+  | -- | Sync — end of an extended query sequence+    Sync+  | -- | Flush — request server to send pending output+    Flush+  | -- | Terminate the connection+    Terminate+  | -- | Simple query protocol+    Query ByteString+  | -- | Password message (cleartext or MD5)+    PasswordMessage ByteString+  | -- | SASL initial response+    SASLInitialResponse+      ByteString+      -- ^ Mechanism name+      ByteString+      -- ^ Client first message+  | -- | SASL response (subsequent messages)+    SASLResponse ByteString+  | -- | CopyData — a chunk of COPY data+    CopyData ByteString+  | -- | CopyDone — end of COPY data from client+    CopyDone+  | -- | CopyFail — client aborts COPY+    CopyFail ByteString+  deriving stock (Show)
+ src/PgWire/Protocol/Oid.hs view
@@ -0,0 +1,143 @@+-- | PostgreSQL type OIDs and convenience constants.+--+-- Used when writing custom 'PgWire.Binary.Types.PgEncode' /+-- 'PgWire.Binary.Types.PgDecode' instances and when inspecting+-- raw 'PgWire.Protocol.Backend.FieldInfo' descriptions. This is+-- part of the public extensibility surface, not an internal module.+module PgWire.Protocol.Oid+  ( Oid (..)+  , oidBool+  , oidBytea+  , oidInt8+  , oidInt2+  , oidInt4+  , oidText+  , oidJson+  , oidFloat4+  , oidFloat8+  , oidVarchar+  , oidDate+  , oidTime+  , oidTimetz+  , oidTimestamp+  , oidTimestamptz+  , oidNumeric+  , oidUuid+  , oidJsonb+  , oidInterval+  , oidInet+  , oidCidr+  , oidMacAddr+  , oidMacAddr8+    -- * Array OIDs+  , oidBoolArray+  , oidByteaArray+  , oidInt2Array+  , oidInt4Array+  , oidInt8Array+  , oidTextArray+  , oidVarcharArray+  , oidFloat4Array+  , oidFloat8Array+  , oidTimestampArray+  , oidTimestamptzArray+  , oidDateArray+  , oidTimeArray+  , oidUuidArray+  , oidJsonArray+  , oidJsonbArray+  , isArrayOid+  , arrayElementOid+  ) where++import Data.Word (Word32)++-- | A PostgreSQL type OID.+newtype Oid = Oid {unOid :: Word32}+  deriving stock (Show, Eq, Ord)++oidBool, oidBytea, oidInt8, oidInt2, oidInt4 :: Oid+oidBool = Oid 16+oidBytea = Oid 17+oidInt8 = Oid 20+oidInt2 = Oid 21+oidInt4 = Oid 23++oidText, oidJson, oidFloat4, oidFloat8 :: Oid+oidText = Oid 25+oidJson = Oid 114+oidFloat4 = Oid 700+oidFloat8 = Oid 701++oidVarchar, oidDate, oidTime, oidTimetz, oidTimestamp, oidTimestamptz :: Oid+oidVarchar = Oid 1043+oidDate = Oid 1082+oidTime = Oid 1083+oidTimetz = Oid 1266+oidTimestamp = Oid 1114+oidTimestamptz = Oid 1184++oidNumeric, oidUuid, oidJsonb, oidInterval, oidInet, oidCidr, oidMacAddr, oidMacAddr8 :: Oid+oidNumeric = Oid 1700+oidUuid = Oid 2950+oidJsonb = Oid 3802+oidInterval = Oid 1186+oidInet = Oid 869+oidCidr = Oid 650+oidMacAddr = Oid 829+oidMacAddr8 = Oid 774++-- Array OIDs ----------------------------------------------------------------++oidBoolArray, oidByteaArray, oidInt2Array, oidInt4Array, oidInt8Array :: Oid+oidBoolArray = Oid 1000+oidByteaArray = Oid 1001+oidInt2Array = Oid 1005+oidInt4Array = Oid 1007+oidInt8Array = Oid 1016++oidTextArray, oidVarcharArray :: Oid+oidTextArray = Oid 1009+oidVarcharArray = Oid 1015++oidFloat4Array, oidFloat8Array :: Oid+oidFloat4Array = Oid 1021+oidFloat8Array = Oid 1022++oidTimestampArray, oidTimestamptzArray, oidDateArray, oidTimeArray :: Oid+oidTimestampArray = Oid 1115+oidTimestamptzArray = Oid 1185+oidDateArray = Oid 1182+oidTimeArray = Oid 1183++oidUuidArray, oidJsonArray, oidJsonbArray :: Oid+oidUuidArray = Oid 2951+oidJsonArray = Oid 199+oidJsonbArray = Oid 3807++-- | Check whether an OID is a known array type.+isArrayOid :: Oid -> Bool+isArrayOid oid = case arrayElementOid oid of+  Just _ -> True+  Nothing -> False++-- | For an array OID, return the element OID.+arrayElementOid :: Oid -> Maybe Oid+arrayElementOid (Oid o) = case o of+  1000 -> Just oidBool+  1001 -> Just oidBytea+  1005 -> Just oidInt2+  1007 -> Just oidInt4+  1009 -> Just oidText+  1015 -> Just oidVarchar+  1016 -> Just oidInt8+  1021 -> Just oidFloat4+  1022 -> Just oidFloat8+  1115 -> Just oidTimestamp+  1182 -> Just oidDate+  1183 -> Just oidTime+  1185 -> Just oidTimestamptz+  2951 -> Just oidUuid+  199 -> Just oidJson+  3807 -> Just oidJsonb+  _ -> Nothing
+ src/PgWire/Protocol/Parsers.hs view
@@ -0,0 +1,355 @@+-- | Parsers for backend (server-to-client) protocol messages.+--+-- __Audience:__ exposed for downstream library authors that decode+-- raw protocol bytes themselves. Most application code consumes+-- decoded results via "PgWire.Connection".+module PgWire.Protocol.Parsers+  ( parseBackendMsg+  , parseCommandTag+  ) where++import Data.Bits (shiftL, (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Unsafe qualified as BU+import Data.Int (Int16, Int32, Int64)+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Vector.Mutable qualified as MV+import Data.Word (Word8, Word32)+import PgWire.Protocol.Backend++-- | Parse a complete backend message from tag + payload bytes.+-- The caller must have already read the tag byte and the 4-byte length,+-- and pass the tag and the remaining payload (length - 4 bytes).+parseBackendMsg :: Word8 -> ByteString -> Either String BackendMsg+parseBackendMsg tag payload = case tag of+  82 {- R -} -> parseAuth payload+  75 {- K -} -> parseBackendKeyData payload+  83 {- S -} -> parseParameterStatus payload+  90 {- Z -} -> parseReadyForQuery payload+  49 {- 1 -} -> Right ParseComplete+  50 {- 2 -} -> Right BindComplete+  51 {- 3 -} -> Right CloseComplete+  67 {- C -} -> parseCommandComplete payload+  68 {- D -} -> parseDataRow payload+  69 {- E -} -> ErrorResponse <$> parseErrorFields payload+  71 {- G -} -> parseCopyInResponse payload+  72 {- H -} -> parseCopyOutResponse payload+  73 {- I -} -> Right EmptyQueryResponse+  99 {- c -} -> Right CopyDoneMsg+  100 {- d -} -> Right (CopyDataMsg payload)+  78 {- N -} -> NoticeResponse . PgNotice <$> parseErrorFields payload+  84 {- T -} -> parseRowDescription payload+  110 {- n -} -> Right NoData+  116 {- t -} -> parseParameterDescription payload+  65 {- A -} -> parseNotification payload+  _ -> Left $ "Unknown backend message tag: " <> show tag++-- Authentication ----------------------------------------------------------++parseAuth :: ByteString -> Either String BackendMsg+parseAuth bs = do+  authType <- getInt32 bs 0+  case authType of+    0 -> Right (Authentication AuthOk)+    3 -> Right (Authentication AuthCleartextPassword)+    5 -> do+      salt <- getBytes bs 4 4+      Right (Authentication (AuthMD5Password salt))+    10 -> do+      -- SASL: list of mechanism names, each NUL-terminated, final NUL+      let mechData = BS.drop 4 bs+          mechs = parseCStringList mechData+      Right (Authentication (AuthSASL mechs))+    11 -> Right (Authentication (AuthSASLContinue (BS.drop 4 bs)))+    12 -> Right (Authentication (AuthSASLFinal (BS.drop 4 bs)))+    _ -> Left $ "Unknown auth type: " <> show authType++-- ParameterStatus ---------------------------------------------------------++parseParameterStatus :: ByteString -> Either String BackendMsg+parseParameterStatus bs =+  let (name, rest) = BS.break (== 0) bs+      value = BS.takeWhile (/= 0) (BS.drop 1 rest)+   in Right (ParameterStatus name value)++-- BackendKeyData ----------------------------------------------------------++parseBackendKeyData :: ByteString -> Either String BackendMsg+parseBackendKeyData bs = do+  pid <- getInt32 bs 0+  key <- getInt32 bs 4+  Right (BackendKeyData pid key)++-- ReadyForQuery -----------------------------------------------------------++parseReadyForQuery :: ByteString -> Either String BackendMsg+parseReadyForQuery bs+  | BS.length bs < 1 = Left "ReadyForQuery: empty payload"+  | otherwise = case BS.index bs 0 of+      73 {- I -} -> Right (ReadyForQuery TxIdle)+      84 {- T -} -> Right (ReadyForQuery TxInTransaction)+      69 {- E -} -> Right (ReadyForQuery TxFailed)+      c -> Left $ "ReadyForQuery: unknown status " <> show c++-- CommandComplete ---------------------------------------------------------++parseCommandComplete :: ByteString -> Either String BackendMsg+parseCommandComplete bs =+  let tag = BS.takeWhile (/= 0) bs+   in Right (CommandComplete (parseCommandTag tag))++parseCommandTag :: ByteString -> CommandTag+parseCommandTag bs+  | "SELECT " `BS.isPrefixOf` bs = SelectTag (readInt64 (BS.drop 7 bs))+  | "INSERT " `BS.isPrefixOf` bs =+      -- INSERT oid count — we want the count (after the second space)+      let afterInsert = BS.drop 7 bs+          afterOid = BS.drop 1 (BS.dropWhile (/= 32) afterInsert)+       in InsertTag (readInt64 afterOid)+  | "UPDATE " `BS.isPrefixOf` bs = UpdateTag (readInt64 (BS.drop 7 bs))+  | "DELETE " `BS.isPrefixOf` bs = DeleteTag (readInt64 (BS.drop 7 bs))+  | otherwise = OtherTag bs++readInt64 :: ByteString -> Int64+readInt64 bs =+  BS.foldl' (\acc w -> acc * 10 + fromIntegral (w - 48)) 0 (BS.takeWhile (\w -> w >= 48 && w <= 57) bs)++-- RowDescription ----------------------------------------------------------++parseRowDescription :: ByteString -> Either String BackendMsg+parseRowDescription bs = do+  nFields <- getInt16 bs 0+  (fields, _) <- parseFields (fromIntegral nFields) (BS.drop 2 bs)+  Right (RowDescription (V.fromList fields))++parseFields :: Int -> ByteString -> Either String ([FieldInfo], ByteString)+parseFields 0 rest = Right ([], rest)+parseFields n bs = do+  let (name, afterName) = BS.break (== 0) bs+      rest = BS.drop 1 afterName -- skip NUL+  tableOid <- getWord32 rest 0+  colNum <- getInt16 rest 4+  typeOid <- getWord32 rest 6+  typeSize <- getInt16 rest 10+  typeMod <- getInt32 rest 12+  fmtCode <- getInt16 rest 16+  let field =+        FieldInfo+          { fiName = name+          , fiTableOid = tableOid+          , fiColumnNum = colNum+          , fiTypeOid = typeOid+          , fiTypeSize = typeSize+          , fiTypeMod = typeMod+          , fiFormatCode = fmtCode+          }+  (fields, rest') <- parseFields (n - 1) (BS.drop 18 rest)+  Right (field : fields, rest')++-- DataRow -----------------------------------------------------------------++-- | Parse a DataRow message directly into a 'Vector', skipping the+-- intermediate list. Column value slices share the underlying buffer+-- via 'BS.take'/'BS.drop' (zero-copy).+parseDataRow :: ByteString -> Either String BackendMsg+parseDataRow bs = do+  nCols <- getInt16 bs 0+  let !n = fromIntegral nCols+  vals <- parseColsDirect n bs 2+  Right (DataRow vals)++parseColsDirect :: Int -> ByteString -> Int -> Either String (Vector (Maybe ByteString))+parseColsDirect n bs startOff+  | BS.length bs < startOff = Left "DataRow: buffer too short for column data"+  | otherwise =+      let !bsLen = BS.length bs+          -- First pass: validate all lengths fit within the buffer.+          validate !i !off+            | i >= n = Right off+            | off + 4 > bsLen = Left "DataRow: truncated column length"+            | otherwise =+                let !len = getInt32Unsafe bs off+                 in if len == -1+                      then validate (i + 1) (off + 4)+                      else+                        let !end = off + 4 + fromIntegral len+                         in if end > bsLen+                              then Left "DataRow: truncated column data"+                              else validate (i + 1) end+       in case validate 0 startOff of+            Left err -> Left err+            Right _ -> Right $ V.create $ do+              mv <- MV.new n+              let go !i !off+                    | i >= n = pure ()+                    | otherwise =+                        let !len = getInt32Unsafe bs off+                         in if len == -1+                              then do+                                MV.write mv i Nothing+                                go (i + 1) (off + 4)+                              else do+                                let !dataLen = fromIntegral len+                                    !val = BU.unsafeTake dataLen (BU.unsafeDrop (off + 4) bs)+                                MV.write mv i (Just val)+                                go (i + 1) (off + 4 + dataLen)+              go 0 startOff+              pure mv+  where+    -- Unchecked Int32 decode — bounds validated by the validation pass.+    getInt32Unsafe :: ByteString -> Int -> Int32+    getInt32Unsafe b off =+      let !b0 = fromIntegral (BU.unsafeIndex b off) :: Int32+          !b1 = fromIntegral (BU.unsafeIndex b (off + 1)) :: Int32+          !b2 = fromIntegral (BU.unsafeIndex b (off + 2)) :: Int32+          !b3 = fromIntegral (BU.unsafeIndex b (off + 3)) :: Int32+       in (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3+    {-# INLINE getInt32Unsafe #-}++-- ParameterDescription ----------------------------------------------------++parseParameterDescription :: ByteString -> Either String BackendMsg+parseParameterDescription bs = do+  nParams <- getInt16 bs 0+  oids <- parseOids (fromIntegral nParams) (BS.drop 2 bs)+  Right (ParameterDescription (V.fromList oids))++parseOids :: Int -> ByteString -> Either String [Word32]+parseOids 0 _ = Right []+parseOids n bs = do+  oid <- getWord32 bs 0+  oids <- parseOids (n - 1) (BS.drop 4 bs)+  Right (oid : oids)++-- ErrorResponse / NoticeResponse ------------------------------------------++parseErrorFields :: ByteString -> Either String PgError+parseErrorFields = go emptyPgError+  where+    go !err bs+      | BS.null bs = Right err+      | BS.index bs 0 == 0 = Right err+      | otherwise =+          let fieldType = BS.index bs 0+              (value, rest) = BS.break (== 0) (BS.drop 1 bs)+              err' = case fieldType of+                83 {- S -} -> err {pgSeverity = value}+                67 {- C -} -> err {pgCode = value}+                77 {- M -} -> err {pgMessage = value}+                68 {- D -} -> err {pgDetail = Just value}+                72 {- H -} -> err {pgHint = Just value}+                80 {- P -} -> err {pgPosition = Just (fromIntegral (readInt64 value))}+                112 {- p -} -> err {pgInternalPosition = Just (fromIntegral (readInt64 value))}+                113 {- q -} -> err {pgInternalQuery = Just value}+                87 {- W -} -> err {pgWhere = Just value}+                115 {- s -} -> err {pgSchema = Just value}+                116 {- t -} -> err {pgTable = Just value}+                99 {- c -} -> err {pgColumn = Just value}+                100 {- d -} -> err {pgDataType = Just value}+                110 {- n -} -> err {pgConstraint = Just value}+                70 {- F -} -> err {pgFile = Just value}+                76 {- L -} -> err {pgLine = Just (fromIntegral (readInt64 value))}+                82 {- R -} -> err {pgRoutine = Just value}+                _ -> err+           in go err' (BS.drop 1 rest) -- skip NUL after value++    emptyPgError =+      PgError+        { pgSeverity = ""+        , pgCode = ""+        , pgMessage = ""+        , pgDetail = Nothing+        , pgHint = Nothing+        , pgPosition = Nothing+        , pgInternalPosition = Nothing+        , pgInternalQuery = Nothing+        , pgWhere = Nothing+        , pgSchema = Nothing+        , pgTable = Nothing+        , pgColumn = Nothing+        , pgDataType = Nothing+        , pgConstraint = Nothing+        , pgFile = Nothing+        , pgLine = Nothing+        , pgRoutine = Nothing+        }++-- CopyInResponse / CopyOutResponse ----------------------------------------++parseCopyInResponse :: ByteString -> Either String BackendMsg+parseCopyInResponse bs = do+  fmt <- getInt8 bs 0+  nCols <- getInt16 bs 1+  fmts <- mapM (\i -> getInt16 bs (3 + i * 2)) [0 .. fromIntegral nCols - 1]+  Right (CopyInResponse (fromIntegral fmt) (V.fromList fmts))++parseCopyOutResponse :: ByteString -> Either String BackendMsg+parseCopyOutResponse bs = do+  fmt <- getInt8 bs 0+  nCols <- getInt16 bs 1+  fmts <- mapM (\i -> getInt16 bs (3 + i * 2)) [0 .. fromIntegral nCols - 1]+  Right (CopyOutResponse (fromIntegral fmt) (V.fromList fmts))++getInt8 :: ByteString -> Int -> Either String Word8+getInt8 bs off+  | BS.length bs < off + 1 = Left "getInt8: buffer too short"+  | otherwise = Right (BS.index bs off)++-- Notification ------------------------------------------------------------++parseNotification :: ByteString -> Either String BackendMsg+parseNotification bs = do+  pid <- getInt32 bs 0+  let (channel, rest1) = BS.break (== 0) (BS.drop 4 bs)+      payload = BS.takeWhile (/= 0) (BS.drop 1 rest1)+  Right (NotificationResponse pid channel payload)++-- Helpers -----------------------------------------------------------------++parseCStringList :: ByteString -> [ByteString]+parseCStringList bs+  | BS.null bs = []+  | BS.index bs 0 == 0 = []+  | otherwise =+      let (s, rest) = BS.break (== 0) bs+       in s : parseCStringList (BS.drop 1 rest)++getInt16 :: ByteString -> Int -> Either String Int16+getInt16 bs off+  | BS.length bs < off + 2 = Left "getInt16: buffer too short"+  | otherwise =+      -- After bounds check, use unsafeIndex for the hot path+      let hi = fromIntegral (BU.unsafeIndex bs off) :: Int16+          lo = fromIntegral (BU.unsafeIndex bs (off + 1)) :: Int16+       in Right (hi `shiftL` 8 .|. lo)+{-# INLINE getInt16 #-}++getInt32 :: ByteString -> Int -> Either String Int32+getInt32 bs off+  | BS.length bs < off + 4 = Left "getInt32: buffer too short"+  | otherwise =+      let b0 = fromIntegral (BU.unsafeIndex bs off) :: Int32+          b1 = fromIntegral (BU.unsafeIndex bs (off + 1)) :: Int32+          b2 = fromIntegral (BU.unsafeIndex bs (off + 2)) :: Int32+          b3 = fromIntegral (BU.unsafeIndex bs (off + 3)) :: Int32+       in Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)+{-# INLINE getInt32 #-}++getWord32 :: ByteString -> Int -> Either String Word32+getWord32 bs off+  | BS.length bs < off + 4 = Left "getWord32: buffer too short"+  | otherwise =+      let b0 = fromIntegral (BU.unsafeIndex bs off) :: Word32+          b1 = fromIntegral (BU.unsafeIndex bs (off + 1)) :: Word32+          b2 = fromIntegral (BU.unsafeIndex bs (off + 2)) :: Word32+          b3 = fromIntegral (BU.unsafeIndex bs (off + 3)) :: Word32+       in Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)+{-# INLINE getWord32 #-}++getBytes :: ByteString -> Int -> Int -> Either String ByteString+getBytes bs off len+  | BS.length bs < off + len = Left "getBytes: buffer too short"+  | otherwise = Right (BS.take len (BS.drop off bs))+{-# INLINE getBytes #-}
+ src/PgWire/TypeCache.hs view
@@ -0,0 +1,60 @@+-- | Pool-level cache for resolved PostgreSQL type metadata.+--+-- When multiple connections encounter the same custom OID (e.g., a PG+-- enum), this cache avoids redundant @pg_type@ round-trips. The cache+-- is shared across all connections in a pool via a 'TVar'.+--+-- The cache grows without bound but is bounded in practice by the number+-- of distinct custom types in the database (typically tens, not thousands).+-- Applications that dynamically create and drop types at high volume may+-- want to monitor 'cachedTypeCount' for unexpected growth.+module PgWire.TypeCache+  ( TypeCache+  , TypeInfo (..)+  , newTypeCache+  , lookupTypeInfo+  , cacheTypeInfo+  , cachedTypeCount+  ) where++import Control.Concurrent.STM (TVar, atomically, modifyTVar', newTVarIO, readTVarIO)+import Data.ByteString (ByteString)+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Word (Word32)++-- | Metadata about a PG type, cached at the pool level.+data TypeInfo = TypeInfo+  { tiName :: !ByteString+  -- ^ Type name from @pg_type.typname@.+  , tiCategory :: !Char+  -- ^ Type category: @\'E\'@ enum, @\'C\'@ composite, @\'R\'@ range,+  --   @\'D\'@ domain, @\'B\'@ base, @\'P\'@ pseudo, etc.+  , tiArrayOid :: !Word32+  -- ^ OID of the array form of this type (0 if none).+  , tiBaseOid :: !Word32+  -- ^ For domains: the underlying type OID. 0 otherwise.+  }+  deriving stock (Show, Eq)++-- | A thread-safe cache mapping OIDs to their resolved type info.+newtype TypeCache = TypeCache (TVar (Map Word32 TypeInfo))++-- | Create an empty type cache.+newTypeCache :: IO TypeCache+newTypeCache = TypeCache <$> newTVarIO Map.empty++-- | Look up cached type info for an OID.+lookupTypeInfo :: TypeCache -> Word32 -> IO (Maybe TypeInfo)+lookupTypeInfo (TypeCache tv) oid = do+  m <- readTVarIO tv+  pure (Map.lookup oid m)++-- | Cache type info for an OID.+cacheTypeInfo :: TypeCache -> Word32 -> TypeInfo -> IO ()+cacheTypeInfo (TypeCache tv) oid info =+  atomically $ modifyTVar' tv (Map.insert oid info)++-- | Return the number of cached entries (for diagnostics).+cachedTypeCount :: TypeCache -> IO Int+cachedTypeCount (TypeCache tv) = Map.size <$> readTVarIO tv
+ src/PgWire/Wire.hs view
@@ -0,0 +1,390 @@+-- | Low-level PostgreSQL wire protocol I\/O.+--+-- 'WireConn' abstracts over a TCP or TLS connection, providing+-- send\/recv operations for protocol messages. This module handles+-- message framing (tag + length + payload), buffering, and+-- TLS upgrades via the SSLRequest subprotocol.+--+-- Most users should use 'PgWire.Connection' instead of this module.+module PgWire.Wire+  ( -- * Wire connection+    WireConn (..)+    -- * Connecting+  , connectTcp+  , connectTcpTimeout+    -- * TLS+  , TlsConfig (..)+  , upgradeTls+    -- * Tracing+  , TraceDirection (..)+    -- * Sending messages+  , sendFrontendMsg+  , sendFrontendMsgs+  , sendRawBytes+    -- * Receiving messages+  , recvBackendMsg+  ) where++import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Char8 qualified as BS8+import Data.ByteString.Unsafe qualified as BU+import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.IORef+import Data.Maybe (fromMaybe)+import Data.Time (NominalDiffTime)+import PgWire.Error (PgWireError (..), throwPgWire)+import System.Timeout (timeout)+import PgWire.Protocol.Backend (BackendMsg)+import PgWire.Protocol.Builders (buildFrontendMsg, buildFrontendMsgsConcat)+import PgWire.Protocol.Frontend (FrontendMsg)+import PgWire.Protocol.Parsers (parseBackendMsg)+import Network.Socket (Socket)+import Network.Socket qualified as NS+import Network.Socket.ByteString qualified as NSB+import Network.Socket (setSocketOption, SocketOption(..))+import Network.TLS qualified as TLS+import Network.TLS (ClientParams (..), ClientHooks (..), Supported (..), Shared (..), Credentials (..))+import Data.X509.CertificateStore (readCertificateStore)+import System.X509 (getSystemCertificateStore)++-- | Abstraction over a Postgres wire connection (plain TCP or TLS).+-- | Direction for protocol trace callback.+data TraceDirection = TraceSend | TraceRecv+  deriving stock (Show, Eq)++data WireConn = WireConn+  { wcSend :: ByteString -> IO ()+  , wcSendMany :: [ByteString] -> IO ()+  -- ^ Vectored I/O: send multiple chunks in one syscall (writev).+  , wcRecv :: Int -> IO ByteString+  , wcClose :: IO ()+  , wcBuffer :: IORef ByteString+  , wcTrace :: IORef (Maybe (TraceDirection -> ByteString -> IO ()))+  }++-- | Connect via TCP to the given host and port (no timeout).+connectTcp :: NS.HostName -> NS.PortNumber -> IO WireConn+connectTcp = connectTcpTimeout 0++-- | Connect via TCP with a timeout in seconds (0 = no timeout).+connectTcpTimeout :: NominalDiffTime -> NS.HostName -> NS.PortNumber -> IO WireConn+connectTcpTimeout timeoutSecs host port = do+  let hints = NS.defaultHints {NS.addrSocketType = NS.Stream}+  addrs <- NS.getAddrInfo (Just hints) (Just host) (Just (show port))+  case addrs of+    [] -> throwPgWire (ConnectionError "No address found")+    (addr : _) -> do+      sock <- NS.socket (NS.addrFamily addr) NS.Stream NS.defaultProtocol+      let doConnect = NS.connect sock (NS.addrAddress addr)+      if timeoutSecs > 0+        then do+          let micros = round (timeoutSecs * 1000000) :: Int+          result <- timeout micros doConnect+          case result of+            Nothing -> do+              NS.close sock+              throwPgWire (ConnectionError "Connect timed out")+            Just () -> pure ()+        else doConnect+      -- Disable Nagle's algorithm for lower latency on small messages+      setSocketOption sock NoDelay 1+      mkWireConn sock++mkWireConn :: Socket -> IO WireConn+mkWireConn sock = do+  buf <- newIORef BS.empty+  traceRef <- newIORef Nothing+  pure+    WireConn+      { wcSend = sendAll sock+      , wcSendMany = NSB.sendMany sock+      , wcRecv = recvExact sock buf+      , wcClose = NS.close sock+      , wcBuffer = buf+      , wcTrace = traceRef+      }++-- | TLS configuration for upgradeTls.+data TlsConfig = TlsConfig+  { tlsHostname :: NS.HostName+  , tlsVerify :: Bool+  -- ^ Whether to verify the server certificate (verify-ca / verify-full)+  , tlsVerifyHostname :: Bool+  -- ^ Whether to verify hostname matches cert (verify-full only)+  , tlsClientCert :: Maybe FilePath+  -- ^ Path to client certificate file+  , tlsClientKey :: Maybe FilePath+  -- ^ Path to client private key file+  , tlsCaCert :: Maybe FilePath+  -- ^ Path to CA certificate file (Nothing = use system store)+  }++-- | Upgrade a TCP WireConn to TLS. Sends the SSLRequest message,+-- checks the server response, and performs the TLS handshake.+-- Returns a new WireConn that sends/receives over TLS.+upgradeTls :: WireConn -> TlsConfig -> IO WireConn+upgradeTls wc tlsCfg = do+  let sslRequest = LBS.toStrict . B.toLazyByteString $+        B.int32BE 8 <> B.int32BE 80877103+  wcSend wc sslRequest++  resp <- wcRecv wc 1+  case BS.index resp 0 of+    83 {- S -} -> do+      -- Load CA certificate store+      systemStore <- getSystemCertificateStore+      caStore <- case tlsCaCert tlsCfg of+        Nothing -> pure systemStore+        Just "system" -> pure systemStore+        Just path -> do+          mStore <- readCertificateStore path+          pure (fromMaybe systemStore mStore)++      -- Load client certificate + key if provided+      clientCred <- case (tlsClientCert tlsCfg, tlsClientKey tlsCfg) of+        (Just certPath, Just keyPath) -> do+          result <- TLS.credentialLoadX509 certPath keyPath+          case result of+            Right cred -> pure (Credentials [cred])+            Left _err -> do+              -- Non-fatal: warn but continue without client cert+              pure (Credentials [])+        _ -> pure (Credentials [])++      let hostname = tlsHostname tlsCfg+          baseParams = TLS.defaultParamsClient hostname ""+          -- Configure certificate validation based on mode+          hooks = (clientHooks baseParams)+            { onServerCertificate =+                if tlsVerify tlsCfg+                  then onServerCertificate (clientHooks baseParams)  -- default validation+                  else \_ _ _ _ -> pure []  -- skip validation (sslmode=require)+            }+          clientParams = baseParams+            { clientSupported = (clientSupported baseParams)+                { supportedVersions = [TLS.TLS13, TLS.TLS12]+                }+            , clientShared = (clientShared baseParams)+                { sharedCAStore = caStore+                , sharedCredentials = clientCred+                }+            , clientHooks = hooks+            }+      ctx <- TLS.contextNew (wireBackend wc) clientParams+      TLS.handshake ctx+      mkTlsWireConn ctx (wcBuffer wc)++    78 {- N -} ->+      throwPgWire (ConnectionError "Server does not support SSL")++    other ->+      throwPgWire (ConnectionError ("Unexpected SSL response: " <> BS8.pack (show other)))++-- | Wrap a WireConn as a TLS backend (for the handshake, which needs+-- raw socket I/O).+wireBackend :: WireConn -> TLS.Backend+wireBackend wc = TLS.Backend+  { TLS.backendFlush = pure ()+  , TLS.backendClose = wcClose wc+  , TLS.backendSend = wcSend wc+  , TLS.backendRecv = \n -> do+      -- TLS library expects exactly n bytes+      wcRecv wc n+  }++mkTlsWireConn :: TLS.Context -> IORef ByteString -> IO WireConn+mkTlsWireConn ctx bufRef = do+  writeIORef bufRef BS.empty+  tlsBuf <- newIORef BS.empty+  traceRef <- newIORef Nothing+  let tlsSend = TLS.sendData ctx . LBS.fromStrict+  pure+    WireConn+      { wcSend = tlsSend+      , wcSendMany = \chunks -> tlsSend (BS.concat chunks)+      , wcRecv = tlsRecvExact ctx tlsBuf+      , wcClose = TLS.bye ctx >> TLS.contextClose ctx+      , wcBuffer = tlsBuf+      , wcTrace = traceRef+      }++-- | Receive exactly n bytes from a TLS context, buffering leftovers.+-- Same optimization as 'recvExact': fast path for buffered data, Builder+-- accumulation for multi-chunk reads.+tlsRecvExact :: TLS.Context -> IORef ByteString -> Int -> IO ByteString+tlsRecvExact ctx bufRef n = do+  buf <- readIORef bufRef+  if BS.length buf >= n+    then do+      let (!taken, !rest) = BS.splitAt n buf+      writeIORef bufRef rest+      pure taken+    else if BS.null buf+      then do+        chunk <- TLS.recvData ctx+        if BS.null chunk+          then throwPgWire (ConnectionError "TLS connection closed")+          else if BS.length chunk >= n+            then do+              let (!taken, !rest) = BS.splitAt n chunk+              writeIORef bufRef rest+              pure taken+            else tlsAccumulate ctx bufRef chunk (n - BS.length chunk) (B.byteString chunk)+      else do+        let !remaining = n - BS.length buf+            !builder0 = B.byteString buf+        chunk <- TLS.recvData ctx+        if BS.null chunk+          then throwPgWire (ConnectionError "TLS connection closed")+          else if BS.length chunk >= remaining+            then do+              let (!taken, !rest) = BS.splitAt remaining chunk+              writeIORef bufRef rest+              pure (LBS.toStrict (B.toLazyByteString (builder0 <> B.byteString taken)))+            else tlsAccumulate ctx bufRef chunk (remaining - BS.length chunk) (builder0 <> B.byteString chunk)++tlsAccumulate :: TLS.Context -> IORef ByteString -> ByteString -> Int -> B.Builder -> IO ByteString+tlsAccumulate ctx bufRef chunk remaining !builder+  | BS.length chunk >= remaining = do+      let (!taken, !rest) = BS.splitAt remaining chunk+      writeIORef bufRef rest+      pure (LBS.toStrict (B.toLazyByteString (builder <> B.byteString taken)))+  | otherwise = do+      let !remaining' = remaining - BS.length chunk+          !builder' = builder <> B.byteString chunk+      next <- TLS.recvData ctx+      if BS.null next+        then throwPgWire (ConnectionError "TLS connection closed")+        else tlsAccumulate ctx bufRef next remaining' builder'++-- | Send a frontend message over the wire.+sendFrontendMsg :: WireConn -> FrontendMsg -> IO ()+sendFrontendMsg wc msg = do+  let bytes = buildFrontendMsg msg+  traceIfEnabled wc TraceSend bytes+  wcSend wc bytes+{-# INLINE sendFrontendMsg #-}++-- | Send multiple frontend messages in a single allocation and syscall.+-- All messages are fused into one 'Builder', materialized once with+-- 'buildFrontendMsgsConcat', and sent with a single @send()@.+sendFrontendMsgs :: WireConn -> [FrontendMsg] -> IO ()+sendFrontendMsgs wc msgs = do+  let !bytes = buildFrontendMsgsConcat msgs+  traceIfEnabled wc TraceSend bytes+  wcSend wc bytes+{-# INLINE sendFrontendMsgs #-}++-- | Send raw bytes (for startup message which has a different format).+sendRawBytes :: WireConn -> ByteString -> IO ()+sendRawBytes wc bs = do+  traceIfEnabled wc TraceSend bs+  wcSend wc bs++-- | Call the trace handler if one is set.+traceIfEnabled :: WireConn -> TraceDirection -> ByteString -> IO ()+traceIfEnabled wc dir bytes = do+  mHandler <- readIORef (wcTrace wc)+  case mHandler of+    Nothing -> pure ()+    Just handler -> handler dir bytes+{-# INLINE traceIfEnabled #-}++-- | Receive and parse a single backend message.+recvBackendMsg :: WireConn -> IO BackendMsg+recvBackendMsg wc = do+  -- Read tag (1 byte) + length (4 bytes) together in a single recv+  header <- wcRecv wc 5+  let !tag = BU.unsafeIndex header 0+      !len = decodeInt32At header 1+      !payloadLen = len - 4+  -- Read payload+  payload <-+    if payloadLen > 0+      then wcRecv wc (fromIntegral payloadLen)+      else if payloadLen < 0+        then throwPgWire (ProtocolError ("Invalid message length: " <> BS8.pack (show len)))+        else pure BS.empty+  -- Trace the raw bytes (header + payload)+  traceIfEnabled wc TraceRecv (header <> payload)+  case parseBackendMsg tag payload of+    Left err -> throwPgWire (ProtocolError (BS8.pack err))+    Right msg -> pure msg+  where+    decodeInt32At :: ByteString -> Int -> Int+    decodeInt32At bs off =+      let !b0 = fromIntegral (BU.unsafeIndex bs off)+          !b1 = fromIntegral (BU.unsafeIndex bs (off + 1))+          !b2 = fromIntegral (BU.unsafeIndex bs (off + 2))+          !b3 = fromIntegral (BU.unsafeIndex bs (off + 3))+       in b0 * 16777216 + b1 * 65536 + b2 * 256 + b3++-- Socket helpers ----------------------------------------------------------++sendAll :: Socket -> ByteString -> IO ()+sendAll sock bs+  | BS.null bs = pure ()+  | otherwise = do+      sent <- NSB.send sock bs+      sendAll sock (BS.drop sent bs)++-- | Socket recv buffer size. Larger values reduce syscalls for big result sets.+recvChunkSize :: Int+recvChunkSize = 32768+{-# INLINE recvChunkSize #-}++-- | Receive exactly @n@ bytes, using the buffer for leftovers.+-- Fast path: if the buffer already has enough data, just splitAt (zero allocation+-- beyond the two slices). Slow path: accumulate chunks via Builder to avoid+-- the O(n) reverse + concat of a list accumulator.+recvExact :: Socket -> IORef ByteString -> Int -> IO ByteString+recvExact sock bufRef n = do+  buf <- readIORef bufRef+  if BS.length buf >= n+    then do+      -- Fast path: buffer has enough data, no syscall needed.+      let (!taken, !rest) = BS.splitAt n buf+      writeIORef bufRef rest+      pure taken+    else if BS.null buf+      then do+        -- Buffer empty: read fresh chunk from socket.+        chunk <- NSB.recv sock (max recvChunkSize n)+        if BS.null chunk+          then throwPgWire (ConnectionError "Connection closed by server")+          else if BS.length chunk >= n+            then do+              let (!taken, !rest) = BS.splitAt n chunk+              writeIORef bufRef rest+              pure taken+            else accumulate chunk (n - BS.length chunk) (B.byteString chunk)+      else do+        -- Buffer has partial data: use it and read more.+        let !remaining = n - BS.length buf+            !builder0 = B.byteString buf+        chunk <- NSB.recv sock (max recvChunkSize remaining)+        if BS.null chunk+          then throwPgWire (ConnectionError "Connection closed by server")+          else if BS.length chunk >= remaining+            then do+              let (!taken, !rest) = BS.splitAt remaining chunk+              writeIORef bufRef rest+              pure (LBS.toStrict (B.toLazyByteString (builder0 <> B.byteString taken)))+            else accumulate chunk (remaining - BS.length chunk) (builder0 <> B.byteString chunk)+  where+    -- Slow path: need multiple recv calls. Uses Builder for O(1) append.+    accumulate :: ByteString -> Int -> B.Builder -> IO ByteString+    accumulate chunk remaining !builder+      | BS.length chunk >= remaining = do+          let (!taken, !rest) = BS.splitAt remaining chunk+          writeIORef bufRef rest+          pure (LBS.toStrict (B.toLazyByteString (builder <> B.byteString taken)))+      | otherwise = do+          let !remaining' = remaining - BS.length chunk+              !builder' = builder <> B.byteString chunk+          next <- NSB.recv sock (max recvChunkSize remaining')+          if BS.null next+            then throwPgWire (ConnectionError "Connection closed by server")+            else accumulate next remaining' builder'
+ test/Main.hs view
@@ -0,0 +1,38 @@+module Main where++import PgWire.Auth.MD5Spec qualified as MD5Spec+import PgWire.Auth.ScramFieldsSpec qualified as ScramFieldsSpec+import PgWire.CancelSpec qualified as CancelSpec+import PgWire.Connection.ConfigSpec qualified as ConfigSpec+import PgWire.ErrorSpec qualified as ErrorSpec+import PgWire.MockServerSpec qualified as MockServerSpec+import PgWire.Pool.ConfigSpec qualified as PoolConfigSpec+import PgWire.Pool.NoThunksSpec qualified as NoThunksSpec+import PgWire.Pool.PropertySpec qualified as PoolPropertySpec+import PgWire.Pool.StateMachineSpec qualified as PoolStateMachineSpec+import PgWire.Protocol.BuildersSpec qualified as BuildersSpec+import PgWire.Protocol.OidSpec qualified as OidSpec+import PgWire.Connection.EscapingSpec qualified as EscapingSpec+import PgWire.Connection.FeaturesSpec qualified as FeaturesSpec+import PgWire.Protocol.FuzzSpec qualified as FuzzSpec+import PgWire.Protocol.ParsersSpec qualified as ParsersSpec+import Test.Hspec++main :: IO ()+main = hspec $ do+  describe "PgWire.Protocol.Builders" BuildersSpec.spec+  describe "PgWire.Protocol.Parsers" ParsersSpec.spec+  describe "PgWire.Protocol.Fuzz" FuzzSpec.spec+  describe "PgWire.Protocol.Oid" OidSpec.spec+  describe "PgWire.Connection.Config" ConfigSpec.spec+  describe "PgWire.Pool.Config" PoolConfigSpec.spec+  describe "PgWire.Error" ErrorSpec.spec+  describe "PgWire.Auth.MD5" MD5Spec.spec+  describe "PgWire.Auth.ScramFields" ScramFieldsSpec.spec+  describe "PgWire.Cancel" CancelSpec.spec+  describe "PgWire.Connection.Features" FeaturesSpec.spec+  describe "PgWire.Connection.Escaping" EscapingSpec.spec+  describe "PgWire.MockServer" MockServerSpec.spec+  describe "PgWire.Pool.NoThunks" NoThunksSpec.spec+  describe "PgWire.Pool.Property" PoolPropertySpec.spec+  describe "PgWire.Pool.StateMachine" PoolStateMachineSpec.spec
+ test/PgWire/Auth/MD5Spec.hs view
@@ -0,0 +1,44 @@+module PgWire.Auth.MD5Spec (spec) where++import Crypto.Hash (MD5 (..), hashWith)+import Data.ByteArray qualified as BA+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Char8 qualified as BS8+import Test.Hspec++spec :: Spec+spec = do+  describe "MD5 password hash" $ do+    it "produces correct md5(md5(password+user)+salt) format" $ do+      -- PG MD5 auth: md5(md5(password + username) + salt)+      -- where md5() returns lowercase hex string+      let user = "testuser" :: ByteString+          password = "testpass" :: ByteString+          salt = BS.pack [0x01, 0x02, 0x03, 0x04]+          inner = md5Hex (password <> user)+          outer = "md5" <> md5Hex (inner <> salt)+      -- The hash should be deterministic+      BS.length outer `shouldBe` 35 -- "md5" + 32 hex chars+      BS8.isPrefixOf "md5" outer `shouldBe` True++    it "different passwords produce different hashes" $ do+      let user = "user" :: ByteString+          salt = BS.pack [0, 0, 0, 0]+          h1 = md5Hex (md5Hex ("pass1" <> user) <> salt)+          h2 = md5Hex (md5Hex ("pass2" <> user) <> salt)+      h1 `shouldNotBe` h2++md5Hex :: ByteString -> ByteString+md5Hex bs = BS8.pack $ concatMap (pad2 . showHex') (BS.unpack (BA.convert (hashWith MD5 bs)))+  where+    showHex' n = showHexByte n ""+    pad2 s = replicate (2 - length s) '0' ++ s+    showHexByte n = showHex n+    showHex 0 = ('0' :)+    showHex n+      | n < 16 = (hexDigit (fromIntegral n) :)+      | otherwise = showHex (n `div` 16) . (hexDigit (fromIntegral (n `mod` 16)) :)+    hexDigit i+      | i < 10 = toEnum (fromEnum '0' + i)+      | otherwise = toEnum (fromEnum 'a' + i - 10)
+ test/PgWire/Auth/ScramFieldsSpec.hs view
@@ -0,0 +1,44 @@+module PgWire.Auth.ScramFieldsSpec (spec) where++import Data.ByteString qualified as BS+import Test.Hspec++-- Test the SCRAM field parser (internal, but we replicate the logic here)+spec :: Spec+spec = do+  describe "SCRAM field parsing" $ do+    it "parses r=nonce,s=salt,i=4096" $ do+      let input = "r=clientnonce+servernonce,s=c2FsdA==,i=4096"+          fields = parseScramFields input+      lookup "r" fields `shouldBe` Just "clientnonce+servernonce"+      lookup "s" fields `shouldBe` Just "c2FsdA=="+      lookup "i" fields `shouldBe` Just "4096"++    it "parses v=signature" $ do+      let input = "v=dGVzdHNpZw=="+          fields = parseScramFields input+      lookup "v" fields `shouldBe` Just "dGVzdHNpZw=="++    it "handles empty input" $ do+      let fields = parseScramFields ""+      fields `shouldBe` []++    it "handles single field" $ do+      let fields = parseScramFields "r=nonce"+      lookup "r" fields `shouldBe` Just "nonce"++    it "handles values containing equals signs" $ do+      -- base64 can contain '=' padding+      let fields = parseScramFields "s=abc=,i=100"+      -- The key is "s", value starts after first '='+      lookup "s" fields `shouldBe` Just "abc="++-- Replicated from ScramSHA256 module for testing+parseScramFields :: BS.ByteString -> [(BS.ByteString, BS.ByteString)]+parseScramFields bs =+  [ (key, val)+  | field <- BS.split 44 bs -- ','+  , let (key, rest) = BS.break (== 61) field -- '='+  , not (BS.null rest)+  , let val = BS.drop 1 rest+  ]
+ test/PgWire/CancelSpec.hs view
@@ -0,0 +1,52 @@+module PgWire.CancelSpec (spec) where++import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.ByteString qualified as BS+import Data.Int (Int32)+import Test.Hspec++spec :: Spec+spec = do+  describe "CancelRequest message format" $ do+    it "is exactly 16 bytes" $ do+      let msg = buildCancelRequest 12345 67890+      BS.length msg `shouldBe` 16++    it "starts with length 16" $ do+      let msg = buildCancelRequest 1 2+      -- First 4 bytes: length = 16 = 0x00000010+      BS.take 4 msg `shouldBe` BS.pack [0, 0, 0, 16]++    it "contains cancel code 80877102" $ do+      let msg = buildCancelRequest 1 2+      -- Bytes 4-7: cancel code = 80877102 = 0x04D2162E+      let code = fromIntegral (BS.index msg 4) * 256 * 256 * 256+               + fromIntegral (BS.index msg 5) * 256 * 256+               + fromIntegral (BS.index msg 6) * 256+               + fromIntegral (BS.index msg 7) :: Int+      code `shouldBe` 80877102++    it "contains PID and key" $ do+      let msg = buildCancelRequest 42 99+      -- Bytes 8-11: PID+      let pid = fromIntegral (BS.index msg 8) * 256 * 256 * 256+              + fromIntegral (BS.index msg 9) * 256 * 256+              + fromIntegral (BS.index msg 10) * 256+              + fromIntegral (BS.index msg 11) :: Int+      pid `shouldBe` 42+      -- Bytes 12-15: key+      let key = fromIntegral (BS.index msg 12) * 256 * 256 * 256+              + fromIntegral (BS.index msg 13) * 256 * 256+              + fromIntegral (BS.index msg 14) * 256+              + fromIntegral (BS.index msg 15) :: Int+      key `shouldBe` 99++-- Replicate the cancel message builder for testing+buildCancelRequest :: Int32 -> Int32 -> BS.ByteString+buildCancelRequest pid key =+  LBS.toStrict . B.toLazyByteString $+    B.int32BE 16+      <> B.int32BE 80877102+      <> B.int32BE pid+      <> B.int32BE key
+ test/PgWire/Connection/ConfigSpec.hs view
@@ -0,0 +1,71 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+module PgWire.Connection.ConfigSpec (spec) where++import PgWire.Connection.Config+import Test.Hspec++spec :: Spec+spec = do+  describe "parseConnString (URI format)" $ do+    it "parses basic postgres:// URL" $ do+      let Right cfg = parseConnString "postgres://user:pass@localhost:5432/mydb"+      ccHost cfg `shouldBe` "localhost"+      ccPort cfg `shouldBe` 5432+      ccUser cfg `shouldBe` "user"+      ccPassword cfg `shouldBe` "pass"+      ccDatabase cfg `shouldBe` "mydb"++    it "parses postgresql:// scheme" $ do+      let Right cfg = parseConnString "postgresql://u:p@host:1234/db"+      ccHost cfg `shouldBe` "host"+      ccPort cfg `shouldBe` 1234++    it "defaults port to 5432" $ do+      let Right cfg = parseConnString "postgres://u:p@myhost/db"+      ccPort cfg `shouldBe` 5432++    it "parses sslmode=require" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslmode=require"+      ccTls cfg `shouldBe` TlsRequire++    it "parses connect_timeout from query params" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db?connect_timeout=30"+      ccConnectTimeout cfg `shouldBe` 30++    it "defaults connect timeout to 10" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db"+      ccConnectTimeout cfg `shouldBe` 10++    it "defaults query timeout to 0 (no timeout)" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db"+      ccQueryTimeout cfg `shouldBe` 0++  describe "parseConnString (key=value format)" $ do+    it "parses key=value pairs" $ do+      let Right cfg = parseConnString "host=myhost port=9999 user=alice password=secret dbname=testdb"+      ccHost cfg `shouldBe` "myhost"+      ccPort cfg `shouldBe` 9999+      ccUser cfg `shouldBe` "alice"+      ccPassword cfg `shouldBe` "secret"+      ccDatabase cfg `shouldBe` "testdb"++    it "defaults missing fields" $ do+      let Right cfg = parseConnString "host=myhost dbname=db"+      ccPort cfg `shouldBe` 5432+      ccUser cfg `shouldBe` ""++    it "parses connect_timeout" $ do+      let Right cfg = parseConnString "host=h connect_timeout=5"+      ccConnectTimeout cfg `shouldBe` 5++    it "parses sslmode" $ do+      let Right cfg = parseConnString "host=h sslmode=prefer"+      ccTls cfg `shouldBe` TlsPrefer++  describe "defaultConnConfig" $ do+    it "has sensible defaults" $ do+      ccHost defaultConnConfig `shouldBe` "localhost"+      ccPort defaultConnConfig `shouldBe` 5432+      ccTls defaultConnConfig `shouldBe` TlsDisable+      ccConnectTimeout defaultConnConfig `shouldBe` 10+      ccQueryTimeout defaultConnConfig `shouldBe` 0
+ test/PgWire/Connection/EscapingSpec.hs view
@@ -0,0 +1,38 @@+module PgWire.Connection.EscapingSpec (spec) where++import PgWire.Connection (escapeLiteral, escapeIdentifier)+import Test.Hspec++spec :: Spec+spec = do+  describe "escapeLiteral" $ do+    it "wraps in single quotes" $+      escapeLiteral undefined "hello" `shouldBe` "'hello'"++    it "escapes single quotes" $+      escapeLiteral undefined "O'Brien" `shouldBe` "'O''Brien'"++    it "escapes backslashes" $+      escapeLiteral undefined "back\\slash" `shouldBe` "'back\\\\slash'"++    it "handles empty string" $+      escapeLiteral undefined "" `shouldBe` "''"++    it "handles multiple quotes" $+      escapeLiteral undefined "it's a 'test'" `shouldBe` "'it''s a ''test'''"++  describe "escapeIdentifier" $ do+    it "wraps in double quotes" $+      escapeIdentifier undefined "column" `shouldBe` "\"column\""++    it "escapes double quotes" $+      escapeIdentifier undefined "my\"col" `shouldBe` "\"my\"\"col\""++    it "handles spaces" $+      escapeIdentifier undefined "user table" `shouldBe` "\"user table\""++    it "handles empty string" $+      escapeIdentifier undefined "" `shouldBe` "\"\""++    it "handles reserved words" $+      escapeIdentifier undefined "select" `shouldBe` "\"select\""
+ test/PgWire/Connection/FeaturesSpec.hs view
@@ -0,0 +1,102 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+module PgWire.Connection.FeaturesSpec (spec) where++import PgWire.Connection.Config+import Test.Hspec++spec :: Spec+spec = do+  describe "parseHosts" $ do+    it "parses single host" $+      parseHosts "localhost" `shouldBe` ["localhost"]++    it "parses multiple hosts" $+      parseHosts "host1,host2,host3" `shouldBe` ["host1", "host2", "host3"]++    it "returns localhost for empty string" $+      parseHosts "" `shouldBe` ["localhost"]++    it "handles two hosts" $+      parseHosts "primary,replica" `shouldBe` ["primary", "replica"]++  describe "TlsMode parsing" $ do+    it "parses verify-ca from URI" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslmode=verify-ca"+      ccTls cfg `shouldBe` TlsVerifyCa++    it "parses verify-full from URI" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslmode=verify-full"+      ccTls cfg `shouldBe` TlsVerifyFull++    it "parses verify-ca from key-value" $ do+      let Right cfg = parseConnString "host=h sslmode=verify-ca"+      ccTls cfg `shouldBe` TlsVerifyCa++    it "parses verify-full from key-value" $ do+      let Right cfg = parseConnString "host=h sslmode=verify-full"+      ccTls cfg `shouldBe` TlsVerifyFull++  describe "SSL cert params" $ do+    it "parses sslcert from URI" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db?sslcert=/path/cert.pem"+      ccSslCert cfg `shouldBe` Just "/path/cert.pem"++    it "parses sslkey from key-value" $ do+      let Right cfg = parseConnString "host=h sslkey=/path/key.pem"+      ccSslKey cfg `shouldBe` Just "/path/key.pem"++    it "parses sslrootcert" $ do+      let Right cfg = parseConnString "host=h sslrootcert=/path/ca.pem"+      ccSslRootCert cfg `shouldBe` Just "/path/ca.pem"++    it "defaults to Nothing" $ do+      let Right cfg = parseConnString "host=h"+      ccSslCert cfg `shouldBe` Nothing+      ccSslKey cfg `shouldBe` Nothing+      ccSslRootCert cfg `shouldBe` Nothing++  describe "keepalive config" $ do+    it "parses keepalives=0 as False" $ do+      let Right cfg = parseConnString "host=h keepalives=0"+      ccKeepalives cfg `shouldBe` False++    it "defaults keepalives to True" $ do+      let Right cfg = parseConnString "host=h"+      ccKeepalives cfg `shouldBe` True++    it "parses keepalives_idle" $ do+      let Right cfg = parseConnString "host=h keepalives_idle=60"+      ccKeepalivesIdle cfg `shouldBe` 60++    it "parses keepalives_interval" $ do+      let Right cfg = parseConnString "host=h keepalives_interval=10"+      ccKeepalivesInterval cfg `shouldBe` 10++    it "parses keepalives_count" $ do+      let Right cfg = parseConnString "host=h keepalives_count=5"+      ccKeepalivesCount cfg `shouldBe` 5++  describe "target_session_attrs" $ do+    it "parses read-write from URI" $ do+      let Right cfg = parseConnString "postgres://u:p@h:5432/db?target_session_attrs=read-write"+      ccTargetSessionAttrs cfg `shouldBe` SessionReadWrite++    it "parses primary from key-value" $ do+      let Right cfg = parseConnString "host=h target_session_attrs=primary"+      ccTargetSessionAttrs cfg `shouldBe` SessionPrimary++    it "parses standby" $ do+      let Right cfg = parseConnString "host=h target_session_attrs=standby"+      ccTargetSessionAttrs cfg `shouldBe` SessionStandby++    it "parses prefer-standby" $ do+      let Right cfg = parseConnString "host=h target_session_attrs=prefer-standby"+      ccTargetSessionAttrs cfg `shouldBe` SessionPreferStandby++    it "parses read-only" $ do+      let Right cfg = parseConnString "host=h target_session_attrs=read-only"+      ccTargetSessionAttrs cfg `shouldBe` SessionReadOnly++    it "defaults to any" $ do+      let Right cfg = parseConnString "host=h"+      ccTargetSessionAttrs cfg `shouldBe` SessionAny
+ test/PgWire/ErrorSpec.hs view
@@ -0,0 +1,62 @@+module PgWire.ErrorSpec (spec) where++import Control.Exception (catch)+import PgWire.Error+import PgWire.Protocol.Backend (PgError (..))+import Test.Hspec++spec :: Spec+spec = do+  describe "PgWireError constructors" $ do+    it "ConnectionError carries message" $ do+      let err = ConnectionError "host unreachable"+      show err `shouldContain` "host unreachable"++    it "AuthError carries message" $ do+      let err = AuthError "bad password"+      show err `shouldContain` "bad password"++    it "ProtocolError carries message" $ do+      let err = ProtocolError "unexpected tag"+      show err `shouldContain` "unexpected tag"++    it "DecodeError carries message" $ do+      let err = DecodeError "int32: wrong size"+      show err `shouldContain` "int32: wrong size"++    it "PoolTimeout has no payload" $+      show PoolTimeout `shouldContain` "PoolTimeout"++    it "PoolClosed has no payload" $+      show PoolClosed `shouldContain` "PoolClosed"++    it "QueryError wraps PgError" $ do+      let pgErr = emptyPgError {pgMessage = "table not found", pgCode = "42P01"}+          err = QueryError pgErr+      show err `shouldContain` "table not found"+      show err `shouldContain` "42P01"++  describe "throwPgWire" $ do+    it "throws as an Exception" $ do+      let action = throwPgWire (ConnectionError "test") :: IO ()+      action `shouldThrow` (== ConnectionError "test")++    it "can be caught with specific pattern" $ do+      result <- (throwPgWire PoolTimeout >> pure "no") `catch` \(e :: PgWireError) ->+        case e of+          PoolTimeout -> pure "caught"+          _ -> pure "wrong"+      result `shouldBe` ("caught" :: String)++  describe "Eq instance" $ do+    it "equal errors are equal" $+      ConnectionError "x" `shouldBe` ConnectionError "x"++    it "different errors are not equal" $+      ConnectionError "x" `shouldNotBe` ConnectionError "y"++    it "different constructors are not equal" $+      ConnectionError "x" `shouldNotBe` AuthError "x"++emptyPgError :: PgError+emptyPgError = PgError "" "" "" Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing Nothing
+ test/PgWire/MockServerSpec.hs view
@@ -0,0 +1,113 @@+module PgWire.MockServerSpec (spec) where++import Control.Exception (try)+import Data.ByteString.Char8 qualified as BS8+import PgWire.Connection (Connection, close, connectString, simpleQuery)+import PgWire.Error (PgWireError (..))+import PgWire.MockServer+import Test.Hspec++connectToMock :: MockConfig -> (Connection -> IO a) -> IO a+connectToMock cfg action =+  withMockServer cfg $ \port -> do+    let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb"+    conn <- connectString connStr+    result <- action conn+    close conn+    pure result++spec :: Spec+spec = do+  describe "MockServer" $ do+    describe "basic protocol" $ do+      it "accepts a connection and responds to simple query" $ do+        connectToMock defaultMockConfig $ \conn -> do+          (rows, _) <- simpleQuery conn "SELECT 1"+          rows `shouldBe` []++      it "returns rows from a custom handler" $ do+        let cfg = defaultMockConfig+              { mockQueryHandler = simpleHandler+                  [ ("SELECT name FROM test",+                     [ [("name", "Alice")]+                     , [("name", "Bob")]+                     ])+                  ]+              }+        connectToMock cfg $ \conn -> do+          (rows, _) <- simpleQuery conn "SELECT name FROM test"+          length rows `shouldBe` 2++      it "handles unknown queries gracefully" $ do+        let cfg = defaultMockConfig+              { mockQueryHandler = simpleHandler []+              }+        connectToMock cfg $ \conn -> do+          (rows, _) <- simpleQuery conn "SELECT nothing"+          rows `shouldBe` []++      it "supports multiple sequential queries" $ do+        connectToMock defaultMockConfig $ \conn -> do+          _ <- simpleQuery conn "SELECT 1"+          _ <- simpleQuery conn "SELECT 2"+          (rows, _) <- simpleQuery conn "SELECT 3"+          rows `shouldBe` []++    describe "error responses" $ do+      it "propagates server errors as QueryError" $ do+        let cfg = defaultMockConfig+              { mockQueryHandler = errorHandler "42P01" "relation \"foo\" does not exist"+              }+        connectToMock cfg $ \conn -> do+          result <- try @PgWireError (simpleQuery conn "SELECT * FROM foo")+          case result of+            Left (QueryError _) -> pure ()+            Left other -> expectationFailure $ "Expected QueryError, got: " <> show other+            Right _ -> expectationFailure "Expected an error"++      it "recovers after an error and accepts subsequent queries" $ do+        let cfg = defaultMockConfig+              { mockQueryHandler = \sql send ->+                  if sql == "BAD"+                    then send $ buildErrorResponse "42601" "syntax error"+                    else send $ buildCommandComplete "SELECT 0"+              }+        connectToMock cfg $ \conn -> do+          -- First query errors+          result1 <- try @PgWireError (simpleQuery conn "BAD")+          case result1 of+            Left (QueryError _) -> pure ()+            _ -> expectationFailure "Expected QueryError for BAD query"+          -- Second query succeeds on the same connection+          (rows, _) <- simpleQuery conn "GOOD"+          rows `shouldBe` []++    describe "multi-row results" $ do+      it "returns correct number of rows for large result sets" $ do+        let manyRows = [ [("id", BS8.pack (show i)), ("val", "x")]+                        | i <- [1 :: Int .. 100]+                        ]+            cfg = defaultMockConfig+              { mockQueryHandler = simpleHandler+                  [("SELECT id, val FROM big", manyRows)]+              }+        connectToMock cfg $ \conn -> do+          (rows, _) <- simpleQuery conn "SELECT id, val FROM big"+          length rows `shouldBe` 100++      it "returns multiple columns correctly" $ do+        let cfg = defaultMockConfig+              { mockQueryHandler = simpleHandler+                  [("SELECT a, b, c FROM t",+                    [ [("a", "1"), ("b", "hello"), ("c", "true")]+                    , [("a", "2"), ("b", "world"), ("c", "false")]+                    ])]+              }+        connectToMock cfg $ \conn -> do+          (rows, _) <- simpleQuery conn "SELECT a, b, c FROM t"+          length rows `shouldBe` 2+          case rows of+            [row1, row2] -> do+              length row1 `shouldBe` 3+              length row2 `shouldBe` 3+            _ -> expectationFailure "Expected exactly 2 rows"
+ test/PgWire/Pool/ConfigSpec.hs view
@@ -0,0 +1,22 @@+module PgWire.Pool.ConfigSpec (spec) where++import PgWire.Pool.Config+import Test.Hspec++spec :: Spec+spec = do+  describe "defaultPoolConfig" $ do+    it "has empty connection string" $+      poolConnString defaultPoolConfig `shouldBe` ""++    it "has pool size 10" $+      poolSize defaultPoolConfig `shouldBe` 10++    it "has idle time 600 seconds" $+      poolIdleTime defaultPoolConfig `shouldBe` 600++    it "has max life 3600 seconds" $+      poolMaxLife defaultPoolConfig `shouldBe` 3600++    it "has acquire timeout 10 seconds" $+      poolAcquireTimeout defaultPoolConfig `shouldBe` 10
+ test/PgWire/Pool/NoThunksSpec.hs view
@@ -0,0 +1,23 @@+module PgWire.Pool.NoThunksSpec (spec) where++import NoThunks.Class (noThunks)+import PgWire.Pool (PoolStats (..))+import Test.Hspec++spec :: Spec+spec = do+  describe "NoThunks" $ do+    it "PoolStats has no thunks when fully evaluated" $ do+      let stats = PoolStats+            { psIdle = 3+            , psInUse = 5+            , psWaiters = 0+            , psMaxSize = 10+            , psTotalCreated = 20+            , psTotalDestroyed = 12+            , psTotalTimeouts = 1+            }+      result <- noThunks [] stats+      case result of+        Nothing -> pure ()+        Just thunkInfo -> expectationFailure $ "Found thunk: " <> show thunkInfo
+ test/PgWire/Pool/PropertySpec.hs view
@@ -0,0 +1,128 @@+module PgWire.Pool.PropertySpec (spec) where++import Control.Concurrent.Async (mapConcurrently)+import Control.Exception (try)+import Data.ByteString.Char8 qualified as BS8+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import PgWire.Connection (simpleQuery)+import PgWire.Error (PgWireError (..))+import PgWire.MockServer (defaultMockConfig, withMockServer)+import PgWire.Pool (Pool, PoolStats (..), closePool, newPool, poolIsAlive, poolStats, resize, withResource)+import PgWire.Pool.Config (PoolConfig (..), defaultPoolConfig)+import Test.Hspec+import Test.Hspec.Hedgehog (hedgehog)++-- | Create a pool backed by the mock server.+withMockPool :: Int -> (Pool -> IO a) -> IO a+withMockPool size action =+  withMockServer defaultMockConfig $ \port -> do+    let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb"+    pool <- newPool defaultPoolConfig+      { poolConnString = connStr+      , poolSize = size+      , poolAcquireTimeout = 5+      }+    result <- action pool+    closePool pool+    pure result++spec :: Spec+spec = do+  describe "Pool invariants" $ do+    it "idle + inUse never exceeds maxSize" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 1 10))+      nOps <- forAll (Gen.int (Range.linear 1 20))+      result <- evalIO $ withMockPool size $ \pool -> do+        -- Do several acquire/release cycles+        mapM_ (\_ -> withResource pool $ \conn ->+          simpleQuery conn "SELECT 1") [1 .. nOps]+        stats <- poolStats pool+        pure (psIdle stats + psInUse stats <= psMaxSize stats)+      assert result++    it "all connections returned after withResource" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 1 8))+      nOps <- forAll (Gen.int (Range.linear 1 15))+      stats <- evalIO $ withMockPool size $ \pool -> do+        mapM_ (\_ -> withResource pool $ \conn ->+          simpleQuery conn "SELECT 1") [1 .. nOps]+        poolStats pool+      psInUse stats === 0++    it "concurrent acquire/release maintains invariant" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 2 8))+      nThreads <- forAll (Gen.int (Range.linear 2 10))+      stats <- evalIO $ withMockPool size $ \pool -> do+        _ <- mapConcurrently (\_ ->+          withResource pool $ \conn ->+            simpleQuery conn "SELECT 1") [1 .. nThreads]+        poolStats pool+      -- After all threads finish, no connections should be in use+      psInUse stats === 0+      -- idle + inUse should not exceed max+      assert (psIdle stats <= size)++    it "totalCreated >= totalDestroyed" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 1 5))+      nOps <- forAll (Gen.int (Range.linear 1 10))+      stats <- evalIO $ withMockPool size $ \pool -> do+        mapM_ (\_ -> withResource pool $ \conn ->+          simpleQuery conn "SELECT 1") [1 .. nOps]+        poolStats pool+      assert (psTotalCreated stats >= psTotalDestroyed stats)++    it "closePool makes pool not alive" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 1 5))+      alive <- evalIO $ withMockServer defaultMockConfig $ \port -> do+        let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb"+        pool <- newPool defaultPoolConfig+          { poolConnString = connStr+          , poolSize = size+          , poolAcquireTimeout = 5+          }+        closePool pool+        poolIsAlive pool+      assert (not alive)++    it "closePool rejects new acquires" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 1 5))+      gotError <- evalIO $ withMockServer defaultMockConfig $ \port -> do+        let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb"+        pool <- newPool defaultPoolConfig+          { poolConnString = connStr+          , poolSize = size+          , poolAcquireTimeout = 5+          }+        closePool pool+        result <- try (withResource pool $ \_ -> pure ())+        pure $ case result of+          Left PoolClosed -> True+          Left _ -> False+          Right _ -> False+      assert gotError++    it "resize changes maxSize" $ hedgehog $ do+      initial <- forAll (Gen.int (Range.linear 2 10))+      newSize <- forAll (Gen.int (Range.linear 1 20))+      stats <- evalIO $ withMockPool initial $ \pool -> do+        resize pool newSize+        poolStats pool+      psMaxSize stats === newSize++    it "pool stats are consistent" $ hedgehog $ do+      size <- forAll (Gen.int (Range.linear 1 8))+      nOps <- forAll (Gen.int (Range.linear 0 10))+      stats <- evalIO $ withMockPool size $ \pool -> do+        mapM_ (\_ -> withResource pool $ \conn ->+          simpleQuery conn "SELECT 1") [1 .. nOps]+        poolStats pool+      -- Basic consistency checks+      assert (psIdle stats >= 0)+      assert (psInUse stats >= 0)+      assert (psWaiters stats >= 0)+      assert (psMaxSize stats > 0)+      assert (psTotalCreated stats >= 0)+      assert (psTotalDestroyed stats >= 0)+      assert (psTotalTimeouts stats >= 0)
+ test/PgWire/Pool/StateMachineSpec.hs view
@@ -0,0 +1,178 @@+-- | Lightweight state-machine tests for the connection pool.+--+-- Existing 'PgWire.Pool.PropertySpec' tests cover individual scenarios+-- (a burst of acquires, a single resize, a single close). This module+-- generates arbitrary sequences of Burst/Resize/Close and checks the+-- model-vs-real invariants after every single step. The model is+-- hand-rolled rather than using 'Hedgehog.Internal.State' so the test+-- drops straight into the existing hspec-hedgehog runner.+module PgWire.Pool.StateMachineSpec (spec) where++import Control.Concurrent.Async (mapConcurrently_)+import Control.Exception (try)+import Data.ByteString.Char8 qualified as BS8+import Hedgehog hiding (Command)+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import PgWire.Connection (simpleQuery)+import PgWire.Error (PgWireError (..))+import PgWire.MockServer (defaultMockConfig, withMockServer)+import PgWire.Pool+  ( Pool+  , PoolStats (..)+  , closePool+  , newPool+  , poolIsAlive+  , poolStats+  , resize+  , withResource+  )+import PgWire.Pool.Config (PoolConfig (..), defaultPoolConfig)+import Test.Hspec+import Test.Hspec.Hedgehog (hedgehog)++data Command+  = CmdBurst !Int+  | CmdResize !Int+  | CmdClose+  deriving (Show)++data Model = Model+  { mMaxSize :: !Int+  , mClosed :: !Bool+  }+  deriving (Show)++initialModel :: Int -> Model+initialModel size = Model {mMaxSize = size, mClosed = False}++stepModel :: Model -> Command -> Model+stepModel m cmd = case cmd of+  CmdBurst _ -> m+  CmdResize n -> m {mMaxSize = n}+  CmdClose -> m {mClosed = True}++-- | Burst issues n concurrent acquire/release cycles. Errors are+-- swallowed: if the pool is closed or the acquire times out, the+-- model already reflects that and the invariant pass after the+-- command settles is what matters.+runCommand :: Pool -> Command -> IO ()+runCommand pool cmd = case cmd of+  CmdBurst n ->+    mapConcurrently_+      ( \_ -> do+          _ <-+            try @PgWireError $+              withResource pool $ \conn -> simpleQuery conn "SELECT 1"+          pure ()+      )+      [1 .. n]+  CmdResize n -> resize pool n+  CmdClose -> closePool pool++genCommand :: Gen Command+genCommand =+  Gen.frequency+    [ (6, CmdBurst <$> Gen.int (Range.linear 1 6))+    , (2, CmdResize <$> Gen.int (Range.linear 1 8))+    , (1, pure CmdClose)+    ]++-- | Generate a program with at most one Close so we don't burn cycles+-- on post-close no-ops. Close may appear at any position or not at all.+genCommands :: Gen [Command]+genCommands = do+  len <- Gen.int (Range.linear 1 12)+  cmds <- Gen.list (Range.singleton len) (Gen.filterT (not . isClose) genCommand)+  includeClose <- Gen.bool+  if includeClose+    then do+      pos <- Gen.int (Range.linear 0 (length cmds))+      let (hd, tl) = splitAt pos cmds+      pure (hd <> [CmdClose] <> tl)+    else pure cmds+  where+    isClose CmdClose = True+    isClose _ = False++-- | A per-step observation collected while running the program, then+-- re-checked inside the Hedgehog property.+data Step = Step+  { stCmd :: !Command+  , stModel :: !Model+  , stStats :: !PoolStats+  , stAlive :: !Bool+  , stClosedProbe :: !(Maybe (Either PgWireError ()))+  -- ^ If the model says closed, probe the real pool by trying to+  -- acquire a connection. Expected: Left PoolClosed.+  }+  deriving (Show)++runProgram :: Pool -> Int -> [Command] -> IO [Step]+runProgram pool initialSize = go (initialModel initialSize) []+  where+    go _ acc [] = pure (reverse acc)+    go m acc (c : cs) = do+      runCommand pool c+      let m' = stepModel m c+      stats <- poolStats pool+      alive <- poolIsAlive pool+      probe <-+        if mClosed m'+          then do+            r <- try @PgWireError (withResource pool $ \_ -> pure ())+            pure (Just r)+          else pure Nothing+      let step = Step c m' stats alive probe+      go m' (step : acc) cs++withMockPool :: Int -> (Pool -> IO a) -> IO a+withMockPool size action =+  withMockServer defaultMockConfig $ \port -> do+    let connStr = "postgres://testuser:testpass@127.0.0.1:" <> BS8.pack (show port) <> "/testdb"+    pool <-+      newPool+        defaultPoolConfig+          { poolConnString = connStr+          , poolSize = size+          , poolAcquireTimeout = 2+          }+    r <- action pool+    closePool pool -- idempotent; safe even if the program closed it+    pure r++checkStep :: (MonadTest m) => Step -> m ()+checkStep Step {stCmd, stModel, stStats, stAlive, stClosedProbe} = do+  annotateShow stCmd+  annotateShow stModel+  annotateShow stStats+  stAlive === not (mClosed stModel)+  psMaxSize stStats === mMaxSize stModel+  -- After a command has settled, nothing should remain in flight.+  psInUse stStats === 0+  -- Idle never exceeds the current cap (resize-down drops excess idle+  -- eagerly, so even a just-executed shrink should satisfy this).+  assert (psIdle stStats <= mMaxSize stModel)+  assert (psIdle stStats >= 0)+  assert (psWaiters stStats >= 0)+  assert (psTotalCreated stStats >= 0)+  assert (psTotalDestroyed stStats >= 0)+  assert (psTotalTimeouts stStats >= 0)+  -- Destroy count can never exceed create count.+  assert (psTotalCreated stStats >= psTotalDestroyed stStats)+  case stClosedProbe of+    Nothing -> pure ()+    Just (Left PoolClosed) -> pure ()+    Just other -> do+      annotateShow other+      failure++spec :: Spec+spec = do+  describe "Pool state machine" $ do+    it "model-vs-pool invariants hold after every command" $ hedgehog $ do+      initialSize <- forAll (Gen.int (Range.linear 1 6))+      program <- forAll genCommands+      steps <- evalIO $ withMockPool initialSize $ \pool ->+        runProgram pool initialSize program+      mapM_ checkStep steps
+ test/PgWire/Protocol/BuildersSpec.hs view
@@ -0,0 +1,181 @@+module PgWire.Protocol.BuildersSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int32)+import Data.Vector qualified as V+import PgWire.Protocol.Builders (buildFrontendMsg, buildStartup)+import PgWire.Protocol.Frontend+import Test.Hspec++spec :: Spec+spec = do+  describe "zero-payload messages" $ do+    it "Sync is tag 'S' + length 4" $ do+      let bs = buildFrontendMsg Sync+      BS.length bs `shouldBe` 5+      BS.index bs 0 `shouldBe` 83++    it "Terminate is tag 'X' + length 4" $ do+      let bs = buildFrontendMsg Terminate+      BS.length bs `shouldBe` 5+      BS.index bs 0 `shouldBe` 88++    it "Flush is tag 'H' + length 4" $ do+      let bs = buildFrontendMsg Flush+      BS.length bs `shouldBe` 5+      BS.index bs 0 `shouldBe` 72++    it "CopyDone is tag 'c' + length 4" $ do+      let bs = buildFrontendMsg CopyDone+      BS.length bs `shouldBe` 5+      BS.index bs 0 `shouldBe` 99++  describe "Query" $ do+    it "has tag 'Q' and NUL-terminated SQL" $ do+      let bs = buildFrontendMsg (Query "SELECT 1")+      BS.index bs 0 `shouldBe` 81+      BS.last bs `shouldBe` 0+      "SELECT 1" `BS.isInfixOf` bs `shouldBe` True++    it "length includes itself + sql + NUL" $ do+      let bs = buildFrontendMsg (Query "SEL")+      -- tag(1) + len(4) + "SEL"(3) + NUL(1) = 9+      -- length field = 4 + 3 + 1 = 8+      BS.length bs `shouldBe` 9+      let lenBytes = BS.take 4 (BS.drop 1 bs)+      decodeInt32 lenBytes `shouldBe` 8++    it "handles empty query" $ do+      let bs = buildFrontendMsg (Query "")+      -- tag(1) + len(4) + NUL(1) = 6+      BS.length bs `shouldBe` 6++  describe "Parse" $ do+    it "has tag 'P'" $ do+      let bs = buildFrontendMsg (Parse "s1" "SELECT $1" (V.singleton 23))+      BS.index bs 0 `shouldBe` 80++    it "contains statement name and SQL" $ do+      let bs = buildFrontendMsg (Parse "mystmt" "SELECT * FROM t" V.empty)+      "mystmt" `BS.isInfixOf` bs `shouldBe` True+      "SELECT * FROM t" `BS.isInfixOf` bs `shouldBe` True++    it "encodes param count and OIDs" $ do+      let bs = buildFrontendMsg (Parse "" "SELECT $1, $2" (V.fromList [23, 25]))+      -- Should contain the Int16 param count (2) and two Int32 OIDs+      BS.length bs `shouldSatisfy` (> 15)++  describe "Bind" $ do+    it "has tag 'B'" $ do+      let bs = buildFrontendMsg (Bind "" "s1" V.empty V.empty V.empty)+      BS.index bs 0 `shouldBe` 66++    it "encodes portal and statement names" $ do+      let bs = buildFrontendMsg (Bind "portal" "stmt" V.empty V.empty V.empty)+      "portal" `BS.isInfixOf` bs `shouldBe` True+      "stmt" `BS.isInfixOf` bs `shouldBe` True++    it "encodes parameter values including NULLs" $ do+      let params = V.fromList [Just "hello", Nothing, Just "world"]+          bs = buildFrontendMsg (Bind "" "" V.empty params V.empty)+      "hello" `BS.isInfixOf` bs `shouldBe` True+      "world" `BS.isInfixOf` bs `shouldBe` True++  describe "Execute" $ do+    it "has tag 'E'" $ do+      let bs = buildFrontendMsg (Execute "" 0)+      BS.index bs 0 `shouldBe` 69++    it "encodes max rows" $ do+      let bs = buildFrontendMsg (Execute "" 100)+      BS.length bs `shouldSatisfy` (> 5)++  describe "Describe" $ do+    it "has tag 'D' for statement" $ do+      let bs = buildFrontendMsg (Describe DescribeStatement "s1")+      BS.index bs 0 `shouldBe` 68++    it "has tag 'D' for portal" $ do+      let bs = buildFrontendMsg (Describe DescribePortal "p1")+      BS.index bs 0 `shouldBe` 68++  describe "Close" $ do+    it "has tag 'C'" $ do+      let bs = buildFrontendMsg (Close DescribeStatement "s1")+      BS.index bs 0 `shouldBe` 67++  describe "PasswordMessage" $ do+    it "has tag 'p'" $ do+      let bs = buildFrontendMsg (PasswordMessage "secret")+      BS.index bs 0 `shouldBe` 112+      "secret" `BS.isInfixOf` bs `shouldBe` True++  describe "CopyData" $ do+    it "has tag 'd' and preserves payload" $ do+      let bs = buildFrontendMsg (CopyData "row data")+      BS.index bs 0 `shouldBe` 100+      "row data" `BS.isInfixOf` bs `shouldBe` True++  describe "CopyFail" $ do+    it "has tag 'f' and includes error message" $ do+      let bs = buildFrontendMsg (CopyFail "bad data")+      BS.index bs 0 `shouldBe` 102+      "bad data" `BS.isInfixOf` bs `shouldBe` True++  describe "SASLInitialResponse" $ do+    it "has tag 'p'" $ do+      let bs = buildFrontendMsg (SASLInitialResponse "SCRAM-SHA-256" "client-first")+      BS.index bs 0 `shouldBe` 112+      "SCRAM-SHA-256" `BS.isInfixOf` bs `shouldBe` True++  describe "SASLResponse" $ do+    it "has tag 'p'" $ do+      let bs = buildFrontendMsg (SASLResponse "client-final")+      BS.index bs 0 `shouldBe` 112++  describe "buildStartup" $ do+    it "starts with length and protocol 3.0" $ do+      let bs = buildStartup (StartupParams "u" "d" "" [])+      BS.length bs `shouldSatisfy` (> 8)+      let version = decodeInt32 (BS.take 4 (BS.drop 4 bs))+      version `shouldBe` 196608++    it "includes user and database" $ do+      let bs = buildStartup (StartupParams "alice" "mydb" "" [])+      "user" `BS.isInfixOf` bs `shouldBe` True+      "alice" `BS.isInfixOf` bs `shouldBe` True+      "database" `BS.isInfixOf` bs `shouldBe` True+      "mydb" `BS.isInfixOf` bs `shouldBe` True++    it "includes application_name when non-empty" $ do+      let bs = buildStartup (StartupParams "u" "d" "myapp" [])+      "application_name" `BS.isInfixOf` bs `shouldBe` True+      "myapp" `BS.isInfixOf` bs `shouldBe` True++    it "omits application_name when empty" $ do+      let bs = buildStartup (StartupParams "u" "d" "" [])+      "application_name" `BS.isInfixOf` bs `shouldBe` False++    it "includes extra params" $ do+      let bs = buildStartup (StartupParams "u" "d" "" [("extra_float_digits", "2")])+      "extra_float_digits" `BS.isInfixOf` bs `shouldBe` True+      "2" `BS.isInfixOf` bs `shouldBe` True++    it "ends with NUL terminator" $ do+      let bs = buildStartup (StartupParams "u" "d" "" [])+      BS.last bs `shouldBe` 0++    it "length field matches actual length" $ do+      let bs = buildStartup (StartupParams "testuser" "testdb" "" [])+          declaredLen = decodeInt32 (BS.take 4 bs)+      declaredLen `shouldBe` fromIntegral (BS.length bs)++-- Helpers -------------------------------------------------------------------++decodeInt32 :: BS.ByteString -> Int32+decodeInt32 bs =+  let b0 = fromIntegral (BS.index bs 0) :: Int32+      b1 = fromIntegral (BS.index bs 1) :: Int32+      b2 = fromIntegral (BS.index bs 2) :: Int32+      b3 = fromIntegral (BS.index bs 3) :: Int32+   in b0 * 256 * 256 * 256 + b1 * 256 * 256 + b2 * 256 + b3
+ test/PgWire/Protocol/FuzzSpec.hs view
@@ -0,0 +1,360 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module PgWire.Protocol.FuzzSpec (spec) where++import Control.DeepSeq (NFData (..), deepseq)+import Control.Exception (evaluate, try, SomeException)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int16, Int32)+import Data.Word (Word8, Word32)+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import PgWire.Protocol.Backend (CommandTag (..))+import PgWire.Protocol.Parsers (parseBackendMsg, parseCommandTag)+import Test.Hspec+import Test.Hspec.Hedgehog (hedgehog)++-- | NFData instance for CommandTag so we can deepseq it.+instance NFData CommandTag where+  rnf (SelectTag n) = n `seq` ()+  rnf (InsertTag n) = n `seq` ()+  rnf (UpdateTag n) = n `seq` ()+  rnf (DeleteTag n) = n `seq` ()+  rnf (OtherTag b) = rnf b++-- | Tags whose parsers are purely bounds-checked (no unbounded recursion).+-- These are safe to fuzz with completely random payloads even on a small+-- stack (the test suite runs with -K8K).+safeTags :: [Word8]+safeTags =+  [ 82  -- R: Authentication+  , 75  -- K: BackendKeyData+  , 83  -- S: ParameterStatus+  , 90  -- Z: ReadyForQuery+  , 49  -- 1: ParseComplete+  , 50  -- 2: BindComplete+  , 51  -- 3: CloseComplete+  , 73  -- I: EmptyQueryResponse+  , 99  -- c: CopyDone+  , 100 -- d: CopyData+  , 110 -- n: NoData+  , 65  -- A: NotificationResponse+  ]++-- | Tags whose parsers may recurse on their payload (ErrorResponse,+-- NoticeResponse) or iterate based on an encoded count (DataRow,+-- RowDescription, ParameterDescription, CopyIn/OutResponse,+-- CommandComplete). We fuzz these with payloads that contain NUL+-- terminators and bounded counts so they do not overflow the tiny+-- test-suite stack.+recursiveTags :: [Word8]+recursiveTags =+  [ 67  -- C: CommandComplete+  , 68  -- D: DataRow+  , 69  -- E: ErrorResponse+  , 71  -- G: CopyInResponse+  , 72  -- H: CopyOutResponse+  , 78  -- N: NoticeResponse+  , 84  -- T: RowDescription+  , 116 -- t: ParameterDescription+  ]++-- | Evaluate a parse result to WHNF (forcing the Either constructor).+-- The module uses StrictData and bang patterns internally, so forcing+-- the outer Either exercises the parser.  Returns True when the parser+-- produces Left or Right without throwing.+doesNotCrash :: Word8 -> ByteString -> IO Bool+doesNotCrash tag payload = do+  result <- try $ evaluate $ case parseBackendMsg tag payload of+    Left _ -> ()+    Right _ -> ()+  case result of+    Left (_ :: SomeException) -> pure False+    Right () -> pure True++spec :: Spec+spec = do+  -- ── 1. No crashes on arbitrary input ──────────────────────────────+  describe "No crashes on arbitrary input" $ do+    it "safe tags with completely random payloads" $ hedgehog $ do+      tag <- forAll (Gen.element safeTags)+      payload <- forAll (Gen.bytes (Range.linear 0 128))+      ok <- evalIO (doesNotCrash tag payload)+      assert ok++    it "unknown tags with random payloads" $ hedgehog $ do+      -- Tags outside the known set always return Left immediately+      tag <- forAll $ Gen.filter (\t -> t `notElem` safeTags && t `notElem` recursiveTags)+                                 (Gen.word8 Range.linearBounded)+      payload <- forAll (Gen.bytes (Range.linear 0 128))+      ok <- evalIO (doesNotCrash tag payload)+      assert ok++    it "ErrorResponse with NUL-bearing payloads" $ hedgehog $ do+      -- Generate payloads that contain NUL bytes so the field parser terminates+      payload <- forAll genNulBearingPayload+      ok <- evalIO (doesNotCrash 69 payload)+      assert ok++    it "NoticeResponse with NUL-bearing payloads" $ hedgehog $ do+      payload <- forAll genNulBearingPayload+      ok <- evalIO (doesNotCrash 78 payload)+      assert ok++    it "CommandComplete with arbitrary tag strings" $ hedgehog $ do+      tagStr <- forAll (Gen.bytes (Range.linear 0 64))+      -- CommandComplete expects a NUL-terminated tag+      let payload = tagStr <> BS.singleton 0+      ok <- evalIO (doesNotCrash 67 payload)+      assert ok++    it "RowDescription with bounded field count" $ hedgehog $ do+      payload <- forAll genRowDescPayload+      ok <- evalIO (doesNotCrash 84 payload)+      assert ok++    it "ParameterDescription with bounded param count" $ hedgehog $ do+      nParams <- forAll (Gen.int16 (Range.linear 0 20))+      trailing <- forAll (Gen.bytes (Range.linear 0 128))+      let payload = int16 nParams <> trailing+      ok <- evalIO (doesNotCrash 116 payload)+      assert ok++    it "CopyInResponse with bounded column count" $ hedgehog $ do+      payload <- forAll genCopyResponsePayload+      ok <- evalIO (doesNotCrash 71 payload)+      assert ok++    it "CopyOutResponse with bounded column count" $ hedgehog $ do+      payload <- forAll genCopyResponsePayload+      ok <- evalIO (doesNotCrash 72 payload)+      assert ok++    it "medium payloads do not hang (safe tags)" $ hedgehog $ do+      tag <- forAll (Gen.element safeTags)+      payload <- forAll (Gen.bytes (Range.linear 64 256))+      ok <- evalIO (doesNotCrash tag payload)+      assert ok++  -- ── 2. No crashes on truncated messages ───────────────────────────+  describe "No crashes on truncated messages" $ do+    it "truncated AuthOk" $ hedgehog $ do+      let fullPayload = int32 0+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 82 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated AuthMD5" $ hedgehog $ do+      let fullPayload = int32 5 <> BS.pack [0xDE, 0xAD, 0xBE, 0xEF]+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 82 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated BackendKeyData" $ hedgehog $ do+      let fullPayload = int32 12345 <> int32 67890+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 75 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated ReadyForQuery" $ hedgehog $ do+      ok <- evalIO (doesNotCrash 90 BS.empty)+      assert ok++    it "truncated DataRow with columns" $ hedgehog $ do+      nCols <- forAll (Gen.int (Range.linear 1 10))+      let colPayload = mconcat [int32 3 <> "abc" | _ <- [1 :: Int .. nCols]]+          fullPayload = int16 (fromIntegral nCols) <> colPayload+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 68 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated RowDescription" $ hedgehog $ do+      let field = "col\0" <> word32 0 <> int16 1 <> word32 23 <> int16 4 <> int32 (-1) <> int16 0+          fullPayload = int16 1 <> field+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 84 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated ParameterDescription" $ hedgehog $ do+      nParams <- forAll (Gen.int (Range.linear 1 5))+      let fullPayload = int16 (fromIntegral nParams)+            <> mconcat [word32 (fromIntegral i) | i <- [1 :: Int .. nParams]]+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 116 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated CopyInResponse" $ hedgehog $ do+      let fullPayload = BS.singleton 0 <> int16 2 <> int16 0 <> int16 1+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 71 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated ErrorResponse" $ hedgehog $ do+      let fullPayload = BS.concat+            [ BS.singleton 83, "ERROR\0"+            , BS.singleton 67, "42P01\0"+            , BS.singleton 77, "relation not found\0"+            , BS.singleton 0+            ]+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 69 (BS.take truncateAt fullPayload))+      assert ok++    it "truncated NotificationResponse" $ hedgehog $ do+      let fullPayload = int32 42 <> "my_channel\0" <> "hello\0"+      truncateAt <- forAll (Gen.int (Range.linear 0 (BS.length fullPayload - 1)))+      ok <- evalIO (doesNotCrash 65 (BS.take truncateAt fullPayload))+      assert ok++  -- ── 3. CommandTag fuzz ────────────────────────────────────────────+  describe "CommandTag fuzz" $ do+    it "parseCommandTag handles arbitrary bytestrings" $ hedgehog $ do+      bs <- forAll (Gen.bytes (Range.linear 0 256))+      ok <- evalIO (commandTagDoesNotCrash bs)+      assert ok++    it "parseCommandTag handles SELECT prefix with garbage" $ hedgehog $ do+      suffix <- forAll (Gen.bytes (Range.linear 0 64))+      ok <- evalIO (commandTagDoesNotCrash ("SELECT " <> suffix))+      assert ok++    it "parseCommandTag handles INSERT prefix with garbage" $ hedgehog $ do+      suffix <- forAll (Gen.bytes (Range.linear 0 64))+      ok <- evalIO (commandTagDoesNotCrash ("INSERT " <> suffix))+      assert ok++    it "parseCommandTag handles UPDATE prefix with garbage" $ hedgehog $ do+      suffix <- forAll (Gen.bytes (Range.linear 0 64))+      ok <- evalIO (commandTagDoesNotCrash ("UPDATE " <> suffix))+      assert ok++    it "parseCommandTag handles DELETE prefix with garbage" $ hedgehog $ do+      suffix <- forAll (Gen.bytes (Range.linear 0 64))+      ok <- evalIO (commandTagDoesNotCrash ("DELETE " <> suffix))+      assert ok++  -- ── 4. DataRow bounds safety ──────────────────────────────────────+  describe "DataRow bounds safety" $ do+    it "well-formed random columns" $ hedgehog $ do+      nCols <- forAll (Gen.int (Range.linear 0 50))+      cols <- forAll (Gen.list (Range.singleton nCols) genColumnData)+      let payload = int16 (fromIntegral nCols) <> mconcat cols+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++    it "declared column count exceeding actual data" $ hedgehog $ do+      declaredCols <- forAll (Gen.int (Range.linear 1 100))+      actualCols <- forAll (Gen.int (Range.linear 0 (declaredCols - 1)))+      cols <- forAll (Gen.list (Range.singleton actualCols) genColumnData)+      let payload = int16 (fromIntegral declaredCols) <> mconcat cols+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++    it "zero-length columns" $ hedgehog $ do+      nCols <- forAll (Gen.int (Range.linear 1 20))+      let payload = int16 (fromIntegral nCols)+            <> mconcat (replicate nCols (int32 0))+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++    it "all NULL columns" $ hedgehog $ do+      nCols <- forAll (Gen.int (Range.linear 1 50))+      let payload = int16 (fromIntegral nCols)+            <> mconcat (replicate nCols (int32 (-1)))+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++    it "very large declared column length" $ hedgehog $ do+      bigLen <- forAll (Gen.int32 (Range.linear 1000 2147483647))+      let payload = int16 1 <> int32 bigLen <> "tiny"+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++    it "mixed valid and truncated columns" $ hedgehog $ do+      nValid <- forAll (Gen.int (Range.linear 1 5))+      validCols <- forAll (Gen.list (Range.singleton nValid) genColumnData)+      -- Append a partial column (length header but no data)+      partialLen <- forAll (Gen.int32 (Range.linear 1 100))+      let payload = int16 (fromIntegral (nValid + 1))+            <> mconcat validCols+            <> int32 partialLen+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++    it "DataRow with random trailing garbage" $ hedgehog $ do+      nCols <- forAll (Gen.int (Range.linear 1 5))+      cols <- forAll (Gen.list (Range.singleton nCols) genColumnData)+      garbage <- forAll (Gen.bytes (Range.linear 0 64))+      let payload = int16 (fromIntegral nCols) <> mconcat cols <> garbage+      ok <- evalIO (doesNotCrash 68 payload)+      assert ok++-- | Evaluate parseCommandTag, fully forcing the result via deepseq.+commandTagDoesNotCrash :: ByteString -> IO Bool+commandTagDoesNotCrash bs = do+  result <- try $ evaluate (parseCommandTag bs `deepseq` ())+  case result of+    Left (_ :: SomeException) -> pure False+    Right () -> pure True++-- Generators ------------------------------------------------------------------++-- | Generate a payload with embedded NUL bytes so recursive parsers+-- (ErrorResponse / NoticeResponse) terminate promptly.+genNulBearingPayload :: Gen ByteString+genNulBearingPayload = do+  -- Build a sequence of (fieldType, value) pairs terminated by NUL+  nFields <- Gen.int (Range.linear 0 10)+  fields <- Gen.list (Range.singleton nFields) $ do+    fieldType <- Gen.word8 Range.linearBounded+    value <- Gen.bytes (Range.linear 0 32)+    -- Ensure value doesn't contain NUL (so NUL terminators are clear)+    let safeValue = BS.map (\b -> if b == 0 then 1 else b) value+    pure $ BS.singleton fieldType <> safeValue <> BS.singleton 0+  pure $ mconcat fields <> BS.singleton 0++-- | Generate a single DataRow column encoding: either NULL (-1) or+-- a length-prefixed bytestring.+genColumnData :: Gen ByteString+genColumnData = Gen.choice+  [ pure (int32 (-1))  -- NULL+  , do+      colBytes <- Gen.bytes (Range.linear 0 128)+      pure (int32 (fromIntegral (BS.length colBytes)) <> colBytes)+  ]++-- | Generate a RowDescription payload with a bounded field count.+genRowDescPayload :: Gen ByteString+genRowDescPayload = do+  nFields <- Gen.int16 (Range.linear 0 5)+  fields <- Gen.list (Range.singleton (fromIntegral nFields)) $ do+    -- name (NUL-terminated) + 18 bytes of field metadata+    name <- Gen.bytes (Range.linear 1 16)+    let safeName = BS.map (\b -> if b == 0 then 65 else b) name+    metadata <- Gen.bytes (Range.singleton 18)+    pure $ safeName <> BS.singleton 0 <> metadata+  pure $ int16 nFields <> mconcat fields++-- | Generate a CopyInResponse / CopyOutResponse payload.+genCopyResponsePayload :: Gen ByteString+genCopyResponsePayload = do+  fmt <- Gen.word8 (Range.linear 0 1)+  nCols <- Gen.int16 (Range.linear 0 10)+  colFmts <- Gen.list (Range.singleton (fromIntegral nCols)) $+    Gen.bytes (Range.singleton 2)+  pure $ BS.singleton fmt <> int16 nCols <> mconcat colFmts++-- Helpers ---------------------------------------------------------------------++int16 :: Int16 -> ByteString+int16 n = LBS.toStrict . B.toLazyByteString $ B.int16BE n++int32 :: Int32 -> ByteString+int32 n = LBS.toStrict . B.toLazyByteString $ B.int32BE n++word32 :: Word32 -> ByteString+word32 n = LBS.toStrict . B.toLazyByteString $ B.word32BE n
+ test/PgWire/Protocol/OidSpec.hs view
@@ -0,0 +1,78 @@+module PgWire.Protocol.OidSpec (spec) where++import PgWire.Protocol.Oid+import Test.Hspec++spec :: Spec+spec = do+  describe "scalar OID constants" $ do+    it "oidBool = 16" $ unOid oidBool `shouldBe` 16+    it "oidBytea = 17" $ unOid oidBytea `shouldBe` 17+    it "oidInt2 = 21" $ unOid oidInt2 `shouldBe` 21+    it "oidInt4 = 23" $ unOid oidInt4 `shouldBe` 23+    it "oidInt8 = 20" $ unOid oidInt8 `shouldBe` 20+    it "oidText = 25" $ unOid oidText `shouldBe` 25+    it "oidJson = 114" $ unOid oidJson `shouldBe` 114+    it "oidFloat4 = 700" $ unOid oidFloat4 `shouldBe` 700+    it "oidFloat8 = 701" $ unOid oidFloat8 `shouldBe` 701+    it "oidVarchar = 1043" $ unOid oidVarchar `shouldBe` 1043+    it "oidDate = 1082" $ unOid oidDate `shouldBe` 1082+    it "oidTime = 1083" $ unOid oidTime `shouldBe` 1083+    it "oidTimestamp = 1114" $ unOid oidTimestamp `shouldBe` 1114+    it "oidTimestamptz = 1184" $ unOid oidTimestamptz `shouldBe` 1184+    it "oidNumeric = 1700" $ unOid oidNumeric `shouldBe` 1700+    it "oidUuid = 2950" $ unOid oidUuid `shouldBe` 2950+    it "oidJsonb = 3802" $ unOid oidJsonb `shouldBe` 3802+    it "oidInterval = 1186" $ unOid oidInterval `shouldBe` 1186++  describe "array OID constants" $ do+    it "oidBoolArray = 1000" $ unOid oidBoolArray `shouldBe` 1000+    it "oidByteaArray = 1001" $ unOid oidByteaArray `shouldBe` 1001+    it "oidInt2Array = 1005" $ unOid oidInt2Array `shouldBe` 1005+    it "oidInt4Array = 1007" $ unOid oidInt4Array `shouldBe` 1007+    it "oidInt8Array = 1016" $ unOid oidInt8Array `shouldBe` 1016+    it "oidTextArray = 1009" $ unOid oidTextArray `shouldBe` 1009+    it "oidVarcharArray = 1015" $ unOid oidVarcharArray `shouldBe` 1015+    it "oidFloat4Array = 1021" $ unOid oidFloat4Array `shouldBe` 1021+    it "oidFloat8Array = 1022" $ unOid oidFloat8Array `shouldBe` 1022+    it "oidTimestampArray = 1115" $ unOid oidTimestampArray `shouldBe` 1115+    it "oidTimestamptzArray = 1185" $ unOid oidTimestamptzArray `shouldBe` 1185+    it "oidDateArray = 1182" $ unOid oidDateArray `shouldBe` 1182+    it "oidTimeArray = 1183" $ unOid oidTimeArray `shouldBe` 1183+    it "oidUuidArray = 2951" $ unOid oidUuidArray `shouldBe` 2951+    it "oidJsonArray = 199" $ unOid oidJsonArray `shouldBe` 199+    it "oidJsonbArray = 3807" $ unOid oidJsonbArray `shouldBe` 3807++  describe "isArrayOid" $ do+    it "True for all 16 array OIDs" $ do+      let arrayOids = [1000, 1001, 1005, 1007, 1009, 1015, 1016,+                        1021, 1022, 1115, 1182, 1183, 1185, 2951, 199, 3807]+      mapM_ (\o -> isArrayOid (Oid o) `shouldBe` True) arrayOids++    it "False for all scalar OIDs" $ do+      let scalarOids = [16, 17, 20, 21, 23, 25, 114, 700, 701,+                         1043, 1082, 1083, 1114, 1184, 1186, 1700, 2950, 3802]+      mapM_ (\o -> isArrayOid (Oid o) `shouldBe` False) scalarOids++    it "False for unknown OID" $+      isArrayOid (Oid 99999) `shouldBe` False++  describe "arrayElementOid" $ do+    it "bool[] → bool" $ arrayElementOid oidBoolArray `shouldBe` Just oidBool+    it "bytea[] → bytea" $ arrayElementOid oidByteaArray `shouldBe` Just oidBytea+    it "int2[] → int2" $ arrayElementOid oidInt2Array `shouldBe` Just oidInt2+    it "int4[] → int4" $ arrayElementOid oidInt4Array `shouldBe` Just oidInt4+    it "int8[] → int8" $ arrayElementOid oidInt8Array `shouldBe` Just oidInt8+    it "text[] → text" $ arrayElementOid oidTextArray `shouldBe` Just oidText+    it "varchar[] → varchar" $ arrayElementOid oidVarcharArray `shouldBe` Just oidVarchar+    it "float4[] → float4" $ arrayElementOid oidFloat4Array `shouldBe` Just oidFloat4+    it "float8[] → float8" $ arrayElementOid oidFloat8Array `shouldBe` Just oidFloat8+    it "timestamp[] → timestamp" $ arrayElementOid oidTimestampArray `shouldBe` Just oidTimestamp+    it "timestamptz[] → timestamptz" $ arrayElementOid oidTimestamptzArray `shouldBe` Just oidTimestamptz+    it "date[] → date" $ arrayElementOid oidDateArray `shouldBe` Just oidDate+    it "time[] → time" $ arrayElementOid oidTimeArray `shouldBe` Just oidTime+    it "uuid[] → uuid" $ arrayElementOid oidUuidArray `shouldBe` Just oidUuid+    it "json[] → json" $ arrayElementOid oidJsonArray `shouldBe` Just oidJson+    it "jsonb[] → jsonb" $ arrayElementOid oidJsonbArray `shouldBe` Just oidJsonb+    it "scalar → Nothing" $ arrayElementOid oidInt4 `shouldBe` Nothing+    it "unknown → Nothing" $ arrayElementOid (Oid 99999) `shouldBe` Nothing
+ test/PgWire/Protocol/ParsersSpec.hs view
@@ -0,0 +1,363 @@+module PgWire.Protocol.ParsersSpec (spec) where++import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int16, Int32)+import Data.Vector qualified as V+import Data.Word (Word32)+import PgWire.Protocol.Backend+import PgWire.Protocol.Parsers (parseBackendMsg, parseCommandTag)+import Test.Hspec++spec :: Spec+spec = do+  -- ── Authentication ──────────────────────────────────────────────+  describe "Authentication" $ do+    it "parses AuthOk" $+      parseBackendMsg 82 (int32 0) `shouldBe` Right (Authentication AuthOk)++    it "parses AuthCleartextPassword" $+      parseBackendMsg 82 (int32 3) `shouldBe` Right (Authentication AuthCleartextPassword)++    it "parses AuthMD5Password with 4-byte salt" $ do+      let salt = BS.pack [0xDE, 0xAD, 0xBE, 0xEF]+      parseBackendMsg 82 (int32 5 <> salt)+        `shouldBe` Right (Authentication (AuthMD5Password salt))++    it "parses AuthSASL with mechanism list" $ do+      let payload = int32 10 <> "SCRAM-SHA-256\0\0"+      case parseBackendMsg 82 payload of+        Right (Authentication (AuthSASL mechs)) ->+          mechs `shouldBe` ["SCRAM-SHA-256"]+        other -> expectationFailure $ "Expected AuthSASL, got: " <> show other++    it "parses AuthSASLContinue" $ do+      let serverData = "r=nonce,s=salt,i=4096"+      parseBackendMsg 82 (int32 11 <> serverData)+        `shouldBe` Right (Authentication (AuthSASLContinue serverData))++    it "parses AuthSASLFinal" $ do+      let serverData = "v=signature"+      parseBackendMsg 82 (int32 12 <> serverData)+        `shouldBe` Right (Authentication (AuthSASLFinal serverData))++    it "rejects unknown auth type" $+      parseBackendMsg 82 (int32 999) `shouldSatisfy` isLeft++  -- ── ReadyForQuery ───────────────────────────────────────────────+  describe "ReadyForQuery" $ do+    it "idle" $ parseBackendMsg 90 (BS.singleton 73) `shouldBe` Right (ReadyForQuery TxIdle)+    it "in transaction" $ parseBackendMsg 90 (BS.singleton 84) `shouldBe` Right (ReadyForQuery TxInTransaction)+    it "failed" $ parseBackendMsg 90 (BS.singleton 69) `shouldBe` Right (ReadyForQuery TxFailed)+    it "rejects unknown status byte" $ parseBackendMsg 90 (BS.singleton 0) `shouldSatisfy` isLeft+    it "rejects empty payload" $ parseBackendMsg 90 BS.empty `shouldSatisfy` isLeft++  -- ── Simple response messages ────────────────────────────────────+  describe "Simple responses" $ do+    it "ParseComplete" $ parseBackendMsg 49 BS.empty `shouldBe` Right ParseComplete+    it "BindComplete" $ parseBackendMsg 50 BS.empty `shouldBe` Right BindComplete+    it "CloseComplete" $ parseBackendMsg 51 BS.empty `shouldBe` Right CloseComplete+    it "NoData" $ parseBackendMsg 110 BS.empty `shouldBe` Right NoData+    it "EmptyQueryResponse" $ parseBackendMsg 73 BS.empty `shouldBe` Right EmptyQueryResponse+    it "CopyDoneMsg" $ parseBackendMsg 99 BS.empty `shouldBe` Right CopyDoneMsg++  -- ── CopyData ────────────────────────────────────────────────────+  describe "CopyDataMsg" $ do+    it "preserves payload" $+      parseBackendMsg 100 "hello world" `shouldBe` Right (CopyDataMsg "hello world")+    it "handles empty payload" $+      parseBackendMsg 100 BS.empty `shouldBe` Right (CopyDataMsg BS.empty)++  -- ── DataRow (hot path) ──────────────────────────────────────────+  describe "DataRow" $ do+    it "parses zero columns" $ do+      let payload = int16 0+      parseBackendMsg 68 payload `shouldBe` Right (DataRow V.empty)++    it "parses single non-null column" $ do+      let payload = int16 1 <> int32 5 <> "hello"+      case parseBackendMsg 68 payload of+        Right (DataRow v) -> do+          V.length v `shouldBe` 1+          v V.! 0 `shouldBe` Just "hello"+        other -> expectationFailure $ show other++    it "parses single NULL column" $ do+      let payload = int16 1 <> int32 (-1)+      case parseBackendMsg 68 payload of+        Right (DataRow v) -> do+          V.length v `shouldBe` 1+          v V.! 0 `shouldBe` Nothing+        other -> expectationFailure $ show other++    it "parses multiple columns with mixed nulls" $ do+      let payload = int16 3+            <> int32 2 <> "ab"    -- col 0: "ab"+            <> int32 (-1)          -- col 1: NULL+            <> int32 4 <> int32 42 -- col 2: 4 bytes (an int32)+      case parseBackendMsg 68 payload of+        Right (DataRow v) -> do+          V.length v `shouldBe` 3+          v V.! 0 `shouldBe` Just "ab"+          v V.! 1 `shouldBe` Nothing+          v V.! 2 `shouldBe` Just (int32 42)+        other -> expectationFailure $ show other++    it "parses empty string column (length 0)" $ do+      let payload = int16 1 <> int32 0+      case parseBackendMsg 68 payload of+        Right (DataRow v) -> do+          V.length v `shouldBe` 1+          v V.! 0 `shouldBe` Just BS.empty+        other -> expectationFailure $ show other++    it "parses many columns" $ do+      let n = 20+          colPayload = int32 1 <> "x"+          payload = int16 n <> BS.concat (replicate (fromIntegral n) colPayload)+      case parseBackendMsg 68 payload of+        Right (DataRow v) -> V.length v `shouldBe` fromIntegral n+        other -> expectationFailure $ show other++    it "fails on truncated column data" $+      parseBackendMsg 68 (int16 1 <> int32 100 <> "short") `shouldSatisfy` isLeft++    it "fails on truncated header" $+      parseBackendMsg 68 (BS.singleton 0) `shouldSatisfy` isLeft++  -- ── RowDescription ──────────────────────────────────────────────+  describe "RowDescription" $ do+    it "parses zero fields" $ do+      let payload = int16 0+      parseBackendMsg 84 payload `shouldBe` Right (RowDescription V.empty)++    it "parses single field" $ do+      let payload = int16 1+            <> "id\0"                -- name (NUL-terminated)+            <> word32 16384          -- table OID+            <> int16 1              -- column number+            <> word32 23            -- type OID (int4)+            <> int16 4              -- type size+            <> int32 (-1)           -- type modifier+            <> int16 0              -- format code (text)+      case parseBackendMsg 84 payload of+        Right (RowDescription v) -> do+          V.length v `shouldBe` 1+          let fi = v V.! 0+          fiName fi `shouldBe` "id"+          fiTableOid fi `shouldBe` 16384+          fiColumnNum fi `shouldBe` 1+          fiTypeOid fi `shouldBe` 23+          fiTypeSize fi `shouldBe` 4+          fiFormatCode fi `shouldBe` 0+        other -> expectationFailure $ show other++    it "parses multiple fields" $ do+      let mkField name oid =+              name <> "\0"+            <> word32 0 <> int16 0 <> word32 oid+            <> int16 (-1) <> int32 (-1) <> int16 1+          payload = int16 3+            <> mkField "id" 23+            <> mkField "name" 25+            <> mkField "email" 25+      case parseBackendMsg 84 payload of+        Right (RowDescription v) -> do+          V.length v `shouldBe` 3+          fiName (v V.! 0) `shouldBe` "id"+          fiName (v V.! 1) `shouldBe` "name"+          fiName (v V.! 2) `shouldBe` "email"+          fiTypeOid (v V.! 0) `shouldBe` 23+          fiTypeOid (v V.! 1) `shouldBe` 25+        other -> expectationFailure $ show other++  -- ── ParameterDescription ────────────────────────────────────────+  describe "ParameterDescription" $ do+    it "parses zero params" $+      parseBackendMsg 116 (int16 0) `shouldBe` Right (ParameterDescription V.empty)++    it "parses multiple param OIDs" $ do+      let payload = int16 2 <> word32 23 <> word32 25+      case parseBackendMsg 116 payload of+        Right (ParameterDescription v) -> do+          V.length v `shouldBe` 2+          v V.! 0 `shouldBe` 23+          v V.! 1 `shouldBe` 25+        other -> expectationFailure $ show other++  -- ── ParameterStatus ─────────────────────────────────────────────+  describe "ParameterStatus" $ do+    it "parses name=value pair" $ do+      let payload = "server_version\0" <> "16.2\0"+      parseBackendMsg 83 payload+        `shouldBe` Right (ParameterStatus "server_version" "16.2")++    it "handles empty value" $ do+      let payload = "param\0\0"+      parseBackendMsg 83 payload+        `shouldBe` Right (ParameterStatus "param" "")++  -- ── BackendKeyData ──────────────────────────────────────────────+  describe "BackendKeyData" $ do+    it "parses PID and key" $ do+      let payload = int32 12345 <> int32 67890+      parseBackendMsg 75 payload `shouldBe` Right (BackendKeyData 12345 67890)++    it "handles zero values" $+      parseBackendMsg 75 (int32 0 <> int32 0) `shouldBe` Right (BackendKeyData 0 0)++  -- ── NotificationResponse ────────────────────────────────────────+  describe "NotificationResponse" $ do+    it "parses PID, channel, payload" $ do+      let payload = int32 42 <> "my_channel\0" <> "hello\0"+      parseBackendMsg 65 payload+        `shouldBe` Right (NotificationResponse 42 "my_channel" "hello")++    it "handles empty payload string" $ do+      let payload = int32 1 <> "ch\0" <> "\0"+      parseBackendMsg 65 payload+        `shouldBe` Right (NotificationResponse 1 "ch" "")++  -- ── CopyInResponse / CopyOutResponse ────────────────────────────+  describe "CopyInResponse" $ do+    it "parses format and column formats" $ do+      let payload = BS.singleton 0 -- overall format: text+            <> int16 2           -- 2 columns+            <> int16 0           -- col 0: text+            <> int16 1           -- col 1: binary+      case parseBackendMsg 71 payload of+        Right (CopyInResponse fmt cols) -> do+          fmt `shouldBe` 0+          V.toList cols `shouldBe` [0, 1]+        other -> expectationFailure $ show other++  describe "CopyOutResponse" $ do+    it "parses format and column formats" $ do+      let payload = BS.singleton 1 <> int16 1 <> int16 1+      case parseBackendMsg 72 payload of+        Right (CopyOutResponse fmt cols) -> do+          fmt `shouldBe` 1+          V.toList cols `shouldBe` [1]+        other -> expectationFailure $ show other++  -- ── CommandComplete ─────────────────────────────────────────────+  describe "CommandComplete" $ do+    it "parses SELECT tag" $ do+      let payload = "SELECT 42\0"+      parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (SelectTag 42))++    it "parses INSERT tag (with OID)" $ do+      let payload = "INSERT 0 5\0"+      parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (InsertTag 5))++    it "parses UPDATE tag" $ do+      let payload = "UPDATE 100\0"+      parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (UpdateTag 100))++    it "parses DELETE tag" $ do+      let payload = "DELETE 0\0"+      parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (DeleteTag 0))++    it "parses CREATE TABLE as OtherTag" $ do+      let payload = "CREATE TABLE\0"+      parseBackendMsg 67 payload `shouldBe` Right (CommandComplete (OtherTag "CREATE TABLE"))++  describe "parseCommandTag" $ do+    it "SELECT 0" $ parseCommandTag "SELECT 0" `shouldBe` SelectTag 0+    it "SELECT 999999" $ parseCommandTag "SELECT 999999" `shouldBe` SelectTag 999999+    it "INSERT 0 1" $ parseCommandTag "INSERT 0 1" `shouldBe` InsertTag 1+    it "UPDATE 1" $ parseCommandTag "UPDATE 1" `shouldBe` UpdateTag 1+    it "DELETE 50" $ parseCommandTag "DELETE 50" `shouldBe` DeleteTag 50+    it "COPY 100" $ parseCommandTag "COPY 100" `shouldBe` OtherTag "COPY 100"+    it "BEGIN" $ parseCommandTag "BEGIN" `shouldBe` OtherTag "BEGIN"++  -- ── ErrorResponse / NoticeResponse ──────────────────────────────+  describe "ErrorResponse" $ do+    it "parses all standard fields" $ do+      let payload = BS.concat+            [ BS.singleton 83, "ERROR\0"              -- S: severity+            , BS.singleton 67, "42P01\0"              -- C: code+            , BS.singleton 77, "relation not found\0" -- M: message+            , BS.singleton 68, "some detail\0"        -- D: detail+            , BS.singleton 72, "try this\0"           -- H: hint+            , BS.singleton 80, "15\0"                 -- P: position+            , BS.singleton 87, "at function foo\0"    -- W: where+            , BS.singleton 115, "public\0"            -- s: schema+            , BS.singleton 116, "users\0"             -- t: table+            , BS.singleton 99, "email\0"              -- c: column+            , BS.singleton 100, "text\0"              -- d: data type+            , BS.singleton 110, "users_pkey\0"        -- n: constraint+            , BS.singleton 70, "parse.c\0"            -- F: file+            , BS.singleton 76, "42\0"                 -- L: line+            , BS.singleton 82, "exec_query\0"         -- R: routine+            , BS.singleton 0                          -- terminator+            ]+      case parseBackendMsg 69 payload of+        Right (ErrorResponse err) -> do+          pgSeverity err `shouldBe` "ERROR"+          pgCode err `shouldBe` "42P01"+          pgMessage err `shouldBe` "relation not found"+          pgDetail err `shouldBe` Just "some detail"+          pgHint err `shouldBe` Just "try this"+          pgPosition err `shouldBe` Just 15+          pgWhere err `shouldBe` Just "at function foo"+          pgSchema err `shouldBe` Just "public"+          pgTable err `shouldBe` Just "users"+          pgColumn err `shouldBe` Just "email"+          pgDataType err `shouldBe` Just "text"+          pgConstraint err `shouldBe` Just "users_pkey"+          pgFile err `shouldBe` Just "parse.c"+          pgLine err `shouldBe` Just 42+          pgRoutine err `shouldBe` Just "exec_query"+        other -> expectationFailure $ show other++    it "handles minimal error (severity + code + message only)" $ do+      let payload = BS.concat+            [ BS.singleton 83, "ERROR\0"+            , BS.singleton 67, "00000\0"+            , BS.singleton 77, "ok\0"+            , BS.singleton 0+            ]+      case parseBackendMsg 69 payload of+        Right (ErrorResponse err) -> do+          pgDetail err `shouldBe` Nothing+          pgHint err `shouldBe` Nothing+          pgPosition err `shouldBe` Nothing+          pgSchema err `shouldBe` Nothing+        other -> expectationFailure $ show other++  describe "NoticeResponse" $ do+    it "parses like ErrorResponse wrapped in PgNotice" $ do+      let payload = BS.concat+            [ BS.singleton 83, "WARNING\0"+            , BS.singleton 67, "01000\0"+            , BS.singleton 77, "be careful\0"+            , BS.singleton 0+            ]+      case parseBackendMsg 78 payload of+        Right (NoticeResponse (PgNotice err)) -> do+          pgSeverity err `shouldBe` "WARNING"+          pgMessage err `shouldBe` "be careful"+        other -> expectationFailure $ show other++  -- ── Unknown tags ────────────────────────────────────────────────+  describe "Unknown tags" $ do+    it "rejects tag 0" $ parseBackendMsg 0 BS.empty `shouldSatisfy` isLeft+    it "rejects tag 255" $ parseBackendMsg 255 BS.empty `shouldSatisfy` isLeft+    it "rejects tag 42" $ parseBackendMsg 42 BS.empty `shouldSatisfy` isLeft++-- Helpers -------------------------------------------------------------------++int16 :: Int16 -> BS.ByteString+int16 n = LBS.toStrict . B.toLazyByteString $ B.int16BE n++int32 :: Int32 -> BS.ByteString+int32 n = LBS.toStrict . B.toLazyByteString $ B.int32BE n++word32 :: Word32 -> BS.ByteString+word32 n = LBS.toStrict . B.toLazyByteString $ B.word32BE n++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False