pqi-native (empty) → 0.0.1.0
raw patch · 16 files changed
+2347/−0 lines, 16 filesdep +basedep +base64-bytestringdep +bytestring
Dependencies added: base, base64-bytestring, bytestring, containers, crypton, hspec, network, pqi, pqi-conformance, pqi-native, ptr-peeker, ptr-poker, transformers
Files
- CHANGELOG.md +0/−0
- LICENSE +22/−0
- README.md +29/−0
- pqi-native.cabal +154/−0
- src/comms/Pqi/Native/Comms.hs +121/−0
- src/library/Pqi/Native.hs +251/−0
- src/library/Pqi/Native/Auth.hs +125/−0
- src/library/Pqi/Native/Connection.hs +347/−0
- src/library/Pqi/Native/LargeObject.hs +131/−0
- src/library/Pqi/Native/Prelude.hs +15/−0
- src/library/Pqi/Native/Query.hs +446/−0
- src/library/Pqi/Native/Types.hs +252/−0
- src/test/Spec.hs +14/−0
- src/transport/Pqi/Native/Transport.hs +119/−0
- src/transport/Pqi/Native/Transport/Message.hs +302/−0
- src/transport/Pqi/Native/Transport/Prelude.hs +19/−0
+ CHANGELOG.md view
+ LICENSE view
@@ -0,0 +1,22 @@+Copyright (c) 2026, Nikita Volkov++Permission is hereby granted, free of charge, to any person+obtaining a copy of this software and associated documentation+files (the "Software"), to deal in the Software without+restriction, including without limitation the rights to use,+copy, modify, merge, publish, distribute, sublicense, and/or sell+copies of the Software, and to permit persons to whom the+Software is furnished to do so, subject to the following+conditions:++The above copyright notice and this permission notice shall be+included in all copies or substantial portions of the Software.++THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES+OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT+HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING+FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR+OTHER DEALINGS IN THE SOFTWARE.
+ README.md view
@@ -0,0 +1,29 @@+# pqi-native++[](https://hackage.haskell.org/package/pqi-native)+[](https://nikita-volkov.github.io/pqi-native/)++A pure-Haskell [`pqi`](https://github.com/nikita-volkov/pqi) adapter+that speaks the PostgreSQL frontend/backend wire protocol directly — no+dependency on the C `libpq` library.++`pqi-native` is an LLM-generated port of the PostgreSQL C client library, [`libpq`](https://www.postgresql.org/docs/current/libpq.html). The upstream [`libpq` source](https://github.com/postgres/postgres/tree/master/src/interfaces/libpq) is the direct reference for the implementation.++## Fidelity goal++The goal is **byte-identical output to `libpq`** for every protocol-derived+value — error message strings, notice text, result status, field metadata,+cell data, and all structured error fields. Fidelity is continuously enforced+by [`pqi-conformance`](https://github.com/nikita-volkov/pqi-conformance), which+runs every operation on both this adapter and a direct+[`postgresql-libpq`](https://hackage.haskell.org/package/postgresql-libpq)+reference connection against the same database and asserts exact equality.++## Status++All classes are implemented. Verified against the `postgresql-libpq` reference+via the conformance differential suite.++Authentication: **trust**, **MD5**, and **SCRAM-SHA-256** are implemented. SCRAM+is verified against a password-auth PostgreSQL 17 container (which defaults to+`scram-sha-256`).
+ pqi-native.cabal view
@@ -0,0 +1,154 @@+cabal-version: 3.0+name: pqi-native+version: 0.0.1.0+category: Database, PostgreSQL+synopsis: Native (pure-Haskell) adapter for pqi+description:+ A @pqi@ adapter that speaks the PostgreSQL frontend\/backend wire+ protocol directly, with no dependency on the C @libpq@ library. The transport+ layer (framing, message serialization\/deserialization, socket I\/O) is+ isolated in an internal @transport@ sub-library built on @ptr-poker@ and+ @ptr-peeker@; the public library is an abstraction over it that implements the+ @pqi@ capability classes.++homepage: https://github.com/nikita-volkov/pqi-native+bug-reports: https://github.com/nikita-volkov/pqi-native/issues+author: Nikita Volkov <nikita.y.volkov@mail.ru>+maintainer: Nikita Volkov <nikita.y.volkov@mail.ru>+copyright: (c) 2026, Nikita Volkov+license: MIT+license-file: LICENSE+extra-doc-files:+ CHANGELOG.md+ LICENSE+ README.md++source-repository head+ type: git+ location: https://github.com/nikita-volkov/pqi-native++common base+ default-language: Haskell2010+ default-extensions:+ ApplicativeDo+ BangPatterns+ BinaryLiterals+ BlockArguments+ ConstraintKinds+ DataKinds+ DefaultSignatures+ DeriveDataTypeable+ DeriveFoldable+ DeriveFunctor+ DeriveTraversable+ DerivingStrategies+ DerivingVia+ DuplicateRecordFields+ EmptyDataDecls+ ExistentialQuantification+ FlexibleContexts+ FlexibleInstances+ FunctionalDependencies+ GADTs+ GeneralizedNewtypeDeriving+ LambdaCase+ LiberalTypeSynonyms+ MagicHash+ MultiParamTypeClasses+ MultiWayIf+ NamedFieldPuns+ NoFieldSelectors+ NoImplicitPrelude+ NoMonomorphismRestriction+ NumericUnderscores+ OverloadedRecordDot+ OverloadedStrings+ ParallelListComp+ PatternGuards+ QuasiQuotes+ RankNTypes+ RecordWildCards+ ScopedTypeVariables+ StandaloneDeriving+ StrictData+ TemplateHaskell+ TupleSections+ TypeApplications+ TypeFamilies+ TypeOperators+ UnboxedTuples+ ViewPatterns++common test+ import: base+ ghc-options:+ -threaded+ -with-rtsopts=-N++library comms+ import: base+ visibility: public+ hs-source-dirs: src/comms+ exposed-modules:+ Pqi.Native.Comms++ build-depends:+ base >=4.11 && <5,+ bytestring >=0.10 && <0.13,+ ptr-peeker ^>=0.2,+ ptr-poker ^>=0.1.3,+ transformers >=0.5 && <0.7,++library transport+ import: base+ visibility: public+ hs-source-dirs: src/transport+ exposed-modules:+ Pqi.Native.Transport+ Pqi.Native.Transport.Message++ other-modules:+ Pqi.Native.Transport.Prelude++ build-depends:+ base >=4.11 && <5,+ bytestring >=0.10 && <0.13,+ network >=3.1 && <3.3,+ pqi-native:comms,+ ptr-poker ^>=0.1.3,+ transformers >=0.5 && <0.7,++library+ import: base+ hs-source-dirs: src/library+ exposed-modules:+ Pqi.Native++ other-modules:+ Pqi.Native.Auth+ Pqi.Native.Connection+ Pqi.Native.LargeObject+ Pqi.Native.Prelude+ Pqi.Native.Query+ Pqi.Native.Types++ build-depends:+ base >=4.11 && <5,+ base64-bytestring >=1.2 && <1.3,+ bytestring >=0.10 && <0.13,+ containers >=0.6 && <0.9,+ crypton >=0.34 && <2,+ pqi ^>=0,+ pqi-native:transport,+ ptr-poker ^>=0.1.3,++test-suite native-test+ import: test+ type: exitcode-stdio-1.0+ hs-source-dirs: src/test+ main-is: Spec.hs+ build-depends:+ base >=4.11 && <5,+ hspec >=2.11 && <2.12,+ pqi-conformance,+ pqi-native,
+ src/comms/Pqi/Native/Comms.hs view
@@ -0,0 +1,121 @@+-- | The serialization abstraction shared by the wire-protocol messages: a type+-- that can be both written (via @ptr-poker@) and decoded (via @ptr-peeker@,+-- with semantic failures carried in an 'ExceptT' layer, since peekers+-- themselves have no failure mechanism).+module Pqi.Native.Comms+ ( Comms (..),+ Decoder,+ DecodingError (..),+ runDecoder,++ -- * Decoder primitives+ liftVariable,+ liftFixed,+ int16,+ int32,+ word8,+ word32,+ cstring,+ bytes,+ remaining,++ -- * Serializable primitives+ CString (..),+ )+where++import Control.Monad.Trans.Class (lift)+import Control.Monad.Trans.Except (ExceptT (..), runExceptT)+import Data.ByteString (ByteString)+import Data.Int (Int16, Int32)+import Data.Word (Word32, Word8)+import qualified PtrPeeker as Peeker+import qualified PtrPoker.Write as Poker+import Prelude++-- | A semantic decoding failure (as opposed to a \"need more bytes\" framing+-- shortfall, which the transport handles separately).+data DecodingError+ = -- | A backend message had an unrecognized type byte.+ UnexpectedMessageType Word8+ | -- | An @Authentication@ message had an unrecognized sub-type.+ UnexpectedAuthType Int32+ | -- | A field could not be interpreted (context, offending tag).+ BadField ByteString+ | -- | The message body was shorter than expected.+ TruncatedInput+ | -- | A general protocol violation.+ ProtocolViolation ByteString+ deriving stock (Eq, Show)++-- | A decoder of a message body: a @ptr-peeker@ 'Peeker.Variable' parser with+-- semantic failures layered on top.+type Decoder = ExceptT DecodingError Peeker.Variable++-- | A type that can be serialized to the wire and decoded back from it.+class Comms a where+ toWrite :: a -> Poker.Write+ decoderOf :: Decoder a++-- | Run a decoder against a complete message body. A framing shortfall (the+-- body being shorter than the decoder consumes) is reported as 'TruncatedInput'.+runDecoder :: Decoder a -> ByteString -> Either DecodingError a+runDecoder decoder body =+ case Peeker.runVariableOnByteString (runExceptT decoder) body of+ Left _needed -> Left TruncatedInput+ Right (Left err) -> Left err+ Right (Right value) -> Right value++liftVariable :: Peeker.Variable a -> Decoder a+liftVariable = lift++liftFixed :: Peeker.Fixed a -> Decoder a+liftFixed = lift . Peeker.fixed++int16 :: Decoder Int16+int16 = liftFixed Peeker.beSignedInt2++int32 :: Decoder Int32+int32 = liftFixed Peeker.beSignedInt4++word8 :: Decoder Word8+word8 = liftFixed Peeker.unsignedInt1++word32 :: Decoder Word32+word32 = liftFixed Peeker.beUnsignedInt4++-- | A null-terminated string.+cstring :: Decoder ByteString+cstring = liftVariable Peeker.nullTerminatedStringAsByteString++-- | Exactly @n@ bytes.+bytes :: Int -> Decoder ByteString+bytes n = liftFixed (Peeker.byteArrayAsByteString n)++-- | All remaining bytes of the body.+remaining :: Decoder ByteString+remaining = liftVariable Peeker.remainderAsByteString++-- | A null-terminated string, as a 'Comms' building block.+newtype CString = CString ByteString+ deriving stock (Eq, Show)++instance Comms CString where+ toWrite (CString value) = Poker.byteString value <> Poker.word8 0+ decoderOf = CString <$> cstring++instance Comms Int16 where+ toWrite = Poker.bInt16+ decoderOf = int16++instance Comms Int32 where+ toWrite = Poker.bInt32+ decoderOf = int32++instance Comms Word32 where+ toWrite = Poker.bWord32+ decoderOf = word32++instance Comms Word8 where+ toWrite = Poker.word8+ decoderOf = word8
+ src/library/Pqi/Native.hs view
@@ -0,0 +1,251 @@+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}++-- | The native (pure-Haskell) @pqi@ adapter.+--+-- 'Connection' speaks the PostgreSQL wire protocol directly.+-- Provides the 'IsConnection' instance.+module Pqi.Native+ ( Connection,+ )+where++import qualified Data.ByteString as ByteString+import qualified Data.ByteString.Char8 as ByteString.Char8+import qualified Data.Map.Strict as Map+import Pqi+import Pqi.Native.Connection (Connection (..))+import qualified Pqi.Native.Connection as Connection+import qualified Pqi.Native.LargeObject as LargeObject+import Pqi.Native.Prelude+import qualified Pqi.Native.Query as Query+import qualified Pqi.Native.Transport as Transport+import Pqi.Native.Transport.Message+ ( BackendMessage (..),+ copyDataMessage,+ copyDoneMessage,+ copyFailMessage,+ flushMessage,+ syncMessage,+ )+import Pqi.Native.Types (NativeCancel (..), NativeResult (..))+import System.Posix.Types (Fd)++instance IsConnection Connection where+ type ResultOf Connection = NativeResult+ type CancelOf Connection = NativeCancel++ connectdb = Connection.establish+ connectStart = Connection.establish+ connectPoll _ = pure PollingOk+ newNullConnection = Connection.nullConnection+ isNullConnection connection = connection.isNull+ finish connection = readIORef connection.transport >>= Transport.close+ reset = Connection.reconnect+ resetStart connection = Connection.reconnect connection $> True+ resetPoll _ = pure PollingOk+ db connection = pure (Just connection.info.database)+ user connection = pure (Just connection.info.user)+ pass connection = pure (Just connection.info.password)+ host connection = pure (Just connection.info.host)+ port connection = pure (Just (ByteString.Char8.pack (show connection.info.port)))+ options _ = pure (Just "")+ status connection = readIORef connection.connStatus+ transactionStatus connection = transactionStatusOf <$> readIORef connection.txStatus+ parameterStatus connection name = Map.lookup name <$> readIORef connection.parameters+ protocolVersion _ = pure 3+ serverVersion connection =+ maybe 0 parseServerVersion . Map.lookup "server_version" <$> readIORef connection.parameters+ errorMessage connection = readIORef connection.lastError+ socket connection = do+ transport <- readIORef connection.transport+ fd <- Transport.socketFd transport+ pure (Just (fromIntegral fd :: Fd))+ backendPID connection = maybe 0 fst <$> readIORef connection.backendKey+ connectionNeedsPassword _ = pure False+ connectionUsedPassword connection = pure (not (ByteString.null connection.info.password))++ exec connection sql = Query.exec connection sql+ execParams connection sql params resultFormat =+ Query.execParams connection sql params resultFormat+ prepare connection name sql parameterTypes =+ Query.prepare connection name sql parameterTypes+ execPrepared connection name params resultFormat =+ Query.execPrepared connection name params resultFormat+ describePrepared connection name = Query.describePrepared connection name+ describePortal connection name = Query.describePortal connection name++ escapeStringConn _ value+ | isValidUtf8 value =+ pure (Just (ByteString.intercalate "''" (ByteString.split 0x27 value)))+ | otherwise = pure Nothing+ escapeByteaConn _ value =+ pure (Just ("\\x" <> hexEncode value))+ escapeIdentifier _ value+ | isValidUtf8 value =+ pure (Just ("\"" <> ByteString.intercalate "\"\"" (ByteString.split 0x22 value) <> "\""))+ | otherwise = pure Nothing++ sendQuery = Query.sendQuery+ sendQueryParams = Query.sendQueryParams+ sendPrepare = Query.sendPrepare+ sendQueryPrepared = Query.sendQueryPrepared+ sendDescribePrepared = Query.sendDescribePrepared+ sendDescribePortal = Query.sendDescribePortal+ getResult connection = Query.getNextResult connection+ consumeInput _ = pure True+ isBusy _ = pure False+ setnonblocking connection flag = writeIORef connection.nonblocking flag $> True+ isnonblocking connection = readIORef connection.nonblocking+ setSingleRowMode connection = do+ pending <- readIORef connection.asyncPending+ if pending+ then writeIORef connection.singleRowMode True $> True+ else pure False+ flush _ = pure FlushOk++ pipelineStatus connection = readIORef connection.pipelineStatus+ enterPipelineMode connection = writeIORef connection.pipelineStatus PipelineOn $> True+ exitPipelineMode connection = do+ pending <- readIORef connection.asyncPending+ if pending+ then pure False+ else writeIORef connection.pipelineStatus PipelineOff $> True+ pipelineSync connection = do+ Connection.sendMessage connection syncMessage+ modifyIORef' connection.pendingSyncs (+ 1)+ writeIORef connection.asyncPending True+ pure True+ sendFlushRequest connection = Connection.sendMessage connection flushMessage $> True++ getCancel connection = do+ key <- readIORef connection.backendKey+ pure+ $ fmap+ ( \(pid, secret) ->+ NativeCancel+ { host = connection.info.host,+ port = connection.info.port,+ pid,+ secret,+ asyncPendingRef = connection.asyncPending,+ pipelineStatusRef = connection.pipelineStatus,+ pendingCommandsRef = connection.pendingCommands+ }+ )+ key++ notifies connection = popFirst connection.pendingNotifications+ disableNoticeReporting connection = writeIORef connection.noticeReporting False+ enableNoticeReporting connection = writeIORef connection.noticeReporting True+ getNotice connection = popFirst connection.notices++ putCopyData connection payload = Connection.sendMessage connection (copyDataMessage payload) $> CopyInOk+ putCopyEnd connection reason = do+ Connection.sendMessage connection (maybe copyDoneMessage copyFailMessage reason)+ writeIORef connection.asyncPending True+ pure CopyInOk+ getCopyData connection _nonBlocking = do+ message <- Connection.nextMessage connection+ case message of+ CopyData payload -> pure (CopyOutRow payload)+ CopyDone -> do+ writeIORef connection.asyncPending True+ pure CopyOutDone+ CommandComplete _ -> drainToReady connection $> CopyOutDone+ ErrorResponse _ -> drainToReady connection $> CopyOutError+ ReadyForQuery txState -> writeIORef connection.txStatus txState $> CopyOutDone+ _ -> getCopyData connection _nonBlocking++ loCreat = LargeObject.loCreat+ loCreate = LargeObject.loCreate+ loImport = LargeObject.loImport+ loImportWithOid = LargeObject.loImportWithOid+ loExport = LargeObject.loExport+ loOpen = LargeObject.loOpen+ loWrite = LargeObject.loWrite+ loRead = LargeObject.loRead+ loSeek = LargeObject.loSeek+ loTell = LargeObject.loTell+ loTruncate = LargeObject.loTruncate+ loClose = LargeObject.loClose+ loUnlink = LargeObject.loUnlink++ clientEncoding connection =+ fromMaybe "SQL_ASCII" . Map.lookup "client_encoding" <$> readIORef connection.parameters+ setClientEncoding connection encoding = do+ result <- Query.exec connection ("SET client_encoding TO '" <> encoding <> "'")+ pure (maybe False (\value -> value.status /= FatalError) result)+ setErrorVerbosity connection verbosity = do+ previous <- readIORef connection.errorVerbosity+ writeIORef connection.errorVerbosity verbosity+ pure previous++-- | Read messages until @ReadyForQuery@, recording the transaction status.+drainToReady :: Connection -> IO ()+drainToReady connection = do+ message <- Connection.nextMessage connection+ case message of+ ReadyForQuery txState -> writeIORef connection.txStatus txState+ _ -> drainToReady connection++-- | Pop the oldest element of a list stored newest-first.+popFirst :: IORef [a] -> IO (Maybe a)+popFirst ref =+ atomicModifyIORef' ref \xs -> case reverse xs of+ [] -> ([], Nothing)+ oldest : rest -> (reverse rest, Just oldest)++transactionStatusOf :: Word8 -> TransactionStatus+transactionStatusOf = \case+ 0x49 -> TransIdle -- 'I'+ 0x54 -> TransInTrans -- 'T'+ 0x45 -> TransInError -- 'E'+ _ -> TransUnknown++-- | Parse the @server_version@ parameter into libpq's @MMmmpp@ integer form+-- (e.g. @\"17.2\"@ -> @170002@, @\"9.6.3\"@ -> @90603@).+parseServerVersion :: ByteString -> Int+parseServerVersion raw =+ case ByteString.Char8.readInt raw of+ Nothing -> 0+ Just (major, rest)+ | major >= 10 -> major * 10000 + nextInt rest+ | otherwise ->+ let minor = nextInt rest+ patch = nextInt (dropInt rest)+ in major * 10000 + minor * 100 + patch+ where+ nextInt bs = case ByteString.Char8.uncons bs of+ Just ('.', remainder) -> maybe 0 fst (ByteString.Char8.readInt remainder)+ _ -> 0+ dropInt bs = case ByteString.Char8.uncons bs of+ Just ('.', remainder) -> case ByteString.Char8.readInt remainder of+ Just (_, leftover) -> leftover+ Nothing -> remainder+ _ -> bs++isValidUtf8 :: ByteString -> Bool+isValidUtf8 = go . ByteString.unpack+ where+ go [] = True+ go (b : bs)+ | b < 0x80 = go bs+ | b < 0xc2 = False+ | b < 0xe0 = cont bs 1+ | b < 0xf0 = cont bs 2+ | b < 0xf5 = cont bs 3+ | otherwise = False+ cont bs (0 :: Int) = go bs+ cont [] _ = False+ cont (b : bs) n+ | b .&. 0xc0 == 0x80 = cont bs (n - 1)+ | otherwise = False++hexEncode :: ByteString -> ByteString+hexEncode = ByteString.Char8.pack . concatMap toHex . ByteString.unpack+ where+ toHex byte = [digit (byte `div` 16), digit (byte `mod` 16)]+ digit n+ | n < 10 = toEnum (fromIntegral n + fromEnum '0')+ | otherwise = toEnum (fromIntegral n - 10 + fromEnum 'a')
+ src/library/Pqi/Native/Auth.hs view
@@ -0,0 +1,125 @@+-- | Authentication helpers: MD5 password hashing and the SASL\/SCRAM-SHA-256+-- exchange.+--+-- The pure crypto lives here; the message round-trip is abstracted as 'SaslStep'+-- so the connection module owns the actual socket I\/O.+module Pqi.Native.Auth+ ( md5Password,+ SaslStep (..),+ SaslMessage (..),+ scram,+ )+where++import Crypto.Hash (Digest, MD5 (..), SHA256 (..), hashWith)+import Crypto.KDF.PBKDF2 (Parameters (..), fastPBKDF2_SHA256)+import Crypto.MAC.HMAC (HMAC, hmac, hmacGetDigest)+import Crypto.Random (getRandomBytes)+import qualified Data.ByteString as ByteString+import qualified Data.ByteString.Base64 as Base64+import qualified Data.ByteString.Char8 as ByteString.Char8+import qualified Data.List as List+import Pqi.Native.Prelude++-- | Compute the response to an @AuthenticationMD5Password@ challenge:+-- @"md5" <> md5hex (md5hex (password <> user) <> salt)@.+md5Password :: ByteString -> ByteString -> ByteString -> ByteString+md5Password user password salt =+ "md5" <> md5Hex (md5Hex (password <> user) <> salt)++md5Hex :: ByteString -> ByteString+md5Hex = ByteString.Char8.pack . show . hashWith MD5++-- | One server SASL\/authentication message, as the SCRAM logic sees it.+data SaslMessage+ = SaslContinue ByteString+ | SaslFinal ByteString+ | SaslOk+ | SaslError ByteString++-- | The message round-trip the SCRAM exchange drives, supplied by the+-- connection module.+data SaslStep = SaslStep+ { sendInitial :: ByteString -> ByteString -> IO (),+ sendResponse :: ByteString -> IO (),+ receive :: IO SaslMessage+ }++mechanismName :: ByteString+mechanismName = "SCRAM-SHA-256"++-- | Run the SCRAM-SHA-256 exchange (without channel binding). Returns @Right ()@+-- once the server accepts the client proof, or @Left@ with a problem+-- description.+scram :: ByteString -> ByteString -> [ByteString] -> SaslStep -> IO (Either ByteString ())+scram _user password mechanisms step+ | mechanismName `notElem` mechanisms =+ pure (Left "server did not offer SCRAM-SHA-256")+ | otherwise = do+ clientNonce <- Base64.encode <$> getRandomBytes 18+ let clientFirstBare = "n=,r=" <> clientNonce+ step.sendInitial mechanismName ("n,," <> clientFirstBare)+ step.receive >>= \case+ SaslError problem -> pure (Left problem)+ SaslContinue serverFirst ->+ case parseServerFirst serverFirst of+ Nothing -> pure (Left "malformed SCRAM server-first message")+ Just (serverNonce, salt, iterations) -> do+ let saltedPassword = fastPBKDF2_SHA256 (Parameters iterations 32) password salt :: ByteString+ clientKey = hmacSha256 saltedPassword "Client Key"+ storedKey = sha256 clientKey+ clientFinalWithoutProof = "c=biws,r=" <> serverNonce+ authMessage =+ ByteString.intercalate "," [clientFirstBare, serverFirst, clientFinalWithoutProof]+ clientSignature = hmacSha256 storedKey authMessage+ clientProof = xorBytes clientKey clientSignature+ clientFinal = clientFinalWithoutProof <> ",p=" <> Base64.encode clientProof+ step.sendResponse clientFinal+ step.receive >>= \case+ SaslFinal _ -> pure (Right ())+ SaslOk -> pure (Right ())+ SaslError problem -> pure (Left problem)+ _ -> pure (Left "unexpected SCRAM server-final message")+ _ -> pure (Left "unexpected SCRAM message")++-- | Parse @r=<nonce>,s=<salt base64>,i=<iterations>@ into the server nonce, the+-- decoded salt, and the iteration count.+parseServerFirst :: ByteString -> Maybe (ByteString, ByteString, Int)+parseServerFirst message = do+ let attributes = ByteString.Char8.split ',' message+ nonce <- attributeValue "r=" attributes+ saltEncoded <- attributeValue "s=" attributes+ salt <- either (const Nothing) Just (Base64.decode saltEncoded)+ iterationsText <- attributeValue "i=" attributes+ (iterations, _) <- ByteString.Char8.readInt iterationsText+ pure (nonce, salt, iterations)+ where+ attributeValue prefix =+ fmap (ByteString.drop (ByteString.length prefix))+ . List.find (ByteString.isPrefixOf prefix)++hmacSha256 :: ByteString -> ByteString -> ByteString+hmacSha256 key message = digestBytes (hmacGetDigest (hmac key message :: HMAC SHA256))++sha256 :: ByteString -> ByteString+sha256 = digestBytes . hashWith SHA256++-- | Extract the raw bytes of a digest via its hexadecimal 'Show' instance,+-- avoiding a direct @memory@ dependency (whose @ByteArrayAccess@ instance for+-- crypton's @Digest@ conflicts across package versions).+digestBytes :: Digest a -> ByteString+digestBytes = hexToBytes . ByteString.Char8.pack . show++hexToBytes :: ByteString -> ByteString+hexToBytes = ByteString.pack . pairs . ByteString.unpack+ where+ pairs (hi : lo : rest) = (hexValue hi * 16 + hexValue lo) : pairs rest+ pairs _ = []+ hexValue w+ | w >= 0x30 && w <= 0x39 = w - 0x30+ | w >= 0x61 && w <= 0x66 = w - 0x57+ | w >= 0x41 && w <= 0x46 = w - 0x37+ | otherwise = 0++xorBytes :: ByteString -> ByteString -> ByteString+xorBytes a b = ByteString.pack (ByteString.zipWith xor a b)
+ src/library/Pqi/Native/Connection.hs view
@@ -0,0 +1,347 @@+-- | The native connection: its mutable state, conninfo parsing, the+-- startup\/authentication handshake, and the interleave-aware receive loop that+-- the higher-level query code is built on.+module Pqi.Native.Connection+ ( Connection (..),+ ConnInfo (..),+ parseConnInfo,+ establish,+ nullConnection,+ reconnect,+ nextMessage,+ sendMessage,+ fieldValue,+ setError,+ )+where++import Control.Exception (IOException, SomeException, catch, try)+import qualified Data.ByteString as ByteString+import qualified Data.ByteString.Char8 as ByteString.Char8+import qualified Data.Map.Strict as Map+import Pqi (ConnStatus (..), Notify (..), PipelineStatus (..), Verbosity (..))+import qualified Pqi.Native.Auth as Auth+import Pqi.Native.Prelude+import Pqi.Native.Transport (Transport)+import qualified Pqi.Native.Transport as Transport+import Pqi.Native.Transport.Message+import Pqi.Native.Types (formatErrorFields)+import qualified PtrPoker.Write as Poker++-- | Parsed connection parameters (the @key=value@ subset we support).+data ConnInfo = ConnInfo+ { host :: ByteString,+ port :: Int,+ user :: ByteString,+ database :: ByteString,+ password :: ByteString+ }+ deriving stock (Eq, Show)++-- | Parse a conninfo string in either @key=value@ or @postgresql:\/\/@ URI+-- format. Unquoted key=value values only; URI values are percent-decoded.+-- Environment variables and @.pgpass@ are not supported.+parseConnInfo :: ByteString -> ConnInfo+parseConnInfo raw+ | "postgresql://" `ByteString.isPrefixOf` raw = parseUri (ByteString.drop 13 raw)+ | "postgres://" `ByteString.isPrefixOf` raw = parseUri (ByteString.drop 11 raw)+ | otherwise = parseKeyValue raw++parseKeyValue :: ByteString -> ConnInfo+parseKeyValue raw =+ ConnInfo+ { host = get "host" "localhost",+ port = maybe 5432 fst (ByteString.Char8.readInt (get "port" "5432")),+ user = theUser,+ database = get "dbname" theUser,+ password = get "password" ""+ }+ where+ pairs = mapMaybe toPair (ByteString.Char8.words raw)+ settings = Map.fromList pairs+ get key def = Map.findWithDefault def key settings+ theUser = get "user" "postgres"+ toPair token = case ByteString.Char8.break (== '=') token of+ (key, value)+ | not (ByteString.null value) -> Just (key, ByteString.drop 1 value)+ _ -> Nothing++-- | Parse the authority+path portion of a @postgresql://@ URI (scheme already+-- stripped). Handles @[user[:password]@][host[:port]][/dbname]@; ignores query+-- parameters other than what appears in those components.+parseUri :: ByteString -> ConnInfo+parseUri withoutScheme =+ ConnInfo {host, port, user, database, password}+ where+ -- Split off optional "userinfo@" prefix. The '@' is unambiguous in this+ -- position: hosts do not contain '@' in practice.+ (userinfoMay, afterAt) = case ByteString.elemIndex 0x40 withoutScheme of+ Just i -> (Just (ByteString.take i withoutScheme), ByteString.drop (i + 1) withoutScheme)+ Nothing -> (Nothing, withoutScheme)++ (user, password) = case userinfoMay of+ Nothing -> ("postgres", "")+ Just ui -> case ByteString.elemIndex 0x3a ui of+ Just c -> (pctDecode (ByteString.take c ui), pctDecode (ByteString.drop (c + 1) ui))+ Nothing -> (pctDecode ui, "")++ -- Split host[:port] from /dbname?params.+ (hostport, pathAndQuery) = ByteString.break (== 0x2f) afterAt++ -- Strip optional leading '/' to get just the dbname (ignoring ?query).+ rawDbname =+ let p = ByteString.drop 1 pathAndQuery+ in ByteString.takeWhile (/= 0x3f) p++ database =+ if ByteString.null rawDbname+ then user+ else pctDecode rawDbname++ -- Parse host and port from "host:port", handling IPv6 "[::1]:port".+ (host, port)+ | not (ByteString.null hostport) && ByteString.head hostport == 0x5b =+ case ByteString.elemIndex 0x5d hostport of+ Just close ->+ let h = ByteString.take close (ByteString.drop 1 hostport)+ portStr = ByteString.drop 2 (ByteString.drop close hostport)+ in (h, readPort portStr)+ Nothing -> (hostport, 5432)+ | otherwise = case ByteString.elemIndexEnd 0x3a hostport of+ Just c ->+ let h = ByteString.take c hostport+ p = ByteString.drop (c + 1) hostport+ in if ByteString.null h then ("localhost", readPort p) else (h, readPort p)+ Nothing ->+ (if ByteString.null hostport then "localhost" else hostport, 5432)++ readPort bs = maybe 5432 fst (ByteString.Char8.readInt bs)++-- | Decode @%XX@ percent-encoding in a URI component.+pctDecode :: ByteString -> ByteString+pctDecode bs+ | ByteString.null bs = ByteString.empty+ | ByteString.head bs == 0x25,+ ByteString.length bs >= 3,+ Just hi <- hexVal (ByteString.index bs 1),+ Just lo <- hexVal (ByteString.index bs 2) =+ ByteString.singleton (fromIntegral (hi * 16 + lo)) <> pctDecode (ByteString.drop 3 bs)+ | otherwise = ByteString.singleton (ByteString.head bs) <> pctDecode (ByteString.tail bs)+ where+ hexVal w+ | w >= 0x30 && w <= 0x39 = Just (fromIntegral w - 0x30 :: Int)+ | w >= 0x41 && w <= 0x46 = Just (fromIntegral w - 0x41 + 10 :: Int)+ | w >= 0x61 && w <= 0x66 = Just (fromIntegral w - 0x61 + 10 :: Int)+ | otherwise = Nothing++-- | A native connection and its mutable state.+data Connection = Connection+ { transport :: IORef Transport,+ info :: ConnInfo,+ isNull :: Bool,+ parameters :: IORef (Map.Map ByteString ByteString),+ backendKey :: IORef (Maybe (Int32, Int32)),+ txStatus :: IORef Word8,+ connStatus :: IORef ConnStatus,+ lastError :: IORef (Maybe ByteString),+ notices :: IORef [ByteString],+ pendingNotifications :: IORef [Notify],+ noticeReporting :: IORef Bool,+ asyncPending :: IORef Bool,+ nonblocking :: IORef Bool,+ pipelineStatus :: IORef PipelineStatus,+ singleRowMode :: IORef Bool,+ singleRowFields :: IORef [FieldDescription],+ pipelineSeparatorPending :: IORef Bool,+ pendingSyncs :: IORef Int,+ pendingCommands :: IORef Int,+ -- | Count of standalone @Parse@ messages sent in pipeline mode (each from+ -- 'sendPrepare'), for which only @ParseComplete@ (no @CommandComplete@) is+ -- expected as a terminal response.+ pendingParses :: IORef Int,+ errorVerbosity :: IORef Verbosity,+ -- | The SQL text of the most recently sent query (set by sendQuery /+ -- sendQueryParams). Used when formatting error messages for async results+ -- so that @LINE N:@ position context can be reproduced.+ currentQuery :: IORef ByteString+ }++-- | Send a serialized frontend message.+sendMessage :: Connection -> Poker.Write -> IO ()+sendMessage connection write = do+ transport <- readIORef connection.transport+ Transport.send transport write++-- | Receive the next /protocol-relevant/ backend message, transparently+-- consuming and recording the asynchronous messages the backend may interleave+-- at any time: @ParameterStatus@ (updates the parameter map), @NoticeResponse@+-- (collected when notice reporting is on), and @NotificationResponse@ (queued).+nextMessage :: Connection -> IO BackendMessage+nextMessage connection = do+ transport <- readIORef connection.transport+ (typeByte, body) <- Transport.receiveFrame transport+ case decodeBackendMessage typeByte body of+ Left err -> ioError (userError ("pqi-native: protocol decode error: " <> show err))+ Right message -> case message of+ ParameterStatus key value -> do+ modifyIORef' connection.parameters (Map.insert key value)+ nextMessage connection+ NoticeResponse fields -> do+ reporting <- readIORef connection.noticeReporting+ when reporting $ do+ let noticeText = formatErrorFields (Map.fromList fields)+ unless (ByteString.null noticeText)+ $ modifyIORef' connection.notices (noticeText :)+ nextMessage connection+ NotificationResponse pid channel payload -> do+ modifyIORef' connection.pendingNotifications (Notify channel pid payload :)+ nextMessage connection+ other -> pure other++-- | Look up an error\/notice field by its single-byte code.+fieldValue :: Word8 -> [(Word8, ByteString)] -> Maybe ByteString+fieldValue code = lookup code++-- | Record a flat error message and mark the connection bad.+setError :: Connection -> ByteString -> IO ()+setError connection message = do+ writeIORef connection.lastError (Just message)+ writeIORef connection.connStatus ConnectionBad++-- | Open a connection: resolve and connect the socket, send the startup+-- message, and run the authentication\/startup handshake. Like libpq, a failed+-- connection (whether due to a network error or a rejected handshake) yields a+-- 'ConnectionBad' connection rather than throwing.+establish :: ByteString -> IO Connection+establish conninfo = do+ let info = parseConnInfo conninfo+ transportResult <- try @IOException (Transport.connect info.host info.port)+ case transportResult of+ Left err -> do+ transport <- Transport.unconnected+ connection <- newConnection False transport info+ setError connection ("could not connect to server: " <> ByteString.Char8.pack (show err))+ pure connection+ Right transport -> do+ connection <- newConnection False transport info+ sendMessage connection (startupMessage [("user", info.user), ("database", info.database)])+ handshake connection+ pure connection++-- | A \"null\" sentinel connection (the analogue of @PQnewNullConnection@): no+-- live socket, permanently in the 'ConnectionBad' state.+nullConnection :: IO Connection+nullConnection = do+ transport <- Transport.unconnected+ conn <- newConnection True transport (parseConnInfo "")+ writeIORef conn.lastError (Just "connection pointer is NULL\n")+ pure conn++-- | Close the current socket and run the startup handshake again on a fresh+-- one, reusing the stored conninfo (the analogue of @PQreset@).+reconnect :: Connection -> IO ()+reconnect connection = do+ oldTransport <- readIORef connection.transport+ Transport.close oldTransport+ newTransport <- Transport.connect connection.info.host connection.info.port+ writeIORef connection.transport newTransport+ writeIORef connection.parameters Map.empty+ writeIORef connection.backendKey Nothing+ writeIORef connection.txStatus 0x49+ writeIORef connection.connStatus ConnectionBad+ writeIORef connection.lastError (Just "")+ sendMessage connection (startupMessage [("user", connection.info.user), ("database", connection.info.database)])+ handshake connection++newConnection :: Bool -> Transport -> ConnInfo -> IO Connection+newConnection isNull transport info = do+ transportRef <- newIORef transport+ Connection transportRef info isNull+ <$> newIORef Map.empty+ <*> newIORef Nothing+ <*> newIORef 0x49 -- 'I'+ <*> newIORef ConnectionBad+ <*> newIORef (Just "")+ <*> newIORef []+ <*> newIORef []+ <*> newIORef False+ <*> newIORef False+ <*> newIORef False+ <*> newIORef PipelineOff+ <*> newIORef False+ <*> newIORef []+ <*> newIORef False+ <*> newIORef 0+ <*> newIORef 0+ <*> newIORef 0+ <*> newIORef ErrorsDefault+ <*> newIORef ""++-- | The startup\/authentication state machine, ending at the first+-- @ReadyForQuery@ (success) or @ErrorResponse@ (failure).+handshake :: Connection -> IO ()+handshake connection = authenticating+ where+ authenticating = do+ message <- nextMessage connection+ case message of+ AuthenticationOk -> startingUp+ AuthenticationCleartextPassword -> do+ sendMessage connection (passwordMessage connection.info.password)+ authenticating+ AuthenticationMD5Password salt -> do+ let response = Auth.md5Password connection.info.user connection.info.password salt+ sendMessage connection (passwordMessage response)+ authenticating+ AuthenticationSASL mechanisms ->+ Auth.scram connection.info.user connection.info.password mechanisms (saslExchange connection) >>= \case+ Left problem -> setError connection problem+ Right () -> startingUp+ ErrorResponse fields -> failWith fields+ other -> failWith [(0x4d, "unexpected authentication message: " <> ByteString.Char8.pack (show other))]+ startingUp = do+ message <- nextMessage connection+ case message of+ BackendKeyData pid secret -> do+ writeIORef connection.backendKey (Just (pid, secret))+ startingUp+ ReadyForQuery txState -> do+ writeIORef connection.txStatus txState+ writeIORef connection.connStatus ConnectionOk+ ErrorResponse fields -> failWith fields+ _ -> startingUp+ failWith fields = do+ let fmtFields = formatErrorFields (Map.fromList fields)+ transport <- readIORef connection.transport+ mIp <- catch (Just <$> Transport.peerIp transport) (\(_ :: SomeException) -> pure Nothing)+ setError connection $ case mIp of+ Nothing -> fmtFields+ Just ip ->+ "connection to server at \""+ <> connection.info.host+ <> "\" ("+ <> ip+ <> "), port "+ <> ByteString.Char8.pack (show connection.info.port)+ <> " failed: "+ <> fmtFields++-- | The SASL message round-trip used by 'Auth.scram': send a client message and+-- receive the next server SASL\/auth message, projected to the bytes the SCRAM+-- logic needs.+saslExchange :: Connection -> Auth.SaslStep+saslExchange connection =+ Auth.SaslStep+ { Auth.sendInitial = \mechanism initial ->+ sendMessage connection (saslInitialResponse mechanism initial),+ Auth.sendResponse = \payload ->+ sendMessage connection (saslResponse payload),+ Auth.receive = do+ message <- nextMessage connection+ pure $ case message of+ AuthenticationSASLContinue payload -> Auth.SaslContinue payload+ AuthenticationSASLFinal payload -> Auth.SaslFinal payload+ AuthenticationOk -> Auth.SaslOk+ ErrorResponse fields -> Auth.SaslError (fromMaybe "SASL error" (fieldValue 0x4d fields))+ other -> Auth.SaslError ("unexpected SASL message: " <> ByteString.Char8.pack (show other))+ }
+ src/library/Pqi/Native/LargeObject.hs view
@@ -0,0 +1,131 @@+-- | The large-object interface, implemented over the server's @lo_*@ SQL+-- functions (rather than libpq's fast-path protocol) — identical results, far+-- less machinery. As with libpq, the open\/read\/write\/close operations must be+-- run inside a transaction managed by the caller.+module Pqi.Native.LargeObject+ ( loCreat,+ loCreate,+ loImport,+ loImportWithOid,+ loExport,+ loOpen,+ loWrite,+ loRead,+ loSeek,+ loTell,+ loTruncate,+ loClose,+ loUnlink,+ )+where++import qualified Data.ByteString.Char8 as ByteString.Char8+import Pqi (Format (..))+import Pqi.Native.Connection (Connection)+import Pqi.Native.Prelude+import qualified Pqi.Native.Query as Query+import Pqi.Native.Types (NativeResult (..))+import System.IO (IOMode (..), SeekMode (..))++loCreat :: Connection -> IO (Maybe Word32)+loCreat connection =+ (>>= parseOid) <$> callText connection "select lo_creat($1 :: integer)" [intParam (-1)]++loCreate :: Connection -> Word32 -> IO (Maybe Word32)+loCreate connection oid =+ (>>= parseOid) <$> callText connection "select lo_create($1 :: oid)" [oidParam oid]++loImport :: Connection -> FilePath -> IO (Maybe Word32)+loImport connection path =+ (>>= parseOid) <$> callText connection "select lo_import($1)" [textParam (ByteString.Char8.pack path)]++loImportWithOid :: Connection -> FilePath -> Word32 -> IO (Maybe Word32)+loImportWithOid connection path oid =+ (>>= parseOid)+ <$> callText connection "select lo_import($1, $2 :: oid)" [textParam (ByteString.Char8.pack path), oidParam oid]++loExport :: Connection -> Word32 -> FilePath -> IO (Maybe ())+loExport connection oid path =+ succeeded <$> callText connection "select lo_export($1 :: oid, $2)" [oidParam oid, textParam (ByteString.Char8.pack path)]++loOpen :: Connection -> Word32 -> IOMode -> IO (Maybe Int32)+loOpen connection oid mode =+ fmap fromIntegral+ . (>>= parseInt)+ <$> callText connection "select lo_open($1 :: oid, $2 :: integer)" [oidParam oid, intParam (ioModeFlag mode)]++loWrite :: Connection -> Int32 -> ByteString -> IO (Maybe Int)+loWrite connection fd payload =+ (>>= parseInt) <$> callText connection "select lowrite($1 :: integer, $2 :: bytea)" [intParam (fromIntegral fd), byteaParam payload]++loRead :: Connection -> Int32 -> Int -> IO (Maybe ByteString)+loRead connection fd len =+ callBinary connection "select loread($1 :: integer, $2 :: integer)" [intParam (fromIntegral fd), intParam len]++loSeek :: Connection -> Int32 -> SeekMode -> Int -> IO (Maybe Int)+loSeek connection fd whence offset =+ (>>= parseInt)+ <$> callText connection "select lo_lseek($1 :: integer, $2 :: integer, $3 :: integer)" [intParam (fromIntegral fd), intParam offset, intParam (seekFlag whence)]++loTell :: Connection -> Int32 -> IO (Maybe Int)+loTell connection fd =+ (>>= parseInt) <$> callText connection "select lo_tell($1 :: integer)" [intParam (fromIntegral fd)]++loTruncate :: Connection -> Int32 -> Int -> IO (Maybe ())+loTruncate connection fd len =+ succeeded <$> callText connection "select lo_truncate($1 :: integer, $2 :: integer)" [intParam (fromIntegral fd), intParam len]++loClose :: Connection -> Int32 -> IO (Maybe ())+loClose connection fd =+ succeeded <$> callText connection "select lo_close($1 :: integer)" [intParam (fromIntegral fd)]++loUnlink :: Connection -> Word32 -> IO (Maybe ())+loUnlink connection oid =+ succeeded <$> callText connection "select lo_unlink($1 :: oid)" [oidParam oid]++-- * Helpers++callText :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> IO (Maybe ByteString)+callText connection sql params = (>>= firstValue) <$> Query.execParams connection sql params Text++callBinary :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> IO (Maybe ByteString)+callBinary connection sql params = (>>= firstValue) <$> Query.execParams connection sql params Binary++firstValue :: NativeResult -> Maybe ByteString+firstValue result = case result.rows of+ (cell : _) : _ -> cell+ _ -> Nothing++succeeded :: Maybe ByteString -> Maybe ()+succeeded = fmap (const ())++intParam :: Int -> Maybe (Word32, ByteString, Format)+intParam value = Just (0, ByteString.Char8.pack (show value), Text)++oidParam :: Word32 -> Maybe (Word32, ByteString, Format)+oidParam value = Just (0, ByteString.Char8.pack (show value), Text)++textParam :: ByteString -> Maybe (Word32, ByteString, Format)+textParam value = Just (0, value, Text)++byteaParam :: ByteString -> Maybe (Word32, ByteString, Format)+byteaParam value = Just (17, value, Binary)++parseInt :: ByteString -> Maybe Int+parseInt = fmap fst . ByteString.Char8.readInt++parseOid :: ByteString -> Maybe Word32+parseOid = fmap (fromIntegral . fst) . ByteString.Char8.readInt++ioModeFlag :: IOMode -> Int+ioModeFlag = \case+ ReadMode -> 0x40000 -- INV_READ+ WriteMode -> 0x20000 -- INV_WRITE+ AppendMode -> 0x20000+ ReadWriteMode -> 0x60000 -- INV_READ | INV_WRITE++seekFlag :: SeekMode -> Int+seekFlag = \case+ AbsoluteSeek -> 0+ RelativeSeek -> 1+ SeekFromEnd -> 2
+ src/library/Pqi/Native/Prelude.hs view
@@ -0,0 +1,15 @@+module Pqi.Native.Prelude+ ( module Exports,+ )+where++import Control.Applicative as Exports+import Control.Monad as Exports+import Data.Bits as Exports+import Data.ByteString as Exports (ByteString)+import Data.Functor as Exports (($>), (<&>))+import Data.IORef as Exports+import Data.Int as Exports+import Data.Maybe as Exports+import Data.Word as Exports+import Prelude as Exports
+ src/library/Pqi/Native/Query.hs view
@@ -0,0 +1,446 @@+-- | Command execution: the simple- and extended-query flows, and the+-- materialization of the backend message stream into a 'NativeResult'.+module Pqi.Native.Query+ ( exec,+ execParams,+ prepare,+ execPrepared,+ describePrepared,+ describePortal,+ sendQuery,+ sendQueryParams,+ sendPrepare,+ sendQueryPrepared,+ sendDescribePrepared,+ sendDescribePortal,+ getNextResult,+ )+where++import qualified Data.Map.Strict as Map+import Pqi (ConnStatus (..), ExecStatus (..), Format (..), PipelineStatus (..))+import Pqi.Native.Connection+import Pqi.Native.Prelude+import Pqi.Native.Transport.Message+import Pqi.Native.Types (NativeResult (..), formatResultError)+import qualified PtrPoker.Write as Poker++-- * Message construction++-- Sync-inclusive variants used by the synchronous exec* functions.++paramsWrite :: ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> Poker.Write+paramsWrite sql params resultFormat =+ asyncParamsWrite sql params resultFormat <> syncMessage++preparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write+preparedWrite name params resultFormat =+ asyncPreparedWrite name params resultFormat <> syncMessage++prepareWrite :: ByteString -> ByteString -> Maybe [Word32] -> Poker.Write+prepareWrite name sql parameterTypes =+ parseMessage name sql (fromMaybe [] parameterTypes) <> syncMessage++-- Sync-free variants used by the async send* functions.+-- In non-pipeline mode sendAsync appends syncMessage; in pipeline mode it does not.++asyncParamsWrite :: ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> Poker.Write+asyncParamsWrite sql params resultFormat =+ parseMessage "" sql (fmap paramOid params)+ <> bindMessage "" "" (fmap paramFormat params) (fmap paramValue params) [formatCode resultFormat]+ <> describePortalMessage ""+ <> executeMessage "" 0++asyncPreparedWrite :: ByteString -> [Maybe (ByteString, Format)] -> Format -> Poker.Write+asyncPreparedWrite name params resultFormat =+ bindMessage "" name (fmap boundFormat params) (fmap boundValue params) [formatCode resultFormat]+ <> describePortalMessage ""+ <> executeMessage "" 0++-- * Synchronous flows++-- | Simple query. Returns the last result, mirroring @PQexec@.+exec :: Connection -> ByteString -> IO (Maybe NativeResult)+exec connection sql = withReady connection do+ sendMessage connection (queryMessage sql)+ lastMaybe <$> collectSimple connection sql++-- | Parameterized query via the extended protocol.+execParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO (Maybe NativeResult)+execParams connection sql params resultFormat = withReady connection do+ sendMessage connection (paramsWrite sql params resultFormat)+ Just <$> collectExtended connection sql++-- | Prepare a named statement.+prepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO (Maybe NativeResult)+prepare connection name sql parameterTypes = withReady connection do+ sendMessage connection (prepareWrite name sql parameterTypes)+ Just <$> collectExtended connection sql++-- | Execute a previously prepared statement.+execPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO (Maybe NativeResult)+execPrepared connection name params resultFormat = withReady connection do+ sendMessage connection (preparedWrite name params resultFormat)+ Just <$> collectExtended connection ""++-- * Asynchronous flows++-- | Send a write in async mode, tracking pending commands for pipeline abort.+sendAsync :: Connection -> ByteString -> Poker.Write -> IO Bool+sendAsync connection sql write = do+ status <- readIORef connection.connStatus+ case status of+ ConnectionOk -> do+ sendMessage connection write+ writeIORef connection.currentQuery sql+ writeIORef connection.asyncPending True+ pipeStatus <- readIORef connection.pipelineStatus+ when (pipeStatus /= PipelineOff) $ modifyIORef' connection.pendingCommands (+ 1)+ pure True+ _ -> pure False++-- | Whether the connection is in pipeline mode.+inPipeline :: Connection -> IO Bool+inPipeline connection = (/= PipelineOff) <$> readIORef connection.pipelineStatus++-- Simple query protocol: no Sync needed (server sends ReadyForQuery on its own).+sendQuery :: Connection -> ByteString -> IO Bool+sendQuery connection sql = sendAsync connection sql (queryMessage sql)++-- Extended query: include Sync when not in pipeline mode; omit Sync in+-- pipeline mode (the caller drives sync boundaries via 'pipelineSync').+sendQueryParams :: Connection -> ByteString -> [Maybe (Word32, ByteString, Format)] -> Format -> IO Bool+sendQueryParams connection sql params resultFormat = do+ pipeline <- inPipeline connection+ sendAsync connection sql+ $ if pipeline+ then asyncParamsWrite sql params resultFormat+ else paramsWrite sql params resultFormat++sendPrepare :: Connection -> ByteString -> ByteString -> Maybe [Word32] -> IO Bool+sendPrepare connection name sql parameterTypes = do+ pipeline <- inPipeline connection+ ok <-+ sendAsync connection sql+ $ if pipeline+ then parseMessage name sql (fromMaybe [] parameterTypes)+ else prepareWrite name sql parameterTypes+ when (ok && pipeline) $ modifyIORef' connection.pendingParses (+ 1)+ pure ok++sendQueryPrepared :: Connection -> ByteString -> [Maybe (ByteString, Format)] -> Format -> IO Bool+sendQueryPrepared connection name params resultFormat = do+ pipeline <- inPipeline connection+ sendAsync connection ""+ $ if pipeline+ then asyncPreparedWrite name params resultFormat+ else preparedWrite name params resultFormat++sendDescribePrepared :: Connection -> ByteString -> IO Bool+sendDescribePrepared connection name = do+ pipeline <- inPipeline connection+ sendAsync connection ""+ $ if pipeline+ then describeStatementMessage name+ else describeStatementMessage name <> syncMessage++sendDescribePortal :: Connection -> ByteString -> IO Bool+sendDescribePortal connection name = do+ pipeline <- inPipeline connection+ sendAsync connection ""+ $ if pipeline+ then describePortalMessage name+ else describePortalMessage name <> syncMessage++-- | Read the next result of an in-flight asynchronous command, or 'Nothing'+-- once @ReadyForQuery@ is reached (clearing the pending flag), mirroring+-- @PQgetResult@.+--+-- In pipeline mode a separator 'Nothing' is returned between each command's+-- result set, and a 'PipelineSync' result is returned for each @Sync@+-- boundary. In single-row mode each data row is delivered as a separate+-- 'SingleTuple' result followed by a final 'TuplesOk' with no rows.+getNextResult :: Connection -> IO (Maybe NativeResult)+getNextResult connection = do+ pending <- readIORef connection.asyncPending+ if not pending+ then pure Nothing+ else do+ sepPending <- readIORef connection.pipelineSeparatorPending+ if sepPending+ then do+ writeIORef connection.pipelineSeparatorPending False+ pure Nothing+ else do+ singleRow <- readIORef connection.singleRowMode+ cachedFields <- readIORef connection.singleRowFields+ let initBuilder =+ if singleRow && not (null cachedFields)+ then emptyBuilder {accFields = cachedFields, accSawRowDescription = True}+ else emptyBuilder+ go singleRow initBuilder+ where+ -- Decrement the pending-command counter and set the separator flag when in+ -- pipeline mode. Called when a "terminal" result is about to be returned.+ finishCommand pipeStatus = do+ when (pipeStatus /= PipelineOff) $ do+ modifyIORef' connection.pendingCommands (subtract 1)+ writeIORef connection.pipelineSeparatorPending True++ go singleRow builder = do+ pipeStatus <- readIORef connection.pipelineStatus+ -- In aborted pipeline mode, if the server has already sent nothing for+ -- the remaining commands (it discards them after the first error), we+ -- generate synthetic PipelineAbort results for each outstanding command+ -- rather than blocking on a wire read that will never come.+ pending <- readIORef connection.pendingCommands+ if pipeStatus == PipelineAborted && pending > 0+ then do+ modifyIORef' connection.pendingCommands (subtract 1)+ writeIORef connection.pipelineSeparatorPending True+ pure (Just (NativeResult PipelineAbort [] [] Nothing Map.empty [] ""))+ else readAndProcess singleRow builder pipeStatus++ readAndProcess singleRow builder pipeStatus = do+ message <- nextMessage connection+ case message of+ RowDescription fs ->+ go singleRow builder {accFields = fs, accSawRowDescription = True, accHadResponse = True}+ ParameterDescription oids ->+ go singleRow builder {accParamOids = oids, accHadResponse = True}+ NoData ->+ go singleRow builder {accHadResponse = True}+ DataRow values ->+ if singleRow+ then do+ writeIORef connection.singleRowFields builder.accFields+ pure (Just (NativeResult SingleTuple builder.accFields [values] Nothing Map.empty [] ""))+ else go singleRow builder {accRevRows = values : builder.accRevRows}+ ParseComplete -> do+ parses <- readIORef connection.pendingParses+ if parses > 0 && pipeStatus /= PipelineOff+ then do+ modifyIORef' connection.pendingParses (subtract 1)+ finishCommand pipeStatus+ pure (Just (NativeResult CommandOk [] [] Nothing Map.empty [] ""))+ else go singleRow builder {accHadResponse = True}+ BindComplete ->+ go singleRow builder {accHadResponse = True}+ CloseComplete ->+ go singleRow builder {accHadResponse = True}+ CommandComplete tag -> do+ if singleRow+ then do+ writeIORef connection.singleRowMode False+ writeIORef connection.singleRowFields []+ writeIORef connection.lastError (Just "")+ pure (Just (NativeResult TuplesOk builder.accFields [] (Just tag) Map.empty [] ""))+ else do+ writeIORef connection.lastError (Just "")+ finishCommand pipeStatus+ pure (Just (commandResult builder (Just tag)))+ EmptyQueryResponse -> do+ writeIORef connection.lastError (Just "")+ finishCommand pipeStatus+ pure (Just (NativeResult EmptyQuery [] [] Nothing Map.empty [] ""))+ ErrorResponse fs -> do+ let errMap = Map.fromList fs+ case pipeStatus of+ PipelineAborted -> do+ -- Should not normally happen (server discards commands in abort+ -- mode) but handle defensively.+ finishCommand pipeStatus+ pure (Just (NativeResult PipelineAbort [] [] Nothing Map.empty [] ""))+ PipelineOn -> do+ writeIORef connection.pipelineStatus PipelineAborted+ finishCommand PipelineOn+ pure (Just (NativeResult FatalError [] [] Nothing errMap [] ""))+ PipelineOff -> do+ sql <- readIORef connection.currentQuery+ writeIORef connection.lastError (Just (formatResultError sql errMap))+ pure (Just (NativeResult FatalError [] [] Nothing errMap [] sql))+ PortalSuspended ->+ pure (Just (commandResult builder Nothing))+ ReadyForQuery txState -> do+ writeIORef connection.txStatus txState+ case pipeStatus of+ PipelineOff -> do+ writeIORef connection.asyncPending False+ if builder.accHadResponse+ then pure (Just (describeResult builder))+ else pure Nothing+ _ -> do+ writeIORef connection.pipelineStatus PipelineOn+ -- A PipelineSync result is its own command boundary: unlike a+ -- normal command result, libpq does not emit a separating NULL+ -- after it, so consecutive syncs are reported back-to-back. We+ -- therefore never set 'pipelineSeparatorPending' here. Only the+ -- final sync clears 'asyncPending'; an earlier one leaves it set+ -- so the next 'getNextResult' reads straight on to the next sync.+ remaining <- atomicModifyIORef' connection.pendingSyncs (\n -> (n - 1, n - 1))+ when (remaining == 0) $ writeIORef connection.asyncPending False+ pure (Just (NativeResult PipelineSync [] [] Nothing Map.empty [] ""))+ _ -> go singleRow builder++-- | Describe a prepared statement.+describePrepared :: Connection -> ByteString -> IO (Maybe NativeResult)+describePrepared connection name = withReady connection do+ sendMessage connection (describeStatementMessage name <> syncMessage)+ Just <$> collectExtended connection ""++-- | Describe a portal.+describePortal :: Connection -> ByteString -> IO (Maybe NativeResult)+describePortal connection name = withReady connection do+ sendMessage connection (describePortalMessage name <> syncMessage)+ Just <$> collectExtended connection ""++-- * Parameter projections++paramOid :: Maybe (Word32, ByteString, Format) -> Word32+paramOid = maybe 0 (\(oid, _, _) -> oid)++paramFormat :: Maybe (Word32, ByteString, Format) -> Int16+paramFormat = maybe 0 (\(_, _, format) -> formatCode format)++paramValue :: Maybe (Word32, ByteString, Format) -> Maybe ByteString+paramValue = fmap (\(_, value, _) -> value)++boundFormat :: Maybe (ByteString, Format) -> Int16+boundFormat = maybe 0 (formatCode . snd)++boundValue :: Maybe (ByteString, Format) -> Maybe ByteString+boundValue = fmap fst++formatCode :: Format -> Int16+formatCode = \case+ Text -> 0+ Binary -> 1++-- * Result collection++-- | Only run a flow on a ready connection; mirror libpq returning no result+-- when the connection is not usable.+withReady :: Connection -> IO (Maybe a) -> IO (Maybe a)+withReady connection action = do+ status <- readIORef connection.connStatus+ case status of+ ConnectionOk -> action+ _ -> pure Nothing++-- accumulator for a result under construction+data Builder = Builder+ { accFields :: [FieldDescription],+ accRevRows :: [[Maybe ByteString]],+ accParamOids :: [Word32],+ accSawRowDescription :: Bool,+ accHadResponse :: Bool+ }++emptyBuilder :: Builder+emptyBuilder = Builder [] [] [] False False++-- | Collect the (possibly several) results of a simple query, up to+-- @ReadyForQuery@. The last is what @PQexec@ returns.+-- @CopyInResponse@ and @CopyOutResponse@ terminate the loop immediately,+-- returning a synthetic result so the caller can enter the copy sub-protocol.+collectSimple :: Connection -> ByteString -> IO [NativeResult]+collectSimple connection sql = go emptyBuilder []+ where+ go builder acc = do+ message <- nextMessage connection+ case message of+ RowDescription fs -> go builder {accFields = fs, accSawRowDescription = True} acc+ DataRow values -> go builder {accRevRows = values : builder.accRevRows} acc+ CommandComplete tag -> do+ writeIORef connection.lastError (Just "")+ go emptyBuilder (commandResult builder (Just tag) : acc)+ EmptyQueryResponse -> do+ writeIORef connection.lastError (Just "")+ go emptyBuilder (NativeResult EmptyQuery [] [] Nothing Map.empty [] "" : acc)+ ErrorResponse fs -> do+ let errMap = Map.fromList fs+ writeIORef connection.lastError (Just (formatResultError sql errMap))+ go emptyBuilder (NativeResult FatalError [] [] Nothing errMap [] sql : acc)+ CopyInResponse _ formats ->+ let fields = map copyField formats+ in pure (reverse (NativeResult CopyIn fields [] Nothing Map.empty [] "" : acc))+ CopyOutResponse _ formats ->+ let fields = map copyField formats+ in pure (reverse (NativeResult CopyOut fields [] Nothing Map.empty [] "" : acc))+ ReadyForQuery txState -> do+ writeIORef connection.txStatus txState+ pure (reverse acc)+ _ -> go builder acc++-- | Collect the single result of an extended-protocol command.+collectExtended :: Connection -> ByteString -> IO NativeResult+collectExtended connection sql = go emptyBuilder Nothing+ where+ go builder finished = do+ message <- nextMessage connection+ case message of+ RowDescription fs -> go builder {accFields = fs, accSawRowDescription = True} finished+ ParameterDescription oids -> go builder {accParamOids = oids} finished+ NoData -> go builder finished+ DataRow values -> go builder {accRevRows = values : builder.accRevRows} finished+ ParseComplete -> go builder finished+ BindComplete -> go builder finished+ CloseComplete -> go builder finished+ PortalSuspended -> go emptyBuilder (finished <|> Just (commandResult builder Nothing))+ CommandComplete tag -> do+ writeIORef connection.lastError (Just "")+ go emptyBuilder (Just (commandResult builder (Just tag)))+ EmptyQueryResponse -> do+ writeIORef connection.lastError (Just "")+ go emptyBuilder (Just (NativeResult EmptyQuery [] [] Nothing Map.empty [] ""))+ ErrorResponse fs -> do+ let errMap = Map.fromList fs+ writeIORef connection.lastError (Just (formatResultError sql errMap))+ go emptyBuilder (Just (NativeResult FatalError [] [] Nothing errMap [] sql))+ ReadyForQuery txState -> do+ writeIORef connection.txStatus txState+ pure (fromMaybe (describeResult builder) finished)+ _ -> go builder finished++-- | A result terminated by @CommandComplete@\/@PortalSuspended@: 'TuplesOk' if a+-- row description was seen, else 'CommandOk'.+commandResult :: Builder -> Maybe ByteString -> NativeResult+commandResult builder tag =+ NativeResult+ (if builder.accSawRowDescription then TuplesOk else CommandOk)+ builder.accFields+ (reverse builder.accRevRows)+ tag+ Map.empty+ builder.accParamOids+ ""++-- | A result with no command completion (a @Describe@\/@Parse@-only flow):+-- 'CommandOk', carrying any column descriptions and parameter OIDs.+describeResult :: Builder -> NativeResult+describeResult builder =+ NativeResult+ CommandOk+ builder.accFields+ (reverse builder.accRevRows)+ Nothing+ Map.empty+ builder.accParamOids+ ""++lastMaybe :: [a] -> Maybe a+lastMaybe = foldl (\_ x -> Just x) Nothing++-- | Build a synthetic 'FieldDescription' from a COPY format code (0=text,+-- 1=binary). COPY results have no column names, table OID, type OID, etc.+copyField :: Int16 -> FieldDescription+copyField fmt =+ FieldDescription+ { name = "",+ tableOid = 0,+ columnAttributeNumber = 0,+ typeOid = 0,+ typeSize = 0,+ typeModifier = 0,+ formatCode = fmt+ }
+ src/library/Pqi/Native/Types.hs view
@@ -0,0 +1,252 @@+-- | Internal types for the native adapter, separated from 'Connection.hs' to+-- avoid orphan-instance warnings for 'IsResult' and 'IsCancel'.+module Pqi.Native.Types+ ( NativeResult (..),+ NativeCancel (..),+ formatErrorFields,+ formatResultError,+ )+where++import Control.Exception (IOException, try)+import qualified Data.ByteString as ByteString+import qualified Data.ByteString.Char8 as ByteString.Char8+import Data.Char (isDigit)+import Data.List (findIndex)+import qualified Data.Map.Strict as Map+import Pqi+ ( ExecStatus (..),+ FieldCode (..),+ Format (..),+ IsCancel (..),+ IsResult (..),+ PipelineStatus (..),+ )+import Pqi.Native.Prelude+import qualified Pqi.Native.Transport as Transport+import Pqi.Native.Transport.Message (FieldDescription (..), cancelRequest)++-- | A fully materialized result. The native adapter buffers the entire result+-- in memory, so the accessors are pure lookups and 'unsafeFreeResult' is a+-- no-op.+data NativeResult = NativeResult+ { status :: ExecStatus,+ fields :: [FieldDescription],+ rows :: [[Maybe ByteString]],+ commandTag :: Maybe ByteString,+ errorFields :: Map.Map Word8 ByteString,+ paramOids :: [Word32],+ -- | The query text that produced this result (used to format @LINE N:@+ -- position context in 'resultErrorMessage', matching libpq's behaviour).+ -- Empty for non-error results and for async results where the query is+ -- unavailable.+ queryText :: ByteString+ }+ deriving stock (Eq, Show)++-- | A standalone cancellation handle for the native adapter.+--+-- Carries references to connection state so the cancel implementation can+-- decide whether a network round-trip is necessary:+--+-- * @asyncPendingRef@: False when nothing is in flight at all — skip the+-- round-trip entirely.+-- * @pipelineStatusRef@ + @pendingCommandsRef@: in pipeline mode,+-- @asyncPending@ stays True even after all @CommandComplete@ messages+-- have arrived (while @ReadyForQuery@ is still unread). Sending a cancel+-- in that window produces a stale signal that can interrupt the very next+-- command (e.g. the @ABORT@ issued during clean-up). Checking+-- @pendingCommands > 0@ instead avoids the stale cancel: once all command+-- completions are in, @pendingCommands@ is 0 and the server is idle.+data NativeCancel = NativeCancel+ { host :: ByteString,+ port :: Int,+ pid :: Int32,+ secret :: Int32,+ asyncPendingRef :: IORef Bool,+ pipelineStatusRef :: IORef PipelineStatus,+ pendingCommandsRef :: IORef Int+ }++-- | Like 'formatResultError' but without a client query text (for+-- connection-level errors, which never carry a statement-position field).+formatErrorFields :: Map.Map Word8 ByteString -> ByteString+formatErrorFields = formatResultError ""++-- | Format error\/notice fields into a message string matching libpq's+-- @PQresultErrorMessage@ at DEFAULT verbosity with @SHOW_CONTEXT_ERRORS@+-- visibility.+--+-- Field order (matching libpq):+-- 1. @\<Severity\>: \<Message\>\\n@+-- 2. @LINE N: \<query line\>\\n \^\\n@ (statement position via @queryText@ + @\'P\'@ field)+-- 3. @LINE N: \<internal query\>\\n \^\\nQUERY: \<internal query\>\\n@ (@\'q\'@+@\'p\'@ fields)+-- 4. @DETAIL: \<D\>\\n@+-- 5. @HINT: \<H\>\\n@+-- 6. @CONTEXT: \<W\>\\n@ (ERROR\/FATAL\/PANIC only)+formatResultError :: ByteString -> Map.Map Word8 ByteString -> ByteString+formatResultError queryText fields =+ case Map.lookup 0x4d fields of+ Nothing -> ""+ Just msg ->+ let sev = Map.findWithDefault "" 0x53 fields+ isError = sev `elem` ["ERROR", "FATAL", "PANIC"]+ line1 = (if ByteString.null sev then msg else sev <> ": " <> msg) <> "\n"+ -- Statement position ('P', 0x50): needs client query text+ stmtCtx = case Map.lookup 0x50 fields of+ Just posStr+ | not (ByteString.null queryText) ->+ maybe "" (positionContext queryText) (readPositiveInt posStr)+ _ -> ""+ -- Internal query+position ('q'=0x71, 'p'=0x70): both in wire fields+ intCtx = case (Map.lookup 0x71 fields, Map.lookup 0x70 fields) of+ (Just intQuery, Just posStr) ->+ maybe+ ""+ (\p -> positionContext intQuery p <> "QUERY: " <> intQuery <> "\n")+ (readPositiveInt posStr)+ _ -> ""+ detLine = case Map.lookup 0x44 fields of+ Just det -> "DETAIL: " <> det <> "\n"+ _ -> ""+ hntLine = case Map.lookup 0x48 fields of+ Just hnt -> "HINT: " <> hnt <> "\n"+ _ -> ""+ ctxLine = case Map.lookup 0x57 fields of+ Just ctx | isError -> "CONTEXT: " <> ctx <> "\n"+ _ -> ""+ in line1 <> stmtCtx <> intCtx <> detLine <> hntLine <> ctxLine++-- | Build the @LINE N: \<text\>\\n \^\\n@ block for a 1-indexed position+-- within a query string, matching libpq's formatting exactly.+positionContext :: ByteString -> Int -> ByteString+positionContext query pos =+ let pos0 = max 0 (pos - 1)+ before = ByteString.take pos0 query+ lineNum = 1 + ByteString.length (ByteString.filter (== 0x0a) before)+ lineStart = maybe 0 (+ 1) (ByteString.elemIndexEnd 0x0a before)+ col = pos0 - lineStart+ rest = ByteString.drop lineStart query+ lineText = ByteString.takeWhile (/= 0x0a) rest+ prefix = ByteString.Char8.pack ("LINE " <> show lineNum <> ": ")+ caret = ByteString.replicate (ByteString.length prefix + col) 0x20 <> "^\n"+ in prefix <> lineText <> "\n" <> caret++readPositiveInt :: ByteString -> Maybe Int+readPositiveInt bs = case ByteString.Char8.readInt bs of+ Just (n, _) | n > 0 -> Just n+ _ -> Nothing++instance IsResult NativeResult where+ resultStatus result = pure result.status+ resultErrorMessage result = pure (Just (formatResultError result.queryText result.errorFields))+ resultErrorField result field = pure (Map.lookup (fieldCodeByte field) result.errorFields)+ unsafeFreeResult _ = pure ()+ ntuples result = pure (fromIntegral (length result.rows))+ nfields result = pure (fromIntegral (length result.fields))+ fname result column = pure $ do+ fd <- atMay result.fields column+ if ByteString.null fd.name then Nothing else Just fd.name+ fnumber result name =+ pure (fromIntegral <$> findIndex (\field -> field.name == folded) result.fields)+ where+ folded = foldIdentifier name+ ftable result column = pure (maybe 0 (.tableOid) (atMay result.fields column))+ ftablecol result column =+ pure (maybe 0 (\field -> fromIntegral (field.columnAttributeNumber :: Int16)) (atMay result.fields column))+ fformat result column =+ pure (maybe Text (\field -> formatOf field.formatCode) (atMay result.fields column))+ ftype result column = pure (maybe 0 (.typeOid) (atMay result.fields column))+ fmod result column = pure (maybe 0 (\field -> fromIntegral (field.typeModifier :: Int32)) (atMay result.fields column))+ fsize result column = pure (maybe 0 (\field -> fromIntegral (field.typeSize :: Int16)) (atMay result.fields column))+ getvalue result row column = pure (join (cellAt result row column))+ getvalue' result row column = pure (join (cellAt result row column))+ getisnull result row column = pure (maybe True isNothing (cellAt result row column))+ getlength result row column =+ pure (maybe 0 (maybe 0 ByteString.length) (cellAt result row column))+ nparams result = pure (fromIntegral (length result.paramOids))+ paramtype result index = pure (fromMaybe 0 (atMay result.paramOids index))+ cmdStatus result = pure (Just (fromMaybe "" result.commandTag))+ cmdTuples result = pure (Just (maybe "" affectedRows result.commandTag))++instance IsCancel NativeCancel where+ cancel nc = do+ pending <- readIORef nc.asyncPendingRef+ if not pending+ then pure (Right ())+ else do+ transport <- Transport.connect nc.host nc.port+ Transport.send transport (cancelRequest nc.pid nc.secret)+ -- Read until EOF to ensure the server has processed the cancel request+ -- before we close the connection. This matches libpq's PQcancel behavior+ -- and prevents the cancel signal from racing with the next query.+ _ <- try @IOException (Transport.readUntilClosed transport)+ Transport.close transport+ pure (Right ())++-- * Helpers for the 'IsResult' instance++atMay :: [a] -> Int32 -> Maybe a+atMay xs i+ | i < 0 = Nothing+ | otherwise = case drop (fromIntegral i) xs of+ x : _ -> Just x+ [] -> Nothing++cellAt :: NativeResult -> Int32 -> Int32 -> Maybe (Maybe ByteString)+cellAt result row column = do+ rowValues <- atMay result.rows row+ atMay rowValues column++formatOf :: Int16 -> Format+formatOf = \case+ 1 -> Binary+ _ -> Text++foldIdentifier :: ByteString -> ByteString+foldIdentifier = ByteString.pack . outside . ByteString.unpack+ where+ quote = 0x22+ outside = \case+ [] -> []+ c : rest+ | c == quote -> inside rest+ | otherwise -> asciiToLower c : outside rest+ inside = \case+ [] -> []+ c : rest+ | c == quote -> case rest of+ c' : rest' | c' == quote -> quote : inside rest'+ _ -> outside rest+ | otherwise -> c : inside rest+ asciiToLower c+ | c >= 0x41 && c <= 0x5a = c + 0x20+ | otherwise = c++fieldCodeByte :: FieldCode -> Word8+fieldCodeByte = \case+ DiagSeverity -> 0x53 -- 'S'+ DiagSqlstate -> 0x43 -- 'C'+ DiagMessagePrimary -> 0x4d -- 'M'+ DiagMessageDetail -> 0x44 -- 'D'+ DiagMessageHint -> 0x48 -- 'H'+ DiagStatementPosition -> 0x50 -- 'P'+ DiagInternalPosition -> 0x70 -- 'p'+ DiagInternalQuery -> 0x71 -- 'q'+ DiagContext -> 0x57 -- 'W'+ DiagSourceFile -> 0x46 -- 'F'+ DiagSourceLine -> 0x4c -- 'L'+ DiagSourceFunction -> 0x52 -- 'R'++-- | The affected-row count from a @CommandComplete@ tag, matching+-- @PQcmdTuples@: the last whitespace-delimited token if it is all digits (which+-- covers @INSERT oid rows@, @UPDATE n@, @SELECT n@, …), else empty.+affectedRows :: ByteString -> ByteString+affectedRows tag =+ case ByteString.Char8.words tag of+ [] -> ""+ tokens ->+ let final = last tokens+ in if not (ByteString.null final) && ByteString.Char8.all isDigit final+ then final+ else ""
+ src/test/Spec.hs view
@@ -0,0 +1,14 @@+-- | The native adapter's conformance test suite: a single delegate to+-- 'Pqi.Conformance.specs', which brings up a throwaway PostgreSQL+-- container, runs the full differential battery (core, capabilities, and+-- SCRAM) against the FFI reference, and tears the container down again.+module Main (main) where++import Data.Proxy (Proxy (..))+import Pqi.Conformance (specs)+import Pqi.Native (Connection)+import Test.Hspec+import Prelude++main :: IO ()+main = hspec (specs (Proxy @Connection))
+ src/transport/Pqi/Native/Transport.hs view
@@ -0,0 +1,119 @@+-- | The byte-level transport: a TCP socket with a read buffer, plus the framing+-- that turns the stream into discrete @[type byte][Int32 length][body]@+-- messages.+module Pqi.Native.Transport+ ( Transport,+ connect,+ unconnected,+ close,+ send,+ receiveFrame,+ socketFd,+ peerIp,+ readUntilClosed,+ )+where++import qualified Data.ByteString as ByteString+import qualified Data.ByteString.Char8 as ByteString.Char8+import Data.IORef+import qualified Network.Socket as Socket+import qualified Network.Socket.ByteString as Socket.ByteString+import Pqi.Native.Transport.Prelude+import qualified PtrPoker.Write as Poker+import System.IO.Error (eofErrorType, mkIOError)++-- | An open connection's byte transport.+data Transport = Transport+ { socket :: Socket.Socket,+ readBuffer :: IORef ByteString+ }++-- | Open a TCP connection to the given host and port.+connect :: ByteString -> Int -> IO Transport+connect host port = do+ let hints = Socket.defaultHints {Socket.addrSocketType = Socket.Stream}+ addresses <-+ Socket.getAddrInfo (Just hints) (Just (ByteString.Char8.unpack host)) (Just (show port))+ case addresses of+ [] -> ioError (userError ("Could not resolve host: " <> ByteString.Char8.unpack host))+ address : _ -> do+ sock <- Socket.socket (Socket.addrFamily address) (Socket.addrSocketType address) (Socket.addrProtocol address)+ Socket.connect sock (Socket.addrAddress address)+ buffer <- newIORef ByteString.empty+ pure Transport {socket = sock, readBuffer = buffer}++-- | An unconnected transport, for representing a \"null\" connection. Its+-- socket is allocated but never connected; it must not be used for I\/O.+unconnected :: IO Transport+unconnected = do+ sock <- Socket.socket Socket.AF_INET Socket.Stream Socket.defaultProtocol+ buffer <- newIORef ByteString.empty+ pure Transport {socket = sock, readBuffer = buffer}++-- | Close the connection.+close :: Transport -> IO ()+close transport = Socket.close transport.socket++-- | The underlying socket file descriptor.+socketFd :: Transport -> IO Int32+socketFd transport = fromIntegral <$> Socket.unsafeFdSocket transport.socket++-- | Send a serialized message.+send :: Transport -> Poker.Write -> IO ()+send transport write = Socket.ByteString.sendAll transport.socket (Poker.toByteString write)++-- | Read exactly @n@ bytes, looping over @recv@ (which yields up to @n@) and+-- buffering any overshoot. Throws on EOF before @n@ bytes arrive.+receiveExactly :: Transport -> Int -> IO ByteString+receiveExactly transport n = do+ buffered <- readIORef transport.readBuffer+ go buffered+ where+ go accumulated+ | ByteString.length accumulated >= n = do+ let (result, rest) = ByteString.splitAt n accumulated+ writeIORef transport.readBuffer rest+ pure result+ | otherwise = do+ chunk <- Socket.ByteString.recv transport.socket (max 4096 (n - ByteString.length accumulated))+ if ByteString.null chunk+ then ioError (mkIOError eofErrorType "pqi-native: connection closed by server" Nothing Nothing)+ else go (accumulated <> chunk)++-- | Receive one framed message: its type byte and its body (the length prefix,+-- which counts itself, is consumed).+receiveFrame :: Transport -> IO (Word8, ByteString)+receiveFrame transport = do+ header <- receiveExactly transport 5+ let typeByte = ByteString.head header+ frameLength = decodeInt32BE (ByteString.drop 1 header)+ bodyLength = frameLength - 4+ body <-+ if bodyLength > 0+ then receiveExactly transport bodyLength+ else pure ByteString.empty+ pure (typeByte, body)++-- | Read bytes until the peer closes the connection (EOF), discarding them.+-- Used by the cancel path to mirror libpq's behaviour: keep the socket open+-- until the server has read the cancel request and closed its end.+readUntilClosed :: Transport -> IO ()+readUntilClosed transport = go+ where+ go = do+ chunk <- Socket.ByteString.recv transport.socket 4096+ if ByteString.null chunk+ then pure ()+ else go++-- | The numeric IP address of the connected peer (e.g. @\"::1\"@ or+-- @\"127.0.0.1\"@). Throws if the socket has no peer (unconnected).+peerIp :: Transport -> IO ByteString+peerIp transport = do+ addr <- Socket.getPeerName transport.socket+ (Just ip, _) <- Socket.getNameInfo [Socket.NI_NUMERICHOST] True False addr+ pure (ByteString.Char8.pack ip)++decodeInt32BE :: ByteString -> Int+decodeInt32BE = ByteString.foldl' (\acc w -> acc * 256 + fromIntegral w) 0 . ByteString.take 4
+ src/transport/Pqi/Native/Transport/Message.hs view
@@ -0,0 +1,302 @@+-- | The PostgreSQL frontend\/backend wire messages: builders ('Poker.Write') for+-- the frontend messages we send, and a decoder for the backend messages we+-- receive. Message framing (@[type byte][Int32 length][body]@, with the+-- startup message lacking a type byte) lives here too.+module Pqi.Native.Transport.Message+ ( -- * Frontend messages+ startupMessage,+ sslRequest,+ cancelRequest,+ passwordMessage,+ saslInitialResponse,+ saslResponse,+ queryMessage,+ parseMessage,+ bindMessage,+ describeStatementMessage,+ describePortalMessage,+ executeMessage,+ closeStatementMessage,+ closePortalMessage,+ syncMessage,+ flushMessage,+ terminateMessage,+ copyDataMessage,+ copyDoneMessage,+ copyFailMessage,++ -- * Backend messages+ BackendMessage (..),+ FieldDescription (..),+ decodeBackendMessage,+ )+where++import qualified Data.ByteString as ByteString+import Pqi.Native.Comms+import Pqi.Native.Transport.Prelude+import qualified PtrPoker.Write as Poker++-- * Framing helpers++-- | Frame a typed frontend message: type byte, then a length covering the+-- length field and the body, then the body.+framed :: Word8 -> Poker.Write -> Poker.Write+framed typeByte body =+ Poker.word8 typeByte+ <> Poker.bInt32 (fromIntegral (4 + Poker.writeSize body))+ <> body++-- | A null-terminated string.+cstr :: ByteString -> Poker.Write+cstr value = Poker.byteString value <> Poker.word8 0++formatCodes :: [Int16] -> Poker.Write+formatCodes codes =+ Poker.bInt16 (fromIntegral (length codes)) <> foldMap Poker.bInt16 codes++-- * Frontend messages++-- | The startup message (no type byte), requesting protocol 3.0 with the given+-- parameters (e.g. @user@, @database@).+startupMessage :: [(ByteString, ByteString)] -> Poker.Write+startupMessage parameters =+ let body =+ Poker.bInt32 196608 -- protocol version 3.0 == (3 << 16)+ <> foldMap (\(k, v) -> cstr k <> cstr v) parameters+ <> Poker.word8 0+ in Poker.bInt32 (fromIntegral (4 + Poker.writeSize body)) <> body++-- | The SSL negotiation request (no type byte).+sslRequest :: Poker.Write+sslRequest = Poker.bInt32 8 <> Poker.bInt32 80877103++-- | A cancel request (no type byte), sent on a fresh connection: the cancel+-- request code followed by the target backend's process ID and secret key.+cancelRequest :: Int32 -> Int32 -> Poker.Write+cancelRequest processId secretKey =+ Poker.bInt32 16 <> Poker.bInt32 80877102 <> Poker.bInt32 processId <> Poker.bInt32 secretKey++-- | A cleartext or MD5 password response (message type @p@).+passwordMessage :: ByteString -> Poker.Write+passwordMessage password = framed 0x70 (cstr password)++-- | A SASL initial response (message type @p@): the chosen mechanism and the+-- client-first message.+saslInitialResponse :: ByteString -> ByteString -> Poker.Write+saslInitialResponse mechanism initialData =+ framed 0x70+ $ cstr mechanism+ <> Poker.bInt32 (fromIntegral (ByteString.length initialData))+ <> Poker.byteString initialData++-- | A subsequent SASL response (message type @p@).+saslResponse :: ByteString -> Poker.Write+saslResponse responseData = framed 0x70 (Poker.byteString responseData)++-- | A simple @Query@ message.+queryMessage :: ByteString -> Poker.Write+queryMessage sql = framed 0x51 (cstr sql)++-- | A @Parse@ message: statement name, SQL, and parameter type OIDs.+parseMessage :: ByteString -> ByteString -> [Word32] -> Poker.Write+parseMessage name sql oids =+ framed 0x50+ $ cstr name+ <> cstr sql+ <> Poker.bInt16 (fromIntegral (length oids))+ <> foldMap Poker.bWord32 oids++-- | A @Bind@ message: portal, statement, per-parameter format codes, parameter+-- values ('Nothing' for SQL @NULL@), and result format codes.+bindMessage :: ByteString -> ByteString -> [Int16] -> [Maybe ByteString] -> [Int16] -> Poker.Write+bindMessage portal statement parameterFormats values resultFormats =+ framed 0x42+ $ cstr portal+ <> cstr statement+ <> formatCodes parameterFormats+ <> Poker.bInt16 (fromIntegral (length values))+ <> foldMap bindValue values+ <> formatCodes resultFormats+ where+ bindValue = \case+ Nothing -> Poker.bInt32 (-1)+ Just value -> Poker.bInt32 (fromIntegral (ByteString.length value)) <> Poker.byteString value++-- | A @Describe@ message targeting a prepared statement.+describeStatementMessage :: ByteString -> Poker.Write+describeStatementMessage name = framed 0x44 (Poker.word8 0x53 <> cstr name)++-- | A @Describe@ message targeting a portal.+describePortalMessage :: ByteString -> Poker.Write+describePortalMessage name = framed 0x44 (Poker.word8 0x50 <> cstr name)++-- | An @Execute@ message: portal name and a maximum row count (0 = unlimited).+executeMessage :: ByteString -> Int32 -> Poker.Write+executeMessage portal maxRows = framed 0x45 (cstr portal <> Poker.bInt32 maxRows)++-- | A @Close@ message targeting a prepared statement.+closeStatementMessage :: ByteString -> Poker.Write+closeStatementMessage name = framed 0x43 (Poker.word8 0x53 <> cstr name)++-- | A @Close@ message targeting a portal.+closePortalMessage :: ByteString -> Poker.Write+closePortalMessage name = framed 0x43 (Poker.word8 0x50 <> cstr name)++-- | A @Sync@ message.+syncMessage :: Poker.Write+syncMessage = framed 0x53 mempty++-- | A @Flush@ message.+flushMessage :: Poker.Write+flushMessage = framed 0x48 mempty++-- | A @Terminate@ message.+terminateMessage :: Poker.Write+terminateMessage = framed 0x58 mempty++-- | A @CopyData@ message.+copyDataMessage :: ByteString -> Poker.Write+copyDataMessage payload = framed 0x64 (Poker.byteString payload)++-- | A @CopyDone@ message.+copyDoneMessage :: Poker.Write+copyDoneMessage = framed 0x63 mempty++-- | A @CopyFail@ message with an error string.+copyFailMessage :: ByteString -> Poker.Write+copyFailMessage reason = framed 0x66 (cstr reason)++-- * Backend messages++-- | The description of one result column, from a @RowDescription@ message.+data FieldDescription = FieldDescription+ { name :: ByteString,+ tableOid :: Word32,+ columnAttributeNumber :: Int16,+ typeOid :: Word32,+ typeSize :: Int16,+ typeModifier :: Int32,+ formatCode :: Int16+ }+ deriving stock (Eq, Show)++-- | A decoded backend message. Only the variants the adapter acts on are+-- distinguished; anything else becomes 'UnknownMessage'.+data BackendMessage+ = AuthenticationOk+ | AuthenticationCleartextPassword+ | AuthenticationMD5Password ByteString+ | AuthenticationSASL [ByteString]+ | AuthenticationSASLContinue ByteString+ | AuthenticationSASLFinal ByteString+ | ParameterStatus ByteString ByteString+ | BackendKeyData Int32 Int32+ | ReadyForQuery Word8+ | RowDescription [FieldDescription]+ | DataRow [Maybe ByteString]+ | CommandComplete ByteString+ | EmptyQueryResponse+ | ErrorResponse [(Word8, ByteString)]+ | NoticeResponse [(Word8, ByteString)]+ | ParseComplete+ | BindComplete+ | CloseComplete+ | ParameterDescription [Word32]+ | NoData+ | PortalSuspended+ | NotificationResponse Int32 ByteString ByteString+ | CopyInResponse Int16 [Int16]+ | CopyOutResponse Int16 [Int16]+ | CopyData ByteString+ | CopyDone+ | FunctionCallResponse (Maybe ByteString)+ | UnknownMessage Word8 ByteString+ deriving stock (Eq, Show)++-- | Decode a backend message from its type byte and body.+decodeBackendMessage :: Word8 -> ByteString -> Either DecodingError BackendMessage+decodeBackendMessage typeByte body =+ case typeByte of+ 0x52 -> runDecoder authentication body+ 0x53 -> runDecoder (ParameterStatus <$> cstring <*> cstring) body+ 0x4b -> runDecoder (BackendKeyData <$> int32 <*> int32) body+ 0x5a -> runDecoder (ReadyForQuery <$> word8) body+ 0x54 -> runDecoder (RowDescription <$> repeatedInt16 fieldDescription) body+ 0x44 -> runDecoder (DataRow <$> repeatedInt16 columnValue) body+ 0x43 -> runDecoder (CommandComplete <$> cstring) body+ 0x49 -> Right EmptyQueryResponse+ 0x45 -> runDecoder (ErrorResponse <$> noticeFields) body+ 0x4e -> runDecoder (NoticeResponse <$> noticeFields) body+ 0x31 -> Right ParseComplete+ 0x32 -> Right BindComplete+ 0x33 -> Right CloseComplete+ 0x74 -> runDecoder (ParameterDescription <$> repeatedInt16 word32) body+ 0x6e -> Right NoData+ 0x73 -> Right PortalSuspended+ 0x41 -> runDecoder (NotificationResponse <$> int32 <*> cstring <*> cstring) body+ 0x47 -> runDecoder (uncurry CopyInResponse <$> copyResponse) body+ 0x48 -> runDecoder (uncurry CopyOutResponse <$> copyResponse) body+ 0x64 -> runDecoder (CopyData <$> remaining) body+ 0x63 -> Right CopyDone+ 0x56 -> runDecoder (FunctionCallResponse <$> columnValue) body+ other -> Right (UnknownMessage other body)++-- | Repeat a decoder an @Int16@-prefixed number of times.+repeatedInt16 :: Decoder a -> Decoder [a]+repeatedInt16 element = do+ count <- int16+ replicateM (fromIntegral count) element++fieldDescription :: Decoder FieldDescription+fieldDescription =+ FieldDescription+ <$> cstring+ <*> word32+ <*> int16+ <*> word32+ <*> int16+ <*> int32+ <*> int16++columnValue :: Decoder (Maybe ByteString)+columnValue = do+ len <- int32+ if len < 0+ then pure Nothing+ else Just <$> bytes (fromIntegral len)++-- | The fields of an @ErrorResponse@\/@NoticeResponse@: @(code, value)@ pairs,+-- terminated by a zero code byte.+noticeFields :: Decoder [(Word8, ByteString)]+noticeFields = do+ code <- word8+ if code == 0+ then pure []+ else do+ value <- cstring+ ((code, value) :) <$> noticeFields++copyResponse :: Decoder (Int16, [Int16])+copyResponse = do+ overall <- fromIntegral <$> word8+ columns <- repeatedInt16 int16+ pure (overall, columns)++authentication :: Decoder BackendMessage+authentication = do+ subtype <- int32+ case subtype of+ 0 -> pure AuthenticationOk+ 3 -> pure AuthenticationCleartextPassword+ 5 -> AuthenticationMD5Password <$> bytes 4+ 10 -> AuthenticationSASL <$> saslMechanisms+ 11 -> AuthenticationSASLContinue <$> remaining+ 12 -> AuthenticationSASLFinal <$> remaining+ other -> throwE (UnexpectedAuthType other)+ where+ saslMechanisms = do+ mechanism <- cstring+ if ByteString.null mechanism+ then pure []+ else (mechanism :) <$> saslMechanisms
+ src/transport/Pqi/Native/Transport/Prelude.hs view
@@ -0,0 +1,19 @@+module Pqi.Native.Transport.Prelude+ ( module Exports,+ )+where++import Control.Applicative as Exports+import Control.Monad as Exports+import Control.Monad.Trans.Class as Exports (lift)+import Control.Monad.Trans.Except as Exports+ ( ExceptT (..),+ except,+ runExceptT,+ throwE,+ withExceptT,+ )+import Data.ByteString as Exports (ByteString)+import Data.Int as Exports+import Data.Word as Exports+import Prelude as Exports