valiant (empty) → 0.1.0.0
raw patch · 80 files changed
+10865/−0 lines, 80 filesdep +aesondep +asyncdep +base
Dependencies added: aeson, async, base, bytestring, containers, criterion, deepseq, hashable, hedgehog, hspec, hspec-hedgehog, network, nothunks, pg-wire, psqueues, scientific, stm, text, time, uuid-types, valiant, vector
Files
- CHANGELOG.md +31/−0
- LICENSE +27/−0
- README.md +51/−0
- bench/BenchAsyncpgCompare.hs +179/−0
- bench/BenchCodec.hs +141/−0
- bench/BenchConcurrent.hs +126/−0
- bench/BenchPool.hs +81/−0
- bench/BenchQuery.hs +211/−0
- bench/Main.hs +22/−0
- bench/TestSupport.hs +107/−0
- integration/ChaosSpec.hs +136/−0
- integration/ConnectionFeaturesSpec.hs +165/−0
- integration/ConnectionSpec.hs +37/−0
- integration/CopySpec.hs +35/−0
- integration/ExecuteSpec.hs +131/−0
- integration/LargeObjectSpec.hs +91/−0
- integration/Main.hs +30/−0
- integration/NotifySpec.hs +51/−0
- integration/PipelineSpec.hs +116/−0
- integration/PoolSpec.hs +467/−0
- integration/SoundnessSpec.hs +285/−0
- integration/StreamingSpec.hs +71/−0
- integration/TestSupport.hs +107/−0
- integration/TransactionSpec.hs +96/−0
- src/Valiant.hs +404/−0
- src/Valiant/Advisory.hs +139/−0
- src/Valiant/Batch.hs +140/−0
- src/Valiant/Binary/Array.hs +108/−0
- src/Valiant/Binary/Composite.hs +90/−0
- src/Valiant/Binary/Decode.hs +203/−0
- src/Valiant/Binary/Encode.hs +221/−0
- src/Valiant/Binary/Enum.hs +44/−0
- src/Valiant/Binary/HStore.hs +130/−0
- src/Valiant/Binary/Inet.hs +198/−0
- src/Valiant/Binary/Interval.hs +71/−0
- src/Valiant/Binary/JSON.hs +44/−0
- src/Valiant/Binary/MacAddr.hs +65/−0
- src/Valiant/Binary/Point.hs +62/−0
- src/Valiant/Binary/Range.hs +135/−0
- src/Valiant/Binary/Scientific.hs +147/−0
- src/Valiant/Binary/UUID.hs +29/−0
- src/Valiant/Binary/Unbounded.hs +167/−0
- src/Valiant/Copy.hs +204/−0
- src/Valiant/Dynamic.hs +140/−0
- src/Valiant/Error.hs +191/−0
- src/Valiant/Execute.hs +817/−0
- src/Valiant/Fold.hs +99/−0
- src/Valiant/FromRow.hs +210/−0
- src/Valiant/FromRowFast.hs +174/−0
- src/Valiant/FromRowStrict.hs +65/−0
- src/Valiant/LargeObject.hs +196/−0
- src/Valiant/Logging.hs +156/−0
- src/Valiant/NamedParams.hs +174/−0
- src/Valiant/Notify.hs +112/−0
- src/Valiant/Pipeline.hs +160/−0
- src/Valiant/Statement.hs +116/−0
- src/Valiant/Streaming.hs +220/−0
- src/Valiant/ToParams.hs +154/−0
- src/Valiant/Transaction.hs +312/−0
- test/Main.hs +44/−0
- test/Valiant/Binary/ArraySpec.hs +84/−0
- test/Valiant/Binary/CompositeSpec.hs +55/−0
- test/Valiant/Binary/DecodeSpec.hs +157/−0
- test/Valiant/Binary/EncodeSpec.hs +120/−0
- test/Valiant/Binary/HStoreSpec.hs +46/−0
- test/Valiant/Binary/InetSpec.hs +87/−0
- test/Valiant/Binary/IntervalSpec.hs +54/−0
- test/Valiant/Binary/JSONSpec.hs +61/−0
- test/Valiant/Binary/MacAddrSpec.hs +53/−0
- test/Valiant/Binary/NewtypeSpec.hs +131/−0
- test/Valiant/Binary/PropertyHedgehogSpec.hs +278/−0
- test/Valiant/Binary/RangeSpec.hs +59/−0
- test/Valiant/Binary/RefineSpec.hs +58/−0
- test/Valiant/Binary/ScientificSpec.hs +48/−0
- test/Valiant/Binary/UUIDSpec.hs +48/−0
- test/Valiant/ErrorSpec.hs +143/−0
- test/Valiant/FromRowSpec.hs +118/−0
- test/Valiant/NamedParamsSpec.hs +180/−0
- test/Valiant/TupleSpec.hs +119/−0
- valiant.cabal +231/−0
+ CHANGELOG.md view
@@ -0,0 +1,31 @@+# 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. Runtime library on top of+[`pg-wire`](https://hackage.haskell.org/package/pg-wire). Re-exports+the connection and pool APIs and adds:++- Binary codecs for the common PostgreSQL types: numerics, text,+ time, UUID, JSON / JSONB, hstore, inet, macaddr, point, ranges,+ arrays, composites, enums.+- A typed `Statement p r` value, query execution+ (`fetchOne`/`fetchAll`/`fetchScalar`/`fetchExists`/...), command+ execution (`execute`/`executeReturning`/`executeBatch`), pipelined+ batch execution.+- Server-side cursor streaming (`withCursor`, `fetchBatch`,+ `fetchAllCursor`).+- COPY in/out (binary and CSV), large objects, LISTEN/NOTIFY,+ advisory locks, transactions with isolation levels.+- `refine` and `refineWith` combinators in `Valiant.Binary.Decode`+ for validating decoded values against extra invariants without a+ full `PgDecode` instance.+- `fetchAllUnboxed`: an unboxed-vector variant of `fetchAllVec` for+ fixed-size primitive result columns. Eliminates per-element pointer+ indirection for `Int32` / `Int64` / `Double` / `Bool` and tuples of+ those.++[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 @@+# valiant++Compile-time checked SQL for Haskell. Runtime library.++## What this package gives you++- The `Statement` type and query execution functions: `execute`,+ `fetchOne`, `fetchAll`, `fetchAllVec`, `fetchAllUnboxed`, server-side+ cursors (`withCursor`, `fetchBatch`, `fetchAllCursor`,+ `executeWithFold`).+- Pipelined batch execution (`runPipeline`) with an `Applicative`+ interface for combining independent reads into one round-trip.+- Binary codecs for all common PostgreSQL types: numeric and text+ primitives, time, UUID, JSON / JSONB, hstore, inet / cidr / macaddr,+ point, ranges, arrays, composites, enums.+- COPY in / out (text, CSV, binary), large objects, LISTEN / NOTIFY,+ advisory locks, transactions (with savepoints and read-only mode),+ prepared statements with statement caching.+- Re-exports from [`pg-wire`](https://hackage.haskell.org/package/pg-wire):+ connection pool, TLS, authentication, error types.++## Compile-time validation++The `Statement` values you execute typically come from the+`Valiant.Plugin` GHC source plugin (in the+[`valiant-plugin`](https://hackage.haskell.org/package/valiant-plugin)+package). The plugin checks parameter and result types against a+`.valiant/` cache produced by the+[`valiant-cli`](https://hackage.haskell.org/package/valiant-cli) tool,+which talks to a live PostgreSQL database via PREPARE / DESCRIBE.++You can also build `Statement` values manually with `mkStatement`,+which skips the compile-time check.++## Installation++```+cabal install valiant+```++## Documentation++- [Top-level README](https://github.com/joshburgess/valiant#readme):+ full project overview, getting started, examples.+- [Tutorial](https://github.com/joshburgess/valiant/blob/main/docs/TUTORIAL.md).+- [Performance notes](https://github.com/joshburgess/valiant/blob/main/docs/PERFORMANCE.md)+ and [benchmark archive](https://github.com/joshburgess/valiant/tree/main/docs/benchmark-results).++## License++BSD-3-Clause. See `LICENSE`.
+ bench/BenchAsyncpgCompare.hs view
@@ -0,0 +1,179 @@+{-# LANGUAGE BangPatterns #-}++-- | Benchmarks matching asyncpg's methodology for direct comparison.+--+-- asyncpg uses: 10 concurrent connections, Unix socket, 30s duration,+-- 7 query patterns. We replicate the same patterns here.+--+-- See: https://github.com/MagicStack/pgbench+module BenchAsyncpgCompare (benchmarks) where++import Control.Concurrent.Async (mapConcurrently)+import Criterion.Main+import Data.Int (Int16, Int32)+import Data.IORef+import Data.Text (Text)+import Data.Text qualified as T+import Valiant+import System.IO.Unsafe (unsafePerformIO)+import TestSupport (requireDatabaseUrl)++{-# NOINLINE globalPool #-}+globalPool :: IORef (Maybe Pool)+globalPool = unsafePerformIO (newIORef Nothing)++getPool :: Int -> IO Pool+getPool size = do+ mp <- readIORef globalPool+ case mp of+ Just p -> pure p+ Nothing -> do+ url <- requireDatabaseUrl+ p <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = size+ , poolAcquireTimeout = 10+ }+ writeIORef globalPool (Just p)+ -- Setup schema for benchmarks+ withResource p $ \conn -> do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS _bench_series CASCADE"+ _ <- simpleQuery conn "DROP TABLE IF EXISTS _bench_insert CASCADE"+ _ <- simpleQuery conn+ "CREATE TABLE _bench_insert (\+ \ a int, b int, c int, d int, e text, f text, g text\+ \)"+ pure ()+ pure p++benchmarks :: [Benchmark]+benchmarks =+ [ bgroup "asyncpg-compare"+ -- Matching asyncpg's 7 benchmark patterns with 10 concurrent connections.+ [ bgroup "select-1+1"+ -- asyncpg benchmark #7: SELECT 1+1 (minimal overhead)+ [ bench "10 conns x 100 queries" $+ whnfIO (concurrentQueries 10 100 queryOnePlusOne)+ ]+ , bgroup "generate-series-1000"+ -- asyncpg benchmark #2: SELECT i FROM generate_series(1, 1000)+ [ bench "10 conns x 10 queries" $+ whnfIO (concurrentQueries 10 10 queryGenerateSeries)+ ]+ , bgroup "pg-type-wide-rows"+ -- asyncpg benchmark #1: wide rows from pg_type (~350 rows, 12 columns)+ [ bench "10 conns x 10 queries" $+ whnfIO (concurrentQueries 10 10 queryPgType)+ ]+ , bgroup "batch-insert-1000"+ -- asyncpg benchmark #6: 1000 individual parameterized INSERTs+ [ bench "10 conns x 1 batch" $+ whnfIO (concurrentQueries 10 1 queryBatchInsert)+ ]+ , bgroup "pg-type-wide-rows-fast"+ -- Same as pg-type but using fetchAllFast (FromRowFast)+ [ bench "10 conns x 10 queries" $+ whnfIO (concurrentQueries 10 10 queryPgTypeFast)+ ]+ , bgroup "batch-insert-1000-pipelined"+ -- Same as above but using valiant's pipelined executeBatch+ [ bench "10 conns x 1 batch" $+ whnfIO (concurrentQueries 10 1 queryBatchInsertPipelined)+ ]++ -- Throughput tests at asyncpg's concurrency level+ , bgroup "throughput"+ [ bench "10 conns x 1000 SELECT 1+1" $+ whnfIO (concurrentQueries 10 1000 queryOnePlusOne)+ , bench "10 conns x 100 fetch-1000-rows" $+ whnfIO (concurrentQueries 10 100 queryGenerateSeries)+ ]+ ]+ ]++-- | N threads, each acquires from pool, runs M iterations of a query function.+concurrentQueries :: Int -> Int -> (Connection -> IO ()) -> IO ()+concurrentQueries nConns nQueries queryFn = do+ pool <- getPool nConns+ _ <- mapConcurrently (\_ ->+ withResource pool $ \conn ->+ mapM_ (\_ -> queryFn conn) [1 :: Int .. nQueries]+ ) [1 .. nConns]+ pure ()++-- Statements ----------------------------------------------------------------++stmtOnePlusOne :: Statement () Int32+stmtOnePlusOne = mkStatement "SELECT (1 + 1)::int4" [] ["?column?"] "<bench>"++stmtGenerateSeries :: Statement Int32 Int32+stmtGenerateSeries = mkStatement+ "SELECT i::int4 FROM generate_series(1, $1) AS i"+ [23] ["i"] "<bench>"++-- pg_type query: 12 columns matching asyncpg's benchmark #1+-- typname(text), typnamespace(oid), typowner(oid), typlen(int2), typbyval(bool),+-- typcategory(text), typispreferred(bool), typisdefined(bool), typdelim(text),+-- typrelid(oid), typelem(oid), typarray(oid)+stmtPgType :: Statement () (Text, Int32, Int32, Int16, Bool, Text, Bool, Bool, Text, Int32, Int32, Int32)+stmtPgType = mkStatement+ "SELECT typname, typnamespace::int4, typowner::int4, typlen, typbyval, \+ \typcategory::text, typispreferred, typisdefined, typdelim::text, typrelid::int4, \+ \typelem::int4, typarray::int4 FROM pg_type WHERE typtypmod = -1 AND typisdefined = true"+ [] ["typname","typnamespace","typowner","typlen","typbyval",+ "typcategory","typispreferred","typisdefined","typdelim",+ "typrelid","typelem","typarray"] "<bench>"++stmtInsertBench :: Statement (Int32, Int32, Int32, Int32, Text, Text, Text) ()+stmtInsertBench = mkStatement+ "INSERT INTO _bench_insert (a, b, c, d, e, f, g) VALUES ($1, $2, $3, $4, $5, $6, $7)"+ [23, 23, 23, 23, 25, 25, 25] [] "<bench>"++-- Query functions -----------------------------------------------------------++queryOnePlusOne :: Connection -> IO ()+queryOnePlusOne conn = do+ !_ <- fetchScalar conn stmtOnePlusOne ()+ pure ()++queryGenerateSeries :: Connection -> IO ()+queryGenerateSeries conn = do+ !_ <- fetchAll conn stmtGenerateSeries (1000 :: Int32)+ pure ()++queryPgType :: Connection -> IO ()+queryPgType conn = do+ -- Match asyncpg's pg_type query: wide rows from system catalog, binary format+ !_ <- fetchAll conn stmtPgType ()+ pure ()++queryPgTypeFast :: Connection -> IO ()+queryPgTypeFast conn = do+ -- Same query but using FromRowFast (no Either allocation per column)+ !_ <- fetchAllFast conn stmtPgType ()+ pure ()++queryBatchInsert :: Connection -> IO ()+queryBatchInsert conn = do+ -- Sequential inserts (matching asyncpg's benchmark #6)+ _ <- simpleQuery conn "TRUNCATE _bench_insert"+ mapM_ (\i ->+ execute conn stmtInsertBench+ (i, i * 2, i * 3, i * 4,+ "val_" <> T.pack (show i),+ "text_" <> T.pack (show i),+ "data_" <> T.pack (show i))+ ) [1 :: Int32 .. 1000]++queryBatchInsertPipelined :: Connection -> IO ()+queryBatchInsertPipelined conn = do+ -- Pipelined inserts (valiant advantage — asyncpg doesn't have this)+ _ <- simpleQuery conn "TRUNCATE _bench_insert"+ _ <- executeBatch conn stmtInsertBench+ [ (i, i * 2, i * 3, i * 4,+ "val_" <> T.pack (show i),+ "text_" <> T.pack (show i),+ "data_" <> T.pack (show i))+ | i <- [1 :: Int32 .. 1000]+ ]+ pure ()
+ bench/BenchCodec.hs view
@@ -0,0 +1,141 @@+module BenchCodec (benchmarks) where++import Criterion.Main+import Data.ByteString qualified as BS+import Data.Int (Int16, Int32, Int64)+import Data.Scientific (Scientific, scientific)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Time+ ( Day+ , LocalTime (..)+ , TimeOfDay (..)+ , UTCTime (..)+ , fromGregorian+ , secondsToDiffTime+ )+import Data.Vector qualified as V+import Valiant.Binary.Array (pgDecodeArray, pgEncodeArray)+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import Data.Time (TimeZone, minutesToTimeZone, ZonedTime (..), utcToZonedTime)+import Valiant.Binary.HStore (PgHStore (..), hstoreFromList)+import Valiant.Binary.Inet (PgInet, ipv4, ipv6Host)+import Valiant.Binary.Interval (PgInterval (..))+import Valiant.Binary.MacAddr (PgMacAddr, macAddr)+import Valiant.Binary.Point (PgPoint (..))+import Valiant.Binary.Scientific ()+import Valiant.Binary.Unbounded (Unbounded (..))+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid++benchmarks :: [Benchmark]+benchmarks =+ [ bgroup "encode"+ [ bench "Bool" $ nf pgEncode True+ , bench "Int16" $ nf pgEncode (12345 :: Int16)+ , bench "Int32" $ nf pgEncode (1234567 :: Int32)+ , bench "Int64" $ nf pgEncode (123456789012 :: Int64)+ , bench "Float" $ nf pgEncode (3.14 :: Float)+ , bench "Double" $ nf pgEncode (3.14159265358979 :: Double)+ , bench "Text/short" $ nf pgEncode ("hello" :: Text)+ , bench "Text/100" $ nf pgEncode (T.replicate 100 "a" :: Text)+ , bench "Text/10000" $ nf pgEncode (T.replicate 10000 "x" :: Text)+ , bench "ByteString/100" $ nf pgEncode (BS.replicate 100 42)+ , bench "ByteString/10000" $ nf pgEncode (BS.replicate 10000 42)+ , bench "Day" $ nf pgEncode (fromGregorian 2024 6 15)+ , bench "TimeOfDay" $ nf pgEncode (TimeOfDay 14 30 45)+ , bench "UTCTime" $ nf pgEncode sampleUTCTime+ , bench "LocalTime" $ nf pgEncode sampleLocalTime+ , bench "Scientific" $ nf pgEncode (scientific 314159 (-5))+ , bench "PgInterval" $ nf pgEncode (PgInterval 3600000000 30 12)+ , bench "PgInet/v4" $ nf pgEncode sampleInet4+ , bench "PgInet/v6" $ nf pgEncode sampleInet6+ , bench "PgMacAddr" $ nf pgEncode sampleMacAddr+ , bench "PgPoint" $ nf pgEncode (PgPoint 1.5 2.5)+ , bench "PgHStore/5" $ nf pgEncode sampleHStore5+ , bench "PgHStore/50" $ nf pgEncode sampleHStore50+ , bench "ZonedTime" $ nf pgEncode sampleZonedTime+ , bench "(TimeOfDay,TimeZone)" $ nf pgEncode sampleTimetz+ , bench "Unbounded UTCTime/finite" $ nf pgEncode (Finite sampleUTCTime)+ , bench "Unbounded UTCTime/inf" $ nf pgEncode (PosInfinity :: Unbounded UTCTime)+ ]+ , bgroup "decode"+ [ bench "Bool" $ nf (pgDecode @Bool) (pgEncode True)+ , bench "Int16" $ nf (pgDecode @Int16) (pgEncode (12345 :: Int16))+ , bench "Int32" $ nf (pgDecode @Int32) (pgEncode (1234567 :: Int32))+ , bench "Int64" $ nf (pgDecode @Int64) (pgEncode (123456789012 :: Int64))+ , bench "Float" $ nf (pgDecode @Float) (pgEncode (3.14 :: Float))+ , bench "Double" $ nf (pgDecode @Double) (pgEncode (3.14159265358979 :: Double))+ , bench "Text/short" $ nf (pgDecode @Text) (pgEncode ("hello" :: Text))+ , bench "Text/100" $ nf (pgDecode @Text) (pgEncode (T.replicate 100 "a"))+ , bench "Day" $ nf (pgDecode @Day) (pgEncode (fromGregorian 2024 6 15))+ , bench "UTCTime" $ nf (pgDecode @UTCTime) (pgEncode sampleUTCTime)+ , bench "Scientific" $ nf (pgDecode @Scientific) encodedScientific+ , bench "PgInterval" $ nf (pgDecode @PgInterval) (pgEncode (PgInterval 3600000000 30 12))+ , bench "PgInet/v4" $ nf (pgDecode @PgInet) (pgEncode sampleInet4)+ , bench "PgMacAddr" $ nf (pgDecode @PgMacAddr) (pgEncode sampleMacAddr)+ , bench "PgPoint" $ nf (pgDecode @PgPoint) (pgEncode (PgPoint 1.5 2.5))+ , bench "PgHStore/5" $ nf (pgDecode @PgHStore) (pgEncode sampleHStore5)+ , bench "PgHStore/50" $ nf (pgDecode @PgHStore) (pgEncode sampleHStore50)+ , bench "ZonedTime" $ nf (pgDecode @ZonedTime) (pgEncode sampleZonedTime)+ , bench "(TimeOfDay,TimeZone)" $ nf (pgDecode @(TimeOfDay, TimeZone)) (pgEncode sampleTimetz)+ , bench "Unbounded UTCTime" $ nf (pgDecode @(Unbounded UTCTime)) (pgEncode (Finite sampleUTCTime))+ ]+ , bgroup "array/encode"+ [ bench "Int32/10" $ nf (pgEncodeArray oidInt4) (V.fromList [1..10 :: Int32])+ , bench "Int32/100" $ nf (pgEncodeArray oidInt4) (V.fromList [1..100 :: Int32])+ , bench "Int32/1000" $ nf (pgEncodeArray oidInt4) (V.fromList [1..1000 :: Int32])+ , bench "Text/100" $ nf (pgEncodeArray oidText) (V.replicate 100 ("hello" :: Text))+ ]+ , bgroup "array/decode"+ [ bench "Int32/10" $ nf (pgDecodeArray @Int32) encodedArr10+ , bench "Int32/100" $ nf (pgDecodeArray @Int32) encodedArr100+ , bench "Int32/1000" $ nf (pgDecodeArray @Int32) encodedArr1000+ , bench "Text/100" $ nf (pgDecodeArray @Text) encodedTextArr100+ ]+ ]++sampleUTCTime :: UTCTime+sampleUTCTime = UTCTime (fromGregorian 2024 6 15) (secondsToDiffTime (12 * 3600 + 30 * 60))++sampleLocalTime :: LocalTime+sampleLocalTime = LocalTime (fromGregorian 2024 6 15) (TimeOfDay 14 30 0)++encodedScientific :: BS.ByteString+encodedScientific = pgEncode (scientific 314159 (-5))++encodedArr10 :: BS.ByteString+encodedArr10 = pgEncodeArray oidInt4 (V.fromList [1..10 :: Int32])++encodedArr100 :: BS.ByteString+encodedArr100 = pgEncodeArray oidInt4 (V.fromList [1..100 :: Int32])++encodedArr1000 :: BS.ByteString+encodedArr1000 = pgEncodeArray oidInt4 (V.fromList [1..1000 :: Int32])++encodedTextArr100 :: BS.ByteString+encodedTextArr100 = pgEncodeArray oidText (V.replicate 100 ("hello" :: Text))++sampleInet4 :: PgInet+sampleInet4 = ipv4 192 168 1 0 24++sampleInet6 :: PgInet+sampleInet6 = ipv6Host (BS.pack [0x20,0x01,0x0d,0xb8,0,0,0,0,0,0,0,0,0,0,0,1])++sampleMacAddr :: PgMacAddr+sampleMacAddr = macAddr 0x08 0x00 0x2b 0x01 0x02 0x03++sampleHStore5 :: PgHStore+sampleHStore5 = hstoreFromList+ [("k1","v1"),("k2","v2"),("k3","v3"),("k4","v4"),("k5","v5")]++sampleHStore50 :: PgHStore+sampleHStore50 = hstoreFromList+ [("key_" <> T.pack (show i), "value_" <> T.pack (show i)) | i <- [1::Int .. 50]]++sampleZonedTime :: ZonedTime+sampleZonedTime = utcToZonedTime (minutesToTimeZone 0) sampleUTCTime++sampleTimetz :: (TimeOfDay, TimeZone)+sampleTimetz = (TimeOfDay 14 30 45, minutesToTimeZone (-300))
+ bench/BenchConcurrent.hs view
@@ -0,0 +1,126 @@+{-# LANGUAGE BangPatterns #-}++-- | Concurrent benchmarks that demonstrate automatic pipelining.+--+-- When N green threads submit queries on the same connection, the+-- writer thread coalesces their messages into a single sendMany+-- syscall and the reader demuxes responses by FIFO order.+-- This is where the async split architecture pays off.+module BenchConcurrent (benchmarks) where++import Control.Concurrent.Async (mapConcurrently)+import Criterion.Main+import Data.Int (Int32, Int64)+import Data.IORef+import Data.Text (Text)+import Valiant+import System.IO.Unsafe (unsafePerformIO)+import TestSupport (requireDatabaseUrl, insertBulkUsers)++{-# NOINLINE globalConn #-}+globalConn :: IORef (Maybe Connection)+globalConn = unsafePerformIO (newIORef Nothing)++{-# NOINLINE globalPool #-}+globalPool :: IORef (Maybe Pool)+globalPool = unsafePerformIO (newIORef Nothing)++getConn :: IO Connection+getConn = do+ mc <- readIORef globalConn+ case mc of+ Just c -> pure c+ Nothing -> do+ url <- requireDatabaseUrl+ c <- connectString url+ _ <- simpleQuery c "CREATE TABLE IF NOT EXISTS users (id SERIAL PRIMARY KEY, name TEXT NOT NULL, email TEXT, is_active BOOLEAN NOT NULL DEFAULT true, created_at TIMESTAMPTZ NOT NULL DEFAULT now())"+ (rows, _) <- simpleQuery c "SELECT count(*) FROM users"+ case rows of+ [[Just "0"]] -> insertBulkUsers c 1000+ _ -> pure ()+ writeIORef globalConn (Just c)+ pure c++getPool :: IO Pool+getPool = do+ mp <- readIORef globalPool+ case mp of+ Just p -> pure p+ Nothing -> do+ url <- requireDatabaseUrl+ p <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 8+ , poolAcquireTimeout = 10+ }+ writeIORef globalPool (Just p)+ pure p++benchmarks :: [Benchmark]+benchmarks =+ [ bgroup "concurrent"+ [ bgroup "single-conn"+ -- N threads sharing ONE connection — the async split coalesces+ -- their messages into batched sends automatically.+ [ bench "1 thread x 100 queries" $+ whnfIO (concurrentOnConn 1 100)+ , bench "4 threads x 100 queries" $+ whnfIO (concurrentOnConn 4 100)+ , bench "16 threads x 100 queries" $+ whnfIO (concurrentOnConn 16 100)+ , bench "32 threads x 100 queries" $+ whnfIO (concurrentOnConn 32 100)+ ]+ , bgroup "pool"+ -- N threads going through the pool — more realistic production usage.+ [ bench "1 thread x 100 queries" $+ whnfIO (concurrentOnPool 1 100)+ , bench "4 threads x 100 queries" $+ whnfIO (concurrentOnPool 4 100)+ , bench "16 threads x 100 queries" $+ whnfIO (concurrentOnPool 16 100)+ , bench "32 threads x 100 queries" $+ whnfIO (concurrentOnPool 32 100)+ ]+ ]+ ]++-- | N threads each run M queries on the SAME connection.+concurrentOnConn :: Int -> Int -> IO ()+concurrentOnConn nThreads nQueries = do+ conn <- getConn+ _ <- mapConcurrently (\_ -> runQueries conn nQueries) [1 .. nThreads]+ pure ()++-- | N threads each acquire from pool and run M queries.+concurrentOnPool :: Int -> Int -> IO ()+concurrentOnPool nThreads nQueries = do+ pool <- getPool+ _ <- mapConcurrently (\_ -> withResource pool $ \conn -> runQueries conn nQueries) [1 .. nThreads]+ pure ()++-- | Run a mix of queries on a connection.+runQueries :: Connection -> Int -> IO ()+runQueries conn n = go n+ where+ go 0 = pure ()+ go !i = do+ -- Alternate between a simple lookup and a count+ if even i+ then do+ _ <- fetchOne conn stmtSelectById (fromIntegral (i `mod` 1000 + 1) :: Int32)+ pure ()+ else do+ !_ <- fetchScalar conn stmtCount ()+ pure ()+ go (i - 1)++-- Statements ----------------------------------------------------------------++stmtSelectById :: Statement Int32 (Int32, Maybe Text)+stmtSelectById = mkStatement+ "SELECT id, name FROM users WHERE id = $1"+ [23] ["id", "name"] "<bench>"++stmtCount :: Statement () Int64+stmtCount = mkStatement "SELECT count(*) FROM users" [] ["count"] "<bench>"
+ bench/BenchPool.hs view
@@ -0,0 +1,81 @@+{-# LANGUAGE BangPatterns #-}++-- | Pool-specific benchmarks: size scaling, contention, recycling methods.+module BenchPool (benchmarks) where++import Control.Concurrent.Async (mapConcurrently)+import Criterion.Main+import Valiant+import TestSupport (requireDatabaseUrl)++benchmarks :: [Benchmark]+benchmarks =+ [ bgroup "pool"+ [ bgroup "acquire-release"+ -- Measure pure pool overhead: acquire a connection, do a trivial+ -- query, release. Varies pool size to show scaling.+ [ bench "pool-size-1" $ whnfIO (poolAcquireRelease 1)+ , bench "pool-size-4" $ whnfIO (poolAcquireRelease 4)+ , bench "pool-size-16" $ whnfIO (poolAcquireRelease 16)+ ]+ , bgroup "contention"+ -- N threads compete for a pool of size M. Shows wait time under load.+ [ bench "8 threads, pool-4" $ whnfIO (poolContention 8 4 10)+ , bench "16 threads, pool-4" $ whnfIO (poolContention 16 4 10)+ , bench "32 threads, pool-4" $ whnfIO (poolContention 32 4 10)+ , bench "32 threads, pool-8" $ whnfIO (poolContention 32 8 10)+ , bench "32 threads, pool-16" $ whnfIO (poolContention 32 16 10)+ , bench "32 threads, pool-32" $ whnfIO (poolContention 32 32 10)+ ]+ , bgroup "recycling"+ -- Compare recycling methods: Fast (no check), Verified (empty query+ -- if idle > threshold), Clean (DISCARD ALL).+ [ bench "RecycleFast" $ whnfIO (poolRecycling RecycleFast)+ , bench "RecycleVerified" $ whnfIO (poolRecycling RecycleVerified)+ , bench "RecycleClean" $ whnfIO (poolRecycling RecycleClean)+ ]+ ]+ ]++-- | Acquire a connection from a fresh pool, run SELECT 1, release.+poolAcquireRelease :: Int -> IO ()+poolAcquireRelease size = do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = size+ , poolAcquireTimeout = 5+ }+ _ <- withResource pool $ \conn ->+ simpleQuery conn "SELECT 1"+ closePool pool++-- | N threads each acquire from pool (size M) and run K queries.+poolContention :: Int -> Int -> Int -> IO ()+poolContention nThreads poolSize nQueries = do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = poolSize+ , poolAcquireTimeout = 10+ }+ _ <- mapConcurrently (\_ ->+ withResource pool $ \conn ->+ mapM_ (\_ -> do _ <- simpleQuery conn "SELECT 1"; pure ()) [1 :: Int .. nQueries]+ ) [1 .. nThreads]+ closePool pool++-- | Benchmark different recycling methods.+poolRecycling :: RecyclingMethod -> IO ()+poolRecycling method = do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ , poolRecyclingMethod = method+ , poolHealthCheckAge = 0 -- force check every time for RecycleVerified+ }+ -- Acquire and release 10 times to exercise recycling+ mapM_ (\_ -> do _ <- withResource pool $ \conn -> simpleQuery conn "SELECT 1"; pure ()) [1 :: Int .. 10]+ closePool pool
+ bench/BenchQuery.hs view
@@ -0,0 +1,211 @@+{-# LANGUAGE BangPatterns #-}++module BenchQuery (benchmarks) where++import Criterion.Main+import Data.Int (Int32, Int64)+import Data.IORef+import Data.Text (Text)+import Data.Text qualified as T+import Valiant+import System.IO.Unsafe (unsafePerformIO)+import TestSupport++-- | Global connection and pool, initialized once via unsafePerformIO.+-- This avoids NFData requirements on Connection/Pool.+{-# NOINLINE globalConn #-}+globalConn :: IORef (Maybe Connection)+globalConn = unsafePerformIO (newIORef Nothing)++{-# NOINLINE globalPool #-}+globalPool :: IORef (Maybe Pool)+globalPool = unsafePerformIO (newIORef Nothing)++getConn :: IO Connection+getConn = do+ mc <- readIORef globalConn+ case mc of+ Just c -> pure c+ Nothing -> do+ url <- requireDatabaseUrl+ c <- connectString url+ writeIORef globalConn (Just c)+ -- Ensure schema + seed data+ _ <- simpleQuery c "CREATE TABLE IF NOT EXISTS users (id SERIAL PRIMARY KEY, name TEXT NOT NULL, email TEXT, is_active BOOLEAN NOT NULL DEFAULT true, created_at TIMESTAMPTZ NOT NULL DEFAULT now())"+ _ <- simpleQuery c "TRUNCATE TABLE users RESTART IDENTITY"+ insertBulkUsers c 10000+ pure c++getPool :: IO Pool+getPool = do+ mp <- readIORef globalPool+ case mp of+ Just p -> pure p+ Nothing -> do+ _ <- getConn -- ensure schema + seed data exist+ url <- requireDatabaseUrl+ p <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ }+ writeIORef globalPool (Just p)+ pure p++benchmarks :: [Benchmark]+benchmarks =+ [ bgroup "query"+ [ bench "SELECT 1 (simple)" $+ whnfIO (getConn >>= \c -> simpleQuery c "SELECT 1")++ , bench "SELECT 1 (extended/prepared)" $+ whnfIO (getConn >>= \c -> fetchScalar c stmtSelectLiteral ())++ , bench "fetchOne by PK" $+ whnfIO (getConn >>= \c -> fetchOne c stmtSelectById (1 :: Int32))++ , bench "fetchAll 5 rows" $+ whnfIO (getConn >>= \c -> fetchAll c stmtListFive ())++ , bench "fetchAll 1000 rows" $+ whnfIO (getConn >>= \c -> fetchAll c stmtListAll1k ())++ , bench "fetchScalar COUNT" $+ whnfIO (getConn >>= \c -> fetchScalar c stmtCount ())++ , bench "execute INSERT (in txn, rolled back)" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx -> do+ _ <- execute (txConn tx) stmtInsertUser ("bench_user", Just ("b@t.com" :: Text))+ -- We let the transaction commit; the user is real but harmless+ pure ()++ , bench "pool acquire/release (SELECT 1)" $+ whnfIO $ do+ p <- getPool+ withResource p $ \c -> simpleQuery c "SELECT 1"++ , bench "fetchAllVec 1000 rows" $+ whnfIO (getConn >>= \c -> fetchAllVec c stmtListAll1k ())++ , bench "executeWithFold 1000 rows" $+ whnfIO (getConn >>= \c -> executeWithFold c stmtListAll1k () (RowFold (0 :: Int) (\n _ -> n + 1)))++ , bench "forEach 1000 rows" $+ whnfIO (getConn >>= \c -> forEach c stmtListAll1k () (\_ -> pure ()))++ , bench "executeBatch 100 inserts (pipelined)" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx -> do+ _ <- executeBatch (txConn tx) stmtInsertUser+ [ ("bench_" <> T.pack (show i), Just ("b@t.com" :: Text))+ | i <- [1 :: Int .. 100]+ ]+ pure ()++ , bench "executeBatch 1000 inserts (pipelined)" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx -> do+ _ <- executeBatch (txConn tx) stmtInsertUser+ [ ("bench_" <> T.pack (show i), Just ("b@t.com" :: Text))+ | i <- [1 :: Int .. 1000]+ ]+ pure ()++ , bench "executeBatch 5000 inserts (pipelined)" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx -> do+ _ <- executeBatch (txConn tx) stmtInsertUser+ [ ("bench_" <> T.pack (show i), Just ("b@t.com" :: Text))+ | i <- [1 :: Int .. 5000]+ ]+ pure ()++ , bench "transaction overhead (BEGIN+COMMIT)" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx ->+ fetchScalar (txConn tx) stmtCount ()++ , bench "cursor 1k rows, batch 100" $+ whnfIO (cursorDrain stmtListAll1k 100)+ , bench "cursor 1k rows, batch 500" $+ whnfIO (cursorDrain stmtListAll1k 500)+ , bench "cursor 1k rows, batch 1000" $+ whnfIO (cursorDrain stmtListAll1k 1000)+ , bench "cursor 10k rows, batch 100" $+ whnfIO (cursorDrain stmtListAll10k 100)+ , bench "cursor 10k rows, batch 1000" $+ whnfIO (cursorDrain stmtListAll10k 1000)+ , bench "cursor 10k rows, batch 10000" $+ whnfIO (cursorDrain stmtListAll10k 10000)++ -- fetchAllCursor variants (decode + materialise the full result list).+ -- Compared against cursorDrain (which only counts), the gap is+ -- decode + accumulation cost.+ , bench "fetchAllCursor 1k rows, batch 100" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx ->+ length <$> fetchAllCursor (txConn tx) stmtListAll1k () 100+ , bench "fetchAllCursor 10k rows, batch 100" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx ->+ length <$> fetchAllCursor (txConn tx) stmtListAll10k () 100+ , bench "fetchAllCursor 10k rows, batch 1000" $+ whnfIO $ do+ p <- getPool+ withTransaction p $ \tx ->+ length <$> fetchAllCursor (txConn tx) stmtListAll10k () 1000+ ]+ ]++-- | Drain a cursor with the given batch size. Returns the row count so+-- the benchmark can't optimize the fetch away.+cursorDrain :: Statement () (Int32, Text) -> Int -> IO Int+cursorDrain !stmt !batch = do+ p <- getPool+ withTransaction p $ \tx ->+ withCursor (txConn tx) stmt () batch $ \cs -> do+ let loop !n = do+ rows <- fetchBatch cs batch+ if null rows then pure n else loop (n + length rows)+ loop 0++-- Statements ----------------------------------------------------------------++stmtSelectLiteral :: Statement () Int32+stmtSelectLiteral = mkStatement "SELECT 1::int4" [] ["?column?"] "<bench>"++stmtSelectById :: Statement Int32 (Int32, Text, Maybe Text)+stmtSelectById = mkStatement+ "SELECT id, name, email FROM users WHERE id = $1"+ [23] ["id", "name", "email"] "<bench>"++stmtListFive :: Statement () (Int32, Text)+stmtListFive = mkStatement+ "SELECT id, name FROM users ORDER BY id LIMIT 5"+ [] ["id", "name"] "<bench>"++stmtListAll1k :: Statement () (Int32, Text)+stmtListAll1k = mkStatement+ "SELECT id, name FROM users ORDER BY id LIMIT 1000"+ [] ["id", "name"] "<bench>"++stmtListAll10k :: Statement () (Int32, Text)+stmtListAll10k = mkStatement+ "SELECT id, name FROM users ORDER BY id LIMIT 10000"+ [] ["id", "name"] "<bench>"++stmtInsertUser :: Statement (Text, Maybe Text) ()+stmtInsertUser = mkStatement+ "INSERT INTO users (name, email) VALUES ($1, $2)"+ [25, 25] [] "<bench>"++stmtCount :: Statement () Int64+stmtCount = mkStatement "SELECT count(*) FROM users" [] ["count"] "<bench>"
+ bench/Main.hs view
@@ -0,0 +1,22 @@+module Main where++import Criterion.Main+import qualified BenchAsyncpgCompare+import qualified BenchCodec+import qualified BenchConcurrent+import qualified BenchPool+import qualified BenchQuery+import System.Environment (lookupEnv)++main :: IO ()+main = do+ mUrl <- lookupEnv "DATABASE_URL"+ let dbBenches = case mUrl of+ Nothing -> []+ Just _ -> BenchQuery.benchmarks+ <> BenchConcurrent.benchmarks+ <> BenchPool.benchmarks+ <> BenchAsyncpgCompare.benchmarks+ defaultMain $+ [ bgroup "codec" BenchCodec.benchmarks+ ] <> dbBenches
+ bench/TestSupport.hs view
@@ -0,0 +1,107 @@+module TestSupport+ ( withTestConnection+ , withTestPool+ , requireDatabaseUrl+ , createSchema+ , dropSchema+ , withSchema+ , insertTestUsers+ , insertBulkUsers+ ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Valiant+import System.Environment (lookupEnv)++-- Connection helpers --------------------------------------------------------++requireDatabaseUrl :: IO ByteString+requireDatabaseUrl = do+ mUrl <- lookupEnv "DATABASE_URL"+ case mUrl of+ Just url -> pure (BS8.pack url)+ Nothing -> error+ "DATABASE_URL is not set.\n\n\+ \ Run: eval $(scripts/pg-setup.sh)\n\+ \ Or: export DATABASE_URL=postgres://user:pass@localhost:5432/mydb\n"++withTestConnection :: (Connection -> IO a) -> IO a+withTestConnection action = do+ url <- requireDatabaseUrl+ conn <- connectString url+ result <- action conn+ close conn+ pure result++withTestPool :: (Pool -> IO a) -> IO a+withTestPool action = do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ }+ result <- action pool+ closePool pool+ pure result++-- Schema --------------------------------------------------------------------++createSchema :: Connection -> IO ()+createSchema conn = do+ _ <- simpleQuery conn+ "CREATE TABLE IF NOT EXISTS users (\+ \ id SERIAL PRIMARY KEY,\+ \ name TEXT NOT NULL,\+ \ email TEXT,\+ \ is_active BOOLEAN NOT NULL DEFAULT true,\+ \ created_at TIMESTAMPTZ NOT NULL DEFAULT now()\+ \)"+ _ <- simpleQuery conn+ "CREATE TABLE IF NOT EXISTS posts (\+ \ id SERIAL PRIMARY KEY,\+ \ author_id INTEGER NOT NULL REFERENCES users(id),\+ \ title TEXT NOT NULL,\+ \ body TEXT,\+ \ published_at TIMESTAMPTZ,\+ \ created_at TIMESTAMPTZ NOT NULL DEFAULT now()\+ \)"+ pure ()++dropSchema :: Connection -> IO ()+dropSchema conn = do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS posts CASCADE"+ _ <- simpleQuery conn "DROP TABLE IF EXISTS users CASCADE"+ pure ()++withSchema :: Connection -> IO a -> IO a+withSchema conn action = do+ dropSchema conn+ createSchema conn+ result <- action+ dropSchema conn+ pure result++-- Fixtures ------------------------------------------------------------------++insertTestUsers :: Connection -> IO ()+insertTestUsers conn = do+ _ <- simpleQuery conn+ "INSERT INTO users (name, email) VALUES \+ \('Alice', 'alice@example.com'),\+ \('Bob', 'bob@example.com'),\+ \('Carol', NULL),\+ \('Dave', 'dave@example.com'),\+ \('Eve', 'eve@example.com')"+ pure ()++insertBulkUsers :: Connection -> Int -> IO ()+insertBulkUsers conn n = do+ let values = BS8.intercalate ","+ [ "('user_" <> BS8.pack (show i) <> "', 'user" <> BS8.pack (show i) <> "@test.com')"+ | i <- [1 .. n]+ ]+ sql = "INSERT INTO users (name, email) VALUES " <> values+ _ <- simpleQuery conn sql+ pure ()
+ integration/ChaosSpec.hs view
@@ -0,0 +1,136 @@+-- | Chaos tests: verify correct behavior under concurrent stress and+-- connection failure scenarios.+module ChaosSpec (spec) where++import Control.Concurrent (threadDelay)+import Control.Concurrent.Async (mapConcurrently_)+import Control.Exception (SomeException, try)+import Data.ByteString.Char8 qualified as BS8+import Valiant+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ poolStressSpec+ terminationRecoverySpec++------------------------------------------------------------------------+-- Pool stress tests+------------------------------------------------------------------------++poolStressSpec :: Spec+poolStressSpec = describe "Pool stress" $ do+ it "handles 100 rapid concurrent acquire/release without corruption" $ do+ withTestPool $ \pool -> do+ mapConcurrently_+ (\_ -> do+ _ <- try @SomeException $ withResource pool $ \conn ->+ simpleQuery conn "SELECT 1"+ pure ()+ )+ [1 :: Int .. 100]++ -- Pool should still be functional+ withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ rows `shouldBe` [[Just "1"]]++ it "statement cache handles 50 concurrent distinct queries" $ do+ withTestPool $ \pool -> do+ mapConcurrently_+ (\i -> withResource pool $ \conn -> do+ let q = "SELECT " <> BS8.pack (show i)+ _ <- simpleQuery conn q+ pure ()+ )+ [1 :: Int .. 50]++ withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ rows `shouldBe` [[Just "1"]]++ it "mixed queries and transactions under contention" $ do+ withTestPool $ \pool -> do+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS chaos_mix CASCADE"+ _ <- simpleQuery conn "CREATE TABLE chaos_mix (id SERIAL PRIMARY KEY, val INT NOT NULL)"+ pure ()++ -- Mix of reads and writes across threads+ mapConcurrently_+ (\i -> do+ if even i+ then do+ -- Write path+ _ <- try @SomeException $ withTransaction pool $ \tx ->+ simpleQuery (txConn tx) ("INSERT INTO chaos_mix (val) VALUES (" <> BS8.pack (show i) <> ")")+ pure ()+ else do+ -- Read path+ _ <- try @SomeException $ withResource pool $ \conn ->+ simpleQuery conn "SELECT count(*) FROM chaos_mix"+ pure ()+ )+ [1 :: Int .. 40]++ withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT count(*) FROM chaos_mix"+ case rows of+ [[Just n]] -> (read (BS8.unpack n) :: Int) `shouldSatisfy` (> 0)+ _ -> expectationFailure $ "Unexpected: " <> show rows+ _ <- simpleQuery conn "DROP TABLE chaos_mix"+ pure ()++------------------------------------------------------------------------+-- Connection termination recovery+------------------------------------------------------------------------++terminationRecoverySpec :: Spec+terminationRecoverySpec = describe "Termination recovery" $ do+ it "pool recovers when idle connection is terminated by server" $ do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 1+ , poolAcquireTimeout = 5+ , poolRecyclingMethod = RecycleVerified+ , poolHealthCheckAge = 0 -- always verify+ }++ -- Populate the pool with one connection+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "SELECT 1"+ pure ()++ -- Terminate the idle connection from outside.+ -- Use a fresh connection (not from the pool) to issue the kill.+ withTestConnection $ \killer -> do+ -- Find PIDs of all connections from our test user+ (rows, _) <- simpleQuery killer+ "SELECT pid FROM pg_stat_activity WHERE usename = 'valiant_test' AND pid != pg_backend_pid() AND state = 'idle'"+ mapM_ (\row -> case row of+ [Just pid] -> do+ _ <- simpleQuery killer ("SELECT pg_terminate_backend(" <> pid <> ")")+ pure ()+ _ -> pure ()+ ) rows++ -- Wait for termination+ threadDelay 300000++ -- Pool should detect dead connection (RecycleVerified + healthCheckAge=0)+ -- and create a fresh one. May need a retry since first acquire gets the+ -- dead connection which fails health check.+ let tryQuery = withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ pure rows+ result1 <- try @SomeException tryQuery+ result2 <- try @SomeException tryQuery+ let succeeded = case (result1, result2) of+ (Right [[Just "1"]], _) -> True+ (_, Right [[Just "1"]]) -> True+ _ -> False+ succeeded `shouldBe` True++ closePool pool
+ integration/ConnectionFeaturesSpec.hs view
@@ -0,0 +1,165 @@+module ConnectionFeaturesSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int32)+import Data.IORef+import Data.Text (Text)+import Data.Maybe (isJust)+import Data.Vector qualified as V+import Valiant+import PgWire.Connection+import PgWire.Connection.Config (parseConnString)+import PgWire.Protocol.Backend (TxStatus (..))+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ describe "connectionStatus" $ do+ it "returns True for a live connection" $ do+ withTestConnection $ \conn -> do+ status <- connectionStatus conn+ status `shouldBe` True++ describe "transactionStatus" $ do+ it "returns TxIdle outside a transaction" $ do+ withTestConnection $ \conn -> do+ status <- transactionStatus conn+ status `shouldBe` TxIdle++ describe "parameterStatus" $ do+ it "returns server_version" $ do+ withTestConnection $ \conn -> do+ version <- parameterStatus conn "server_version"+ version `shouldSatisfy` isJust++ it "returns server_encoding" $ do+ withTestConnection $ \conn -> do+ enc <- parameterStatus conn "server_encoding"+ enc `shouldSatisfy` isJust++ it "returns Nothing for unknown params" $ do+ withTestConnection $ \conn -> do+ val <- parameterStatus conn "nonexistent_param"+ val `shouldBe` Nothing++ describe "serverVersion" $ do+ it "returns a positive integer" $ do+ withTestConnection $ \conn -> do+ version <- serverVersion conn+ case version of+ Just v -> v `shouldSatisfy` (> 100000) -- at least PG 10+ Nothing -> expectationFailure "Expected version"++ describe "backendPid" $ do+ it "returns a positive PID" $ do+ withTestConnection $ \conn -> do+ let pid = backendPid conn+ pid `shouldSatisfy` (> 0)++ describe "reset" $ do+ it "reconnects and returns a working connection" $ do+ url <- requireDatabaseUrl+ conn <- connectString url+ conn2 <- reset conn+ -- conn is now closed; conn2 is the new connection+ (rows, _) <- simpleQuery conn2 "SELECT 1"+ length rows `shouldBe` 1+ close conn2++ describe "checkNotification" $ do+ it "returns Nothing when no notifications pending" $ do+ withTestConnection $ \conn -> do+ result <- checkNotification conn+ result `shouldBe` Nothing++ describe "describePrepared" $ do+ it "describes a prepared statement" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ -- Prepare a statement manually+ _ <- simpleQuery conn "PREPARE test_stmt (int4) AS SELECT id, name FROM users WHERE id = $1"+ (params, cols) <- describePrepared conn "test_stmt"+ -- Should have 1 param (int4 = OID 23)+ length (pdOids params) `shouldBe` 1+ -- Should have 2 columns+ length (cdFields cols) `shouldBe` 2+ _ <- simpleQuery conn "DEALLOCATE test_stmt"+ pure ()++ describe "escapeLiteral" $ do+ it "produces valid SQL when used in a query" $ do+ withTestConnection $ \conn -> do+ let escaped = escapeLiteral conn "O'Brien"+ (rows, _) <- simpleQuery conn ("SELECT " <> escaped)+ case rows of+ [[Just val]] -> val `shouldBe` "O'Brien"+ _ -> expectationFailure $ "Unexpected: " <> show rows++ describe "escapeIdentifier" $ do+ it "produces valid SQL for identifiers" $ do+ withTestConnection $ \conn -> do+ let escaped = escapeIdentifier conn "my col"+ (rows, _) <- simpleQuery conn ("SELECT 1 AS " <> escaped)+ length rows `shouldBe` 1++ describe "ping" $ do+ it "returns True for a reachable server" $ do+ url <- requireDatabaseUrl+ case parseConnString url of+ Right cfg -> do+ result <- ping cfg+ result `shouldBe` True+ Left _ -> expectationFailure "Failed to parse URL"++ describe "encryptPassword" $ do+ it "produces md5-prefixed hash" $ do+ let hash = encryptPassword "testuser" "testpass"+ BS.isPrefixOf "md5" hash `shouldBe` True+ BS.length hash `shouldBe` 35 -- "md5" + 32 hex chars++ it "produces different hashes for different passwords" $ do+ let h1 = encryptPassword "user" "pass1"+ h2 = encryptPassword "user" "pass2"+ h1 `shouldNotBe` h2++ describe "setTraceHandler" $ do+ it "captures send and recv traffic" $ do+ withTestConnection $ \conn -> do+ sendRef <- newIORef (0 :: Int)+ recvRef <- newIORef (0 :: Int)+ setTraceHandler conn $ \isSend _bytes ->+ if isSend+ then modifyIORef' sendRef (+ 1)+ else modifyIORef' recvRef (+ 1)+ _ <- simpleQuery conn "SELECT 1"+ sends <- readIORef sendRef+ recvs <- readIORef recvRef+ sends `shouldSatisfy` (> 0)+ recvs `shouldSatisfy` (> 0)++ describe "executeWithFold" $ do+ it "folds rows in constant memory" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ let countStmt :: Statement () (Int32, Text, Maybe Text)+ countStmt = mkStatement+ "SELECT id, name, email FROM users"+ [] ["id", "name", "email"] "<test>"+ count <- executeWithFold conn countStmt () $+ RowFold (0 :: Int) (\acc _ -> acc + 1)+ count `shouldBe` 5++ describe "copyInBinary" $ do+ it "bulk inserts via binary COPY" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ result <- copyInBinary conn+ "COPY users (name, email) FROM STDIN WITH (FORMAT binary)"+ 2+ $ \sendRow -> do+ sendRow (V.fromList [Just (pgEncode ("BinAlice" :: BS.ByteString)), Just (pgEncode ("alice@bin.com" :: BS.ByteString))])+ sendRow (V.fromList [Just (pgEncode ("BinBob" :: BS.ByteString)), Nothing])+ copyRows result `shouldBe` 2+ (rows, _) <- simpleQuery conn "SELECT count(*) FROM users"+ case rows of+ [[Just n]] -> n `shouldBe` "2"+ _ -> expectationFailure $ "Expected 2, got: " <> show rows
+ integration/ConnectionSpec.hs view
@@ -0,0 +1,37 @@+module ConnectionSpec (spec) where++import Valiant+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ describe "connectString" $ do+ it "connects and disconnects cleanly" $ do+ withTestConnection $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ length rows `shouldBe` 1++ it "returns server version via parameter status" $ do+ withTestConnection $ \conn -> do+ (rows, _) <- simpleQuery conn "SHOW server_version"+ length rows `shouldBe` 1++ describe "simpleQuery" $ do+ it "handles empty query" $ do+ withTestConnection $ \conn -> do+ (rows, _) <- simpleQuery conn ""+ rows `shouldBe` []++ it "returns multiple rows" $ do+ withTestConnection $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT generate_series(1, 5)"+ length rows `shouldBe` 5++ it "returns NULL values" $ do+ withTestConnection $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT NULL::text"+ length rows `shouldBe` 1+ case rows of+ [[Nothing]] -> pure ()+ _ -> expectationFailure $ "Expected [[Nothing]], got: " <> show rows
+ integration/CopySpec.hs view
@@ -0,0 +1,35 @@+module CopySpec (spec) where++import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Valiant+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ describe "copyIn" $ do+ it "bulk inserts via COPY FROM STDIN" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ result <- copyIn conn+ "COPY users (name, email) FROM STDIN WITH (FORMAT csv)"+ $ \send -> do+ send "Alice,alice@example.com\n"+ send "Bob,bob@example.com\n"+ send "Carol,\n"+ copyRows result `shouldBe` 3++ describe "copyOut" $ do+ it "exports via COPY TO STDOUT" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ _ <- simpleQuery conn+ "INSERT INTO users (name, email) VALUES ('Alice', 'alice@example.com'), ('Bob', 'bob@example.com')"++ chunks <- newIORef ([] :: [BS8.ByteString])+ _ <- copyOut conn+ "COPY users (name, email) TO STDOUT WITH (FORMAT csv)"+ $ \chunk -> modifyIORef' chunks (chunk :)++ allData <- BS8.concat . reverse <$> readIORef chunks+ BS8.isInfixOf "Alice" allData `shouldBe` True+ BS8.isInfixOf "Bob" allData `shouldBe` True
+ integration/ExecuteSpec.hs view
@@ -0,0 +1,131 @@+module ExecuteSpec (spec) where++import Control.Exception (try)+import Data.Int (Int32, Int64)+import Data.Text (Text)+import Data.Vector.Unboxed qualified as U+import Valiant+import PgWire.Error (PgWireError (..))+import TestSupport+import Test.Hspec++stmtSelectOne :: Statement Int32 (Int32, Text, Maybe Text)+stmtSelectOne = mkStatement+ "SELECT id, name, email FROM users WHERE id = $1"+ [23] ["id", "name", "email"] "<test>"++stmtListAll :: Statement () (Int32, Text)+stmtListAll = mkStatement+ "SELECT id, name FROM users ORDER BY id"+ [] ["id", "name"] "<test>"++stmtListIds :: Statement () Int32+stmtListIds = mkStatement+ "SELECT id FROM users ORDER BY id"+ [] ["id"] "<test>"++stmtInsert :: Statement (Text, Maybe Text) ()+stmtInsert = mkStatement+ "INSERT INTO users (name, email) VALUES ($1, $2)"+ [25, 25] [] "<test>"++stmtCount :: Statement () Int64+stmtCount = mkStatement+ "SELECT count(*) FROM users"+ [] ["count"] "<test>"++stmtInsertReturningId :: Statement (Text, Maybe Text) Int32+stmtInsertReturningId = mkStatement+ "INSERT INTO users (name, email) VALUES ($1, $2) RETURNING id"+ [25, 25] ["id"] "<test>"++stmtInsertMultiReturningId :: Statement () Int32+stmtInsertMultiReturningId = mkStatement+ "INSERT INTO users (name, email) VALUES ('X', 'x@test.com'), ('Y', 'y@test.com'), ('Z', 'z@test.com') RETURNING id"+ [] ["id"] "<test>"++spec :: Spec+spec = do+ describe "fetchOne" $ do+ it "returns Nothing for missing row" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ result <- fetchOne conn stmtSelectOne 999+ result `shouldBe` Nothing++ it "returns Just for existing row" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ result <- fetchOne conn stmtSelectOne 1+ case result of+ Just (_, name, email) -> do+ name `shouldBe` "Alice"+ email `shouldBe` Just "alice@example.com"+ Nothing -> expectationFailure "Expected Just, got Nothing"++ describe "fetchAll" $ do+ it "returns empty list for no rows" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ rows <- fetchAll conn stmtListAll ()+ rows `shouldBe` []++ it "returns all rows" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ rows <- fetchAll conn stmtListAll ()+ length rows `shouldBe` 5++ describe "fetchAllUnboxed" $ do+ it "returns empty unboxed vector for no rows" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ ids <- fetchAllUnboxed conn stmtListIds () :: IO (U.Vector Int32)+ U.length ids `shouldBe` 0++ it "returns all rows in an unboxed vector" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ ids <- fetchAllUnboxed conn stmtListIds () :: IO (U.Vector Int32)+ U.length ids `shouldBe` 5+ U.toList ids `shouldSatisfy` all (> 0)++ describe "execute" $ do+ it "inserts a row and returns rows affected" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ n <- execute conn stmtInsert ("TestUser", Just "test@example.com")+ n `shouldBe` 1++ describe "fetchScalar" $ do+ it "returns a scalar count" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ count <- fetchScalar conn stmtCount ()+ count `shouldBe` 5++ describe "executeReturning" $ do+ it "returns row count and decoded rows for single insert" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ (n, ids) <- executeReturning conn stmtInsertReturningId ("NewUser", Just "new@example.com")+ n `shouldBe` 1+ length ids `shouldBe` 1++ it "returns row count and decoded rows for multi-value insert" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ (n, ids) <- executeReturning conn stmtInsertMultiReturningId ()+ n `shouldBe` 3+ length ids `shouldBe` 3++ describe "fetchOneOrThrow" $ do+ it "returns the value for an existing row" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ (uid, name, email) <- fetchOneOrThrow conn stmtSelectOne 1+ name `shouldBe` "Alice"+ email `shouldBe` Just "alice@example.com"+ uid `shouldSatisfy` (> 0)++ it "throws on a missing row" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ result <- try (fetchOneOrThrow conn stmtSelectOne (-1))+ case result of+ Left (DecodeError _) -> pure ()+ Left err -> expectationFailure $ "Expected DecodeError, got: " <> show err+ Right _ -> expectationFailure "Expected an exception, but got a result"
+ integration/LargeObjectSpec.hs view
@@ -0,0 +1,91 @@+module LargeObjectSpec (spec) where++import Control.Exception (bracket_)+import Data.ByteString qualified as BS+import Valiant+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ describe "loCreate / loUnlink" $ do+ it "round-trips OIDs without error" $ do+ withTestPool $ \pool -> do+ withTransaction pool $ \tx -> do+ oid <- loCreate (txConn tx)+ oid `shouldSatisfy` (> 0)+ loUnlink (txConn tx) oid++ describe "withLargeObject + loRead/loWrite" $ do+ it "writes and reads back binary data" $ do+ withTestPool $ \pool -> do+ let payload = BS.pack [0 .. 255]+ oid <- withTransaction pool $ \tx -> do+ o <- loCreate (txConn tx)+ withLargeObject (txConn tx) o WriteMode $ \fd -> do+ n <- loWrite (txConn tx) fd payload+ n `shouldBe` fromIntegral (BS.length payload)+ pure o++ readBack <- withTransaction pool $ \tx ->+ withLargeObject (txConn tx) oid ReadMode $ \fd ->+ loRead (txConn tx) fd 1024+ readBack `shouldBe` payload++ withTransaction pool $ \tx -> loUnlink (txConn tx) oid++ it "supports loSeek + partial reads" $ do+ withTestPool $ \pool -> do+ let payload = BS.pack [10, 20, 30, 40, 50, 60, 70, 80, 90, 100]+ oid <- withTransaction pool $ \tx -> do+ o <- loCreate (txConn tx)+ withLargeObject (txConn tx) o WriteMode $ \fd -> do+ _ <- loWrite (txConn tx) fd payload+ pure ()+ pure o++ chunk <- withTransaction pool $ \tx ->+ withLargeObject (txConn tx) oid ReadMode $ \fd -> do+ _ <- loSeek (txConn tx) fd 4 0 -- SEEK_SET+ loRead (txConn tx) fd 3+ chunk `shouldBe` BS.pack [50, 60, 70]++ withTransaction pool $ \tx -> loUnlink (txConn tx) oid++ it "loTell reports the current cursor position" $ do+ withTestPool $ \pool -> do+ let payload = BS.replicate 100 0xAB+ oid <- withTransaction pool $ \tx -> do+ o <- loCreate (txConn tx)+ withLargeObject (txConn tx) o WriteMode $ \fd -> do+ _ <- loWrite (txConn tx) fd payload+ pure ()+ pure o++ position <- withTransaction pool $ \tx ->+ withLargeObject (txConn tx) oid ReadMode $ \fd -> do+ _ <- loRead (txConn tx) fd 25+ loTell (txConn tx) fd+ position `shouldBe` 25++ withTransaction pool $ \tx -> loUnlink (txConn tx) oid++ describe "withLargeObject bracket safety" $ do+ it "closes the descriptor when the action throws" $ do+ withTestPool $ \pool -> do+ oid <- withTransaction pool $ \tx -> loCreate (txConn tx)+ -- The bracket should propagate exceptions while still closing+ -- the underlying lo descriptor on the server. Round-trip via a+ -- second transaction proves the OID is still readable.+ _ <- withTransaction pool $ \tx ->+ withLargeObject (txConn tx) oid WriteMode $ \fd ->+ bracket_ (pure ()) (pure ()) $ do+ _ <- loWrite (txConn tx) fd "hello"+ pure ()++ readBack <- withTransaction pool $ \tx ->+ withLargeObject (txConn tx) oid ReadMode $ \fd ->+ loRead (txConn tx) fd 5+ readBack `shouldBe` "hello"++ withTransaction pool $ \tx -> loUnlink (txConn tx) oid
+ integration/Main.hs view
@@ -0,0 +1,30 @@+module Main where++import ChaosSpec qualified+import ConnectionSpec qualified+import ConnectionFeaturesSpec qualified+import CopySpec qualified+import ExecuteSpec qualified+import LargeObjectSpec qualified+import NotifySpec qualified+import PipelineSpec qualified+import PoolSpec qualified+import SoundnessSpec qualified+import StreamingSpec qualified+import TransactionSpec qualified+import Test.Hspec++main :: IO ()+main = hspec $ do+ describe "Connection" ConnectionSpec.spec+ describe "Connection.Features" ConnectionFeaturesSpec.spec+ describe "Execute" ExecuteSpec.spec+ describe "Pool" PoolSpec.spec+ describe "Transaction" TransactionSpec.spec+ describe "Copy" CopySpec.spec+ describe "Pipeline" PipelineSpec.spec+ describe "Streaming" StreamingSpec.spec+ describe "Notify" NotifySpec.spec+ describe "LargeObject" LargeObjectSpec.spec+ describe "Soundness" SoundnessSpec.spec+ describe "Chaos" ChaosSpec.spec
+ integration/NotifySpec.hs view
@@ -0,0 +1,51 @@+module NotifySpec (spec) where++import Control.Concurrent (threadDelay)+import Valiant+import TestSupport+import Test.Hspec++-- NOTE: waitForNotification is edge-triggered by the next query; see+-- Valiant.Notify haddock. These tests send the NOTIFY first, wait a+-- beat for the server to push it onto the listener's socket, then call+-- waitForNotificationTimeout so a hung delivery fails fast instead of+-- stalling CI.+spec :: Spec+spec = do+ describe "listen / waitForNotificationTimeout" $ do+ it "delivers a notification queued before the wait" $ do+ withTestConnection $ \listenerConn -> withTestConnection $ \senderConn -> do+ listen listenerConn "valiant_test_chan"+ _ <- simpleQuery senderConn "NOTIFY valiant_test_chan, 'hello'"+ threadDelay 50000+ result <- waitForNotificationTimeout listenerConn 1.0+ case result of+ Just n -> do+ notifChannel n `shouldBe` "valiant_test_chan"+ notifPayload n `shouldBe` "hello"+ Nothing -> expectationFailure "expected notification, got timeout"+ unlisten listenerConn "valiant_test_chan"++ it "returns Nothing when no notification is sent" $ do+ withTestConnection $ \conn -> do+ listen conn "valiant_test_silent"+ result <- waitForNotificationTimeout conn 0.2+ result `shouldBe` Nothing+ unlisten conn "valiant_test_silent"++ it "ignores notifications on other channels" $ do+ withTestConnection $ \listenerConn -> withTestConnection $ \senderConn -> do+ listen listenerConn "valiant_test_a"+ _ <- simpleQuery senderConn "NOTIFY valiant_test_b, 'wrong-channel'"+ threadDelay 50000+ result <- waitForNotificationTimeout listenerConn 0.3+ result `shouldBe` Nothing+ unlisten listenerConn "valiant_test_a"++ describe "channel name escaping" $ do+ it "accepts channel names that need quoting" $ do+ withTestConnection $ \conn -> do+ -- Mixed-case names require quoting; the LISTEN/UNLISTEN helpers+ -- quote identifiers automatically.+ listen conn "MixedCase"+ unlisten conn "MixedCase"
+ integration/PipelineSpec.hs view
@@ -0,0 +1,116 @@+module PipelineSpec (spec) where++import Data.Int (Int32, Int64)+import Data.Text (Text)+import Valiant+import PgWire.Protocol.Oid (oidInt4)+import TestSupport+import Test.Hspec++stmtFetchOne :: Statement Int32 (Int32, Text, Maybe Text)+stmtFetchOne = mkStatement+ "SELECT id, name, email FROM users WHERE id = $1"+ [23] ["id", "name", "email"] "<test>"++stmtCount :: Statement () Int64+stmtCount = mkStatement+ "SELECT count(*) FROM users"+ [] ["count"] "<test>"++stmtListAll :: Statement () (Int32, Text)+stmtListAll = mkStatement+ "SELECT id, name FROM users ORDER BY id"+ [] ["id", "name"] "<test>"++spec :: Spec+spec = do+ describe "runPipeline" $ do+ it "executes multiple queries in one round-trip" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ (mUser, count) <- runPipeline conn $ (,)+ <$> pipeFetchOne stmtFetchOne 1+ <*> pipeFetchScalar stmtCount ()+ case mUser of+ Just (_, name, _) -> name `shouldBe` "Alice"+ Nothing -> expectationFailure "Expected user"+ count `shouldBe` 5++ it "handles mixed result types" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ (users, count, mMissing) <- runPipeline conn $ (,,)+ <$> pipeFetchAll stmtListAll ()+ <*> pipeFetchScalar stmtCount ()+ <*> pipeFetchOne stmtFetchOne 999+ length users `shouldBe` 5+ count `shouldBe` 5+ mMissing `shouldBe` Nothing++ it "handles empty pipeline" $ do+ withTestConnection $ \conn -> do+ result <- runPipeline conn (pure 42 :: Pipeline Int)+ result `shouldBe` 42++ it "handles single query" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ count <- runPipeline conn $ pipeFetchScalar stmtCount ()+ count `shouldBe` 5++ describe "fetchBatchOne" $ do+ it "fetches multiple rows by different params" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ results <- fetchBatchOne conn stmtFetchOne [1, 2, 3, 999]+ length results `shouldBe` 4+ -- First 3 should be Just, last should be Nothing+ case results of+ [Just (_, n1, _), Just (_, n2, _), Just (_, n3, _), Nothing] -> do+ n1 `shouldBe` "Alice"+ n2 `shouldBe` "Bob"+ n3 `shouldBe` "Carol"+ _ -> expectationFailure $ "Unexpected: " <> show (length results)++ it "handles empty list" $ do+ withTestConnection $ \conn -> do+ results <- fetchBatchOne conn stmtFetchOne ([] :: [Int32])+ results `shouldBe` []++ describe "fetchBatchAll" $ do+ it "fetches multiple result sets" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ results <- fetchBatchAll conn stmtListAll [(), ()]+ length results `shouldBe` 2+ case results of+ (r : _) -> length r `shouldBe` 5+ [] -> expectationFailure "Expected non-empty results"+ length (results !! 1) `shouldBe` 5++ describe "fetchByIds" $ do+ it "fetches multiple rows by array parameter" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ users <- fetchByIds conn+ "SELECT id, name, email FROM users WHERE id = ANY($1::int4[])"+ oidInt4+ [1 :: Int32, 2, 3]+ length (users :: [(Int32, Text, Maybe Text)]) `shouldBe` 3++ it "returns empty for no matches" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn+ users <- fetchByIds conn+ "SELECT id, name, email FROM users WHERE id = ANY($1::int4[])"+ oidInt4+ [999 :: Int32, 998]+ length (users :: [(Int32, Text, Maybe Text)]) `shouldBe` 0++ it "handles empty ID list" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ users <- fetchByIds conn+ "SELECT id, name, email FROM users WHERE id = ANY($1::int4[])"+ oidInt4+ ([] :: [Int32])+ length (users :: [(Int32, Text, Maybe Text)]) `shouldBe` 0
+ integration/PoolSpec.hs view
@@ -0,0 +1,467 @@+module PoolSpec (spec) where++import Control.Concurrent (threadDelay)+import Control.Concurrent.Async (async, mapConcurrently, wait)+import Control.Concurrent.MVar+import Control.Concurrent.STM (TVar, newTVarIO, readTVarIO, modifyTVar', atomically)+import Control.Exception (SomeException, try)+import Control.Monad (replicateM_)+import Data.IORef+import NoThunks.Class (NoThunks, ThunkInfo, noThunks)+import PgWire.Pool (pActive, pClosed, pConnMeta, pEffectiveSize, pIdle, pTotalCreated, pTotalDestroyed, pTotalTimeouts)+import Valiant+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ describe "withResource" $ do+ it "acquires and releases a connection" $ do+ withTestPool $ \pool -> do+ result <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 42"+ pure (length rows)+ result `shouldBe` 1++ it "handles concurrent access" $ do+ withTestPool $ \pool -> do+ results <- mapConcurrently (\_ -> withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ pure (length rows)) [1 :: Int .. 20]+ all (== 1) results `shouldBe` True++ it "reuses connections" $ do+ withTestPool $ \pool -> do+ pid1 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ pid2 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ pid1 `shouldBe` pid2++ describe "concurrent acquire does not exceed poolSize" $ do+ it "never creates more connections than poolSize" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 3+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Track the maximum concurrent active connections+ maxActive <- newTVarIO (0 :: Int)+ active <- newTVarIO (0 :: Int)+ _ <- mapConcurrently (\_ -> withResource pool $ \conn -> do+ atomically $ modifyTVar' active (+ 1)+ cur <- readTVarIO active+ atomically $ modifyTVar' maxActive (\m -> max m cur)+ _ <- simpleQuery conn "SELECT pg_sleep(0.05)"+ atomically $ modifyTVar' active (subtract 1)+ pure ()+ ) [1 :: Int .. 10]+ closePool pool+ maxA <- readTVarIO maxActive+ maxA `shouldSatisfy` (<= 3)++ describe "shutdown wakes blocked waiters" $ do+ it "blocked waiters get PoolClosed" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 1+ , poolAcquireTimeout = 10+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Hold the only connection+ barrier <- newEmptyMVar+ done <- newEmptyMVar+ _ <- async $ withResource pool $ \_ -> do+ putMVar barrier ()+ -- Wait until test says to continue+ takeMVar done+ -- Wait for the connection to be acquired+ takeMVar barrier+ -- Try to acquire in another thread — should block+ waiterResult <- async $ try @SomeException $ withResource pool $ \_ -> pure ()+ -- Give waiter time to block+ threadDelay 50000+ -- Close pool — should wake the waiter+ closePool pool+ putMVar done ()+ result <- wait waiterResult+ case result of+ Left e -> show e `shouldContain` "PoolClosed"+ Right () -> expectationFailure "Expected PoolClosed but got success"++ describe "poolStats" $ do+ it "returns sane values" $ do+ withTestPool $ \pool -> do+ -- Initial state: all idle after first warm-up+ _ <- withResource pool $ \conn -> simpleQuery conn "SELECT 1"+ stats <- poolStats pool+ psIdle stats `shouldSatisfy` (>= 1)+ psInUse stats `shouldBe` 0+ psWaiters stats `shouldBe` 0+ psTotalCreated stats `shouldSatisfy` (>= 1)++ it "tracks in-use connections" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ barrier <- newEmptyMVar+ done <- newEmptyMVar+ _ <- async $ withResource pool $ \_ -> do+ putMVar barrier ()+ takeMVar done+ takeMVar barrier+ stats <- poolStats pool+ psInUse stats `shouldBe` 1+ putMVar done ()+ threadDelay 50000+ stats2 <- poolStats pool+ psInUse stats2 `shouldBe` 0+ closePool pool++ describe "resize" $ do+ it "shrinks idle connections" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Create 3 idle connections+ pids <- mapConcurrently (\_ -> withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows) [1 :: Int .. 3]+ length pids `shouldBe` 3+ stats1 <- poolStats pool+ psIdle stats1 `shouldSatisfy` (>= 2)+ -- Resize down to 1+ resize pool 1+ threadDelay 50000+ stats2 <- poolStats pool+ psIdle stats2 `shouldSatisfy` (<= 1)+ psMaxSize stats2 `shouldBe` 1+ closePool pool++ it "allows growth after resize up" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 1+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ resize pool 3+ stats <- poolStats pool+ psMaxSize stats `shouldBe` 3+ -- Should be able to create multiple concurrent connections now+ results <- mapConcurrently (\_ -> withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ pure (length rows)) [1 :: Int .. 3]+ all (== 1) results `shouldBe` True+ closePool pool++ describe "recycling methods" $ do+ it "RecycleFast reuses connections without query" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 2+ , poolRecyclingMethod = RecycleFast+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ pid1 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ pid2 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ pid1 `shouldBe` pid2+ closePool pool++ it "RecycleVerified reuses recent connections without query" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 2+ , poolRecyclingMethod = RecycleVerified+ , poolHealthCheckAge = 60+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ pid1 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ -- Immediately reacquire — should skip health check (used < 60s ago)+ pid2 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ pid1 `shouldBe` pid2+ closePool pool++ it "RecycleClean resets session state" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 2+ , poolRecyclingMethod = RecycleClean+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Set a session variable+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "SET application_name TO 'test_app'"+ pure ()+ -- Reacquire — RecycleClean runs DISCARD ALL which resets session state+ appName <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SHOW application_name"+ pure rows+ -- After DISCARD ALL, application_name should be reset to default+ -- (empty string or the connConfig app name, not 'test_app')+ case appName of+ [[Just name]] -> name `shouldNotBe` "test_app"+ _ -> expectationFailure "Expected one row"+ closePool pool++ describe "lifecycle hooks" $ do+ it "fires hooks in correct order" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 2+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ events <- newIORef ([] :: [String])+ setPostCreateHook pool $ \_ -> modifyIORef' events ("create" :)+ setOnAcquireHook pool $ \_ -> modifyIORef' events ("acquire" :)+ setPreReleaseHook pool $ \_ -> modifyIORef' events ("release" :)+ -- First acquire: create + acquire+ _ <- withResource pool $ \conn -> simpleQuery conn "SELECT 1"+ -- Second acquire (reuse): acquire only+ _ <- withResource pool $ \conn -> simpleQuery conn "SELECT 1"+ evts <- reverse <$> readIORef events+ -- First use: create, acquire, release+ -- Second use: acquire, release+ evts `shouldBe` ["create", "acquire", "release", "acquire", "release"]+ closePool pool++ it "hook failure doesn't break the pool" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 2+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ setOnAcquireHook pool $ \_ -> error "hook failure"+ -- Should still work despite the hook throwing+ result <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 42"+ pure (length rows)+ result `shouldBe` 1+ closePool pool++ describe "retain" $ do+ it "filters idle connections" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Create 3 idle connections and get their PIDs+ pids <- mapConcurrently (\_ -> withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows) [1 :: Int .. 3]+ length pids `shouldBe` 3+ stats1 <- poolStats pool+ let idleBefore = psIdle stats1+ idleBefore `shouldSatisfy` (>= 2)+ -- Keep only the first connection (drop all others by returning False)+ callCount <- newIORef (0 :: Int)+ retain pool $ \_ -> do+ n <- readIORef callCount+ modifyIORef' callCount (+ 1)+ pure (n == 0) -- keep first, drop rest+ stats2 <- poolStats pool+ psIdle stats2 `shouldSatisfy` (<= 1)+ closePool pool++ describe "QueueMode" $ do+ it "QueueLIFO reuses the most recently released connection" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolQueueMode = QueueLIFO+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Create 2 connections, then release them in order+ pid1 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure rows+ pid2 <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT pg_backend_pid()"+ -- this should be the same as pid1 (LIFO: reuses most recent)+ pure rows+ -- In LIFO mode, we should get the same connection back+ pid1 `shouldBe` pid2+ closePool pool++ it "QueueFIFO rotates through connections" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolQueueMode = QueueFIFO+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Create 2 idle connections first+ barrier <- newEmptyMVar+ done1 <- newEmptyMVar+ done2 <- newEmptyMVar+ _ <- async $ withResource pool $ \_ -> do+ putMVar barrier ()+ takeMVar done1+ takeMVar barrier+ _ <- async $ withResource pool $ \_ -> do+ putMVar barrier ()+ takeMVar done2+ takeMVar barrier+ -- Now release both+ putMVar done1 ()+ putMVar done2 ()+ threadDelay 50000+ -- With 2 idle, FIFO should give oldest first+ stats <- poolStats pool+ psIdle stats `shouldSatisfy` (>= 2)+ -- Acquire twice and verify we get different PIDs+ pid1 <- withResource pool $ \conn -> do+ ([[Just p]], _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure p+ pid2 <- withResource pool $ \conn -> do+ ([[Just p]], _) <- simpleQuery conn "SELECT pg_backend_pid()"+ pure p+ -- FIFO should rotate, not reuse the same connection+ -- (Note: the first might equal the second if there's only one idle)+ -- At minimum, verify both acquisitions work+ pid1 `shouldNotBe` ""+ pid2 `shouldNotBe` ""+ closePool pool++ describe "pool timeout" $ do+ it "times out when pool is exhausted" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 1+ , poolAcquireTimeout = 0.5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ barrier <- newEmptyMVar+ done <- newEmptyMVar+ -- Hold the only connection+ _ <- async $ withResource pool $ \_ -> do+ putMVar barrier ()+ takeMVar done+ takeMVar barrier+ -- Try to acquire — should timeout+ result <- try @SomeException $ withResource pool $ \_ -> pure ()+ putMVar done ()+ closePool pool+ case result of+ Left e -> show e `shouldContain` "PoolTimeout"+ Right () -> expectationFailure "Expected PoolTimeout"++ it "timeout increments counter" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 1+ , poolAcquireTimeout = 0.2+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ barrier <- newEmptyMVar+ done <- newEmptyMVar+ _ <- async $ withResource pool $ \_ -> do+ putMVar barrier ()+ takeMVar done+ takeMVar barrier+ _ <- try @SomeException $ withResource pool $ \_ -> pure ()+ putMVar done ()+ threadDelay 50000+ stats <- poolStats pool+ psTotalTimeouts stats `shouldBe` 1+ closePool pool++ describe "no-thunks invariant under churn" $ do+ -- Regression guard. Catches future bugs of the form+ -- writeTVar (pTotalCreated pool) (oldVal + 1)+ -- where the stored value is a thunk rather than WHNF. Existing+ -- code uses modifyTVar' / strict-evaluated values everywhere, so+ -- this should remain green.+ it "pool bookkeeping TVars stay thunk-free after churn" $ do+ url <- requireDatabaseUrl+ let cfg = defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ , poolMaxLifeJitter = 0+ }+ pool <- newPool cfg+ -- Serial churn: forces counter increments and idle/active updates.+ replicateM_ 200 $ withResource pool $ \conn -> do+ _ <- simpleQuery conn "SELECT 1"+ pure ()+ -- Concurrent churn: stresses contention paths.+ _ <- mapConcurrently+ (\_ -> withResource pool $ \conn -> simpleQuery conn "SELECT 1")+ [1 :: Int .. 50]+ thunks <- concat <$> sequence+ [ checkTVar "pIdle" (pIdle pool)+ , checkTVar "pActive" (pActive pool)+ , checkTVar "pClosed" (pClosed pool)+ , checkTVar "pEffectiveSize" (pEffectiveSize pool)+ , checkTVar "pTotalCreated" (pTotalCreated pool)+ , checkTVar "pTotalDestroyed" (pTotalDestroyed pool)+ , checkTVar "pTotalTimeouts" (pTotalTimeouts pool)+ , checkTVar "pConnMeta" (pConnMeta pool)+ ]+ closePool pool+ case thunks of+ [] -> pure ()+ leaks -> expectationFailure $ "Found thunks in pool state: " <> show leaks++checkTVar :: NoThunks a => String -> TVar a -> IO [(String, ThunkInfo)]+checkTVar name tvar = do+ v <- readTVarIO tvar+ result <- noThunks [name] v+ pure $ maybe [] (\info -> [(name, info)]) result
+ integration/SoundnessSpec.hs view
@@ -0,0 +1,285 @@+-- | Integration tests for the soundness fixes from the correctness audit.+--+-- These tests exercise error paths, exception safety, and resource cleanup+-- that are difficult to verify from code review alone.+module SoundnessSpec (spec) where++import Control.Exception (SomeException, try)+import Data.ByteString.Char8 qualified as BS8+import Data.Int (Int32)+import Valiant+import PgWire.Connection.Config (parseConnString)+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ commitFailureSpec+ savepointCleanupSpec+ copyExceptionSpec+ cursorCleanupSpec+ advisoryLockSpec+ cacheEvictionSpec+ recycleCleanCacheSpec+ stmtCacheSizeSpec++------------------------------------------------------------------------+-- Transaction COMMIT failure (#4 from second-pass audit)+------------------------------------------------------------------------++commitFailureSpec :: Spec+commitFailureSpec = describe "COMMIT failure recovery" $ do+ it "rolls back on deferred constraint violation so connection is reusable" $ do+ withTestPool $ \pool -> do+ -- Set up a table with a deferred unique constraint+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS soundness_deferred CASCADE"+ _ <- simpleQuery conn+ "CREATE TABLE soundness_deferred (\+ \ id INTEGER NOT NULL,\+ \ CONSTRAINT deferred_unique UNIQUE (id) DEFERRABLE INITIALLY DEFERRED\+ \)"+ pure ()++ -- Insert a duplicate — the constraint is deferred, so the INSERT+ -- succeeds but COMMIT fails.+ result <- try @SomeException $ withTransaction pool $ \tx -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO soundness_deferred (id) VALUES (1)"+ _ <- simpleQuery (txConn tx) "INSERT INTO soundness_deferred (id) VALUES (1)"+ pure ()+ result `shouldSatisfy` isLeft++ -- The connection should be returned to the pool in a clean state.+ -- If COMMIT failure wasn't handled, this would fail with+ -- "current transaction is aborted".+ withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ rows `shouldBe` [[Just "1"]]++ -- Cleanup+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS soundness_deferred"+ pure ()++------------------------------------------------------------------------+-- Savepoint cleanup (#8 from first audit)+------------------------------------------------------------------------++savepointCleanupSpec :: Spec+savepointCleanupSpec = describe "Savepoint cleanup" $ do+ it "rolled-back savepoint doesn't corrupt subsequent queries" $ do+ withTestPool $ \pool -> do+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS soundness_sp CASCADE"+ _ <- simpleQuery conn "CREATE TABLE soundness_sp (id INTEGER NOT NULL)"+ pure ()++ -- Multiple savepoint failures in a row should all be cleanly handled+ withTransaction pool $ \tx -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO soundness_sp (id) VALUES (1)"++ -- First failed savepoint+ _ <- try @SomeException $ withSavepoint tx $ \_ ->+ error "sp failure 1"++ -- Second failed savepoint (tests that first cleanup was correct)+ _ <- try @SomeException $ withSavepoint tx $ \_ ->+ error "sp failure 2"++ -- Should still be able to query inside the transaction+ (rows, _) <- simpleQuery (txConn tx) "SELECT count(*) FROM soundness_sp"+ case rows of+ [[Just n]] -> n `shouldBe` "1"+ _ -> expectationFailure $ "Unexpected: " <> show rows++ -- Verify commit worked+ withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT count(*) FROM soundness_sp"+ case rows of+ [[Just n]] -> n `shouldBe` "1"+ _ -> expectationFailure $ "Unexpected: " <> show rows+ _ <- simpleQuery conn "DROP TABLE soundness_sp"+ pure ()++------------------------------------------------------------------------+-- COPY exception handling (#4 from first audit)+------------------------------------------------------------------------++copyExceptionSpec :: Spec+copyExceptionSpec = describe "COPY exception safety" $ do+ it "connection is usable after producer throws during copyIn" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ -- Producer throws mid-COPY+ result <- try @SomeException $+ copyIn conn "COPY users (name, email) FROM STDIN WITH (FORMAT csv)" $ \send -> do+ send "Alice,alice@example.com\n"+ error "producer failure"+ result `shouldSatisfy` isLeft++ -- Connection should still be usable (CopyFail was sent)+ (rows, _) <- simpleQuery conn "SELECT 1"+ rows `shouldBe` [[Just "1"]]++ it "preserves original exception message" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ result <- try @SomeException $+ copyIn conn "COPY users (name, email) FROM STDIN WITH (FORMAT csv)" $ \send -> do+ send "Alice,alice@example.com\n"+ error "specific error message"+ case result of+ Left e -> show e `shouldSatisfy` BS8.isInfixOf "specific error message" . BS8.pack+ Right _ -> expectationFailure "Expected exception"++------------------------------------------------------------------------+-- Cursor cleanup on exception (#3 from first audit)+------------------------------------------------------------------------++cursorCleanupSpec :: Spec+cursorCleanupSpec = describe "Fold cleanup on exception" $ do+ it "connection is usable after exception during fold" $ do+ withTestConnection $ \conn -> withSchema conn $ do+ insertTestUsers conn++ -- executeWithFold uses exclusive mode. The fold step throws on the+ -- first row, exercising the exception cleanup path (drain to+ -- ReadyForQuery before releasing exclusive mode).+ result <- try @SomeException $+ executeWithFold conn countQuery ()+ (RowFold (0 :: Int) (\_ _ -> error "mid-fold failure"))+ result `shouldSatisfy` isLeft++ -- Connection should still be usable after the fold exception+ (rows, _) <- simpleQuery conn "SELECT 1"+ rows `shouldBe` [[Just "1"]]++ where+ countQuery :: Statement () Int32+ countQuery = mkStatement+ "SELECT count(*)::int4 FROM users"+ []+ ["count"]+ "<test>"++------------------------------------------------------------------------+-- Advisory lock safety (#5 from first audit, #19 from MEDIUM fixes)+------------------------------------------------------------------------++advisoryLockSpec :: Spec+advisoryLockSpec = describe "Advisory locks" $ do+ it "session-scoped lock is released even when action throws" $ do+ withTestConnection $ \conn -> do+ -- Acquire lock, throw inside+ _ <- try @SomeException $ withAdvisoryLock conn 99999 $+ error "inside lock"++ -- Lock should be released — try to acquire it again immediately+ -- If the lock leaked, this would deadlock+ withAdvisoryLock conn 99999 $ pure ()++ it "try-lock releases on exception" $ do+ withTestConnection $ \conn -> do+ _ <- try @SomeException $ withAdvisoryLockTry conn 88888 $+ (error "inside try lock" :: IO ())++ -- Should be able to acquire again+ result <- withAdvisoryLockTry conn 88888 $ pure (42 :: Int)+ result `shouldBe` Just 42++ it "transaction-scoped lock requires active transaction" $ do+ withTestConnection $ \conn -> do+ result <- try @SomeException $ withAdvisoryLockTx conn 77777 $+ pure ()+ result `shouldSatisfy` isLeft++------------------------------------------------------------------------+-- Statement cache LRU eviction+------------------------------------------------------------------------++cacheEvictionSpec :: Spec+cacheEvictionSpec = describe "Statement cache eviction" $ do+ it "handles more unique queries than cache size without error" $ do+ withTestConnection $ \conn -> do+ -- Generate more unique queries than the default cache size (256).+ -- Each query has a different constant, forcing a new prepared statement.+ -- This exercises the LRU eviction path.+ let n = 300+ results <- mapM (\i -> do+ let q = "SELECT " <> BS8.pack (show i)+ (rows, _) <- simpleQuery conn q+ pure rows+ ) [1 :: Int .. n]+ length results `shouldBe` n++ it "frequently-used statements survive eviction" $ do+ withTestConnection $ \conn -> do+ -- Execute a "hot" query repeatedly, interleaved with many unique queries.+ -- The hot query should remain cached (LRU keeps it) and not cause errors.+ _ <- simpleQuery conn "CREATE TEMP TABLE cache_test (x INT)"+ let hotQuery = "SELECT count(*) FROM cache_test"+ -- Warm the hot query+ _ <- simpleQuery conn hotQuery++ -- Flood with unique queries to trigger eviction+ mapM_ (\i -> simpleQuery conn ("SELECT " <> BS8.pack (show i))) [1 :: Int .. 300]++ -- Hot query should still work (re-prepared if evicted, but no error)+ (rows, _) <- simpleQuery conn hotQuery+ rows `shouldBe` [[Just "0"]]++------------------------------------------------------------------------+-- RecycleClean clears statement cache (#15)+------------------------------------------------------------------------++recycleCleanCacheSpec :: Spec+recycleCleanCacheSpec = describe "RecycleClean cache invalidation" $ do+ it "DISCARD ALL clears statement cache so no stale references" $ do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 1+ , poolRecyclingMethod = RecycleClean+ , poolIdleTime = 600+ , poolHealthCheckAge = 0 -- always check+ }++ -- First use: prepare a statement+ withResource pool $ \conn -> do+ _ <- simpleQuery conn "SELECT 1"+ pure ()++ -- Return connection, re-acquire (triggers RecycleClean → DISCARD ALL).+ -- The cache should be cleared. Using the same query should re-prepare+ -- without "prepared statement does not exist".+ withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT 1"+ rows `shouldBe` [[Just "1"]]++ closePool pool++------------------------------------------------------------------------+-- Configurable statement cache size+------------------------------------------------------------------------++stmtCacheSizeSpec :: Spec+stmtCacheSizeSpec = describe "Statement cache size" $ do+ it "respects custom cache size from config" $ do+ url <- requireDatabaseUrl+ case parseConnString url of+ Left err -> expectationFailure err+ Right cfg -> do+ let cfg' = cfg { ccMaxPreparedStatements = 10 }+ conn <- connect cfg'+ -- Run more unique queries than the cache size+ mapM_ (\i -> simpleQuery conn ("SELECT " <> BS8.pack (show i))) [1 :: Int .. 20]+ -- Should still work (eviction happening correctly)+ (rows, _) <- simpleQuery conn "SELECT 42"+ rows `shouldBe` [[Just "42"]]+ close conn++------------------------------------------------------------------------+-- Helpers+------------------------------------------------------------------------++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ integration/StreamingSpec.hs view
@@ -0,0 +1,71 @@+module StreamingSpec (spec) where++import Data.Int (Int32)+import Data.Text (Text)+import Valiant+import TestSupport+import Test.Hspec++stmtListAll :: Statement () (Int32, Text)+stmtListAll = mkStatement+ "SELECT id, name FROM users ORDER BY id"+ [] ["id", "name"] "<test>"++stmtListByPrefix :: Statement Text (Int32, Text)+stmtListByPrefix = mkStatement+ "SELECT id, name FROM users WHERE name LIKE $1 ORDER BY id"+ [25] ["id", "name"] "<test>"++spec :: Spec+spec = do+ describe "withCursor / fetchBatch" $ do+ it "drains a small result set across multiple batches" $ do+ withTestPool $ \pool -> do+ withTransaction pool $ \tx -> withSchema (txConn tx) $ do+ insertTestUsers (txConn tx)+ rows <- withCursor (txConn tx) stmtListAll () 2 $ \cs ->+ let loop !acc = do+ batch <- fetchBatch cs 2+ if null batch+ then pure (reverse acc)+ else loop (batch : acc)+ in loop []+ length (concat rows) `shouldBe` 5++ it "returns empty after exhaustion without re-querying" $ do+ withTestPool $ \pool -> do+ withTransaction pool $ \tx -> withSchema (txConn tx) $ do+ insertTestUsers (txConn tx)+ (firstBatch, secondBatch, thirdBatch) <-+ withCursor (txConn tx) stmtListAll () 100 $ \cs -> do+ b1 <- fetchBatch cs 100+ b2 <- fetchBatch cs 100+ b3 <- fetchBatch cs 100+ pure (b1, b2, b3)+ length firstBatch `shouldBe` 5+ secondBatch `shouldBe` []+ thirdBatch `shouldBe` []++ it "respects parameter binding" $ do+ withTestPool $ \pool -> do+ withTransaction pool $ \tx -> withSchema (txConn tx) $ do+ insertTestUsers (txConn tx)+ rows <- withCursor (txConn tx) stmtListByPrefix "A%" 100 $ \cs ->+ fetchBatch cs 100+ length rows `shouldBe` 1++ describe "fetchAllCursor" $ do+ it "materialises every row in order" $ do+ withTestPool $ \pool -> do+ withTransaction pool $ \tx -> withSchema (txConn tx) $ do+ insertBulkUsers (txConn tx) 1000+ rows <- fetchAllCursor (txConn tx) stmtListAll () 100+ length rows `shouldBe` 1000+ map fst rows `shouldBe` [1 .. 1000]++ it "handles batch sizes larger than the result set" $ do+ withTestPool $ \pool -> do+ withTransaction pool $ \tx -> withSchema (txConn tx) $ do+ insertTestUsers (txConn tx)+ rows <- fetchAllCursor (txConn tx) stmtListAll () 10000+ length rows `shouldBe` 5
+ integration/TestSupport.hs view
@@ -0,0 +1,107 @@+module TestSupport+ ( withTestConnection+ , withTestPool+ , requireDatabaseUrl+ , createSchema+ , dropSchema+ , withSchema+ , insertTestUsers+ , insertBulkUsers+ ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Valiant+import System.Environment (lookupEnv)++-- Connection helpers --------------------------------------------------------++requireDatabaseUrl :: IO ByteString+requireDatabaseUrl = do+ mUrl <- lookupEnv "DATABASE_URL"+ case mUrl of+ Just url -> pure (BS8.pack url)+ Nothing -> error+ "DATABASE_URL is not set.\n\n\+ \ Run: eval $(scripts/pg-setup.sh)\n\+ \ Or: export DATABASE_URL=postgres://user:pass@localhost:5432/mydb\n"++withTestConnection :: (Connection -> IO a) -> IO a+withTestConnection action = do+ url <- requireDatabaseUrl+ conn <- connectString url+ result <- action conn+ close conn+ pure result++withTestPool :: (Pool -> IO a) -> IO a+withTestPool action = do+ url <- requireDatabaseUrl+ pool <- newPool defaultPoolConfig+ { poolConnString = url+ , poolSize = 4+ , poolAcquireTimeout = 5+ }+ result <- action pool+ closePool pool+ pure result++-- Schema --------------------------------------------------------------------++createSchema :: Connection -> IO ()+createSchema conn = do+ _ <- simpleQuery conn+ "CREATE TABLE IF NOT EXISTS users (\+ \ id SERIAL PRIMARY KEY,\+ \ name TEXT NOT NULL,\+ \ email TEXT,\+ \ is_active BOOLEAN NOT NULL DEFAULT true,\+ \ created_at TIMESTAMPTZ NOT NULL DEFAULT now()\+ \)"+ _ <- simpleQuery conn+ "CREATE TABLE IF NOT EXISTS posts (\+ \ id SERIAL PRIMARY KEY,\+ \ author_id INTEGER NOT NULL REFERENCES users(id),\+ \ title TEXT NOT NULL,\+ \ body TEXT,\+ \ published_at TIMESTAMPTZ,\+ \ created_at TIMESTAMPTZ NOT NULL DEFAULT now()\+ \)"+ pure ()++dropSchema :: Connection -> IO ()+dropSchema conn = do+ _ <- simpleQuery conn "DROP TABLE IF EXISTS posts CASCADE"+ _ <- simpleQuery conn "DROP TABLE IF EXISTS users CASCADE"+ pure ()++withSchema :: Connection -> IO a -> IO a+withSchema conn action = do+ dropSchema conn+ createSchema conn+ result <- action+ dropSchema conn+ pure result++-- Fixtures ------------------------------------------------------------------++insertTestUsers :: Connection -> IO ()+insertTestUsers conn = do+ _ <- simpleQuery conn+ "INSERT INTO users (name, email) VALUES \+ \('Alice', 'alice@example.com'),\+ \('Bob', 'bob@example.com'),\+ \('Carol', NULL),\+ \('Dave', 'dave@example.com'),\+ \('Eve', 'eve@example.com')"+ pure ()++insertBulkUsers :: Connection -> Int -> IO ()+insertBulkUsers conn n = do+ let values = BS8.intercalate ","+ [ "('user_" <> BS8.pack (show i) <> "', 'user" <> BS8.pack (show i) <> "@test.com')"+ | i <- [1 .. n]+ ]+ sql = "INSERT INTO users (name, email) VALUES " <> values+ _ <- simpleQuery conn sql+ pure ()
+ integration/TransactionSpec.hs view
@@ -0,0 +1,96 @@+module TransactionSpec (spec) where++import Control.Exception (SomeException, try)+import Valiant+import TestSupport+import Test.Hspec++spec :: Spec+spec = do+ describe "withTransaction" $ do+ it "commits on success" $ do+ withTestPool $ \pool -> do+ withResource pool $ \conn -> do+ dropSchema conn+ createSchema conn++ _ <- withTransaction pool $ \tx ->+ simpleQuery (txConn tx) "INSERT INTO users (name) VALUES ('Alice')"++ count <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT count(*) FROM users"+ pure rows+ case count of+ [[Just n]] -> n `shouldBe` "1"+ _ -> expectationFailure $ "Unexpected: " <> show count++ withResource pool $ \conn -> dropSchema conn++ it "rolls back on exception" $ do+ withTestPool $ \pool -> do+ withResource pool $ \conn -> do+ dropSchema conn+ createSchema conn++ _ <- try @SomeException $ withTransaction pool $ \tx -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO users (name) VALUES ('Bob')"+ error "intentional failure"++ count <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT count(*) FROM users"+ pure rows+ case count of+ [[Just n]] -> n `shouldBe` "0"+ _ -> expectationFailure $ "Unexpected: " <> show count++ withResource pool $ \conn -> dropSchema conn++ describe "withSavepoint" $ do+ it "preserves outer transaction on savepoint failure" $ do+ withTestPool $ \pool -> do+ withResource pool $ \conn -> do+ dropSchema conn+ createSchema conn++ withTransaction pool $ \tx -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO users (name) VALUES ('Alice')"++ -- This savepoint should fail but not kill the transaction+ _ <- try @SomeException $ withSavepoint tx $ \_ -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO users (name) VALUES ('Bob')"+ error "savepoint failure"++ -- Alice's insert should still be visible+ (rows, _) <- simpleQuery (txConn tx) "SELECT count(*) FROM users"+ case rows of+ [[Just n]] -> n `shouldBe` "1"+ _ -> expectationFailure $ "Expected 1, got: " <> show rows++ -- After commit, Alice should be in the table+ count <- withResource pool $ \conn -> do+ (rows, _) <- simpleQuery conn "SELECT count(*) FROM users"+ pure rows+ case count of+ [[Just n]] -> n `shouldBe` "1"+ _ -> expectationFailure $ "Expected 1, got: " <> show count++ withResource pool $ \conn -> dropSchema conn++ it "commits savepoint on success" $ do+ withTestPool $ \pool -> do+ withResource pool $ \conn -> do+ dropSchema conn+ createSchema conn++ withTransaction pool $ \tx -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO users (name) VALUES ('Alice')"+ withSavepoint tx $ \_ -> do+ _ <- simpleQuery (txConn tx) "INSERT INTO users (name) VALUES ('Bob')"+ pure ()++ (rows, _) <- simpleQuery (txConn tx) "SELECT count(*) FROM users"+ case rows of+ [[Just n]] -> n `shouldBe` "2"+ _ -> expectationFailure $ "Expected 2, got: " <> show rows++ withResource pool $ \conn -> dropSchema conn
+ src/Valiant.hs view
@@ -0,0 +1,404 @@+-- | Valiant — compile-time checked SQL for Haskell.+--+-- This is the main entry point for the runtime library. Import this module+-- to get access to all user-facing types and functions.+--+-- == Quick start+--+-- @+-- {\-\# OPTIONS_GHC -fplugin=Valiant.Plugin+-- -fplugin-opt=Valiant.Plugin:sql-dir=sql \#-\}+--+-- module MyApp.Queries.Users where+--+-- import Valiant+--+-- -- sql\/users\/find_by_id.sql:+-- -- SELECT id, name, email FROM users WHERE id = $1+-- findById :: Statement Int32 (Maybe (Int32, Text, Maybe Text))+-- findById = queryFile \"users\/find_by_id.sql\"+-- @+--+-- == Runtime usage+--+-- @+-- pool <- 'newPool' 'defaultPoolConfig' { poolConnString = \"postgres:\/\/...\" }+--+-- -- Fetch one row+-- mUser <- 'withResource' pool $ \\conn ->+-- 'fetchOne' conn findById 42+--+-- -- Batch insert (pipelined)+-- 'withResource' pool $ \\conn ->+-- 'executeBatch' conn insertStmt [(\"Alice\", email1), (\"Bob\", email2)]+--+-- -- Transaction+-- 'withTransaction' pool $ \\tx ->+-- 'execute' ('txConn' tx) insertStmt (\"Carol\", email3)+-- @+module Valiant+ ( -- * Statement+ -- | A 'Statement' represents a compile-time validated SQL query with+ -- typed parameters and results. Create them with 'queryFile' (validated+ -- by the GHC plugin) or 'mkStatement' (for manual construction).+ Statement (..)+ , query+ , queryFile+ , queryFileAs+ , mkStatement++ -- * Execution+ -- | Execute statements against a 'Connection'. All functions use the+ -- PostgreSQL extended query protocol with binary format encoding.+ , fetchOne+ , fetchAll+ , fetchAllVec+ , fetchAllUnboxed+ , fetchAllWith+ , fetchScalar+ , fetchOneOrThrow+ , fetchOneOr+ , fetchFirst+ , fetchExists+ , forEach+ -- * Fast decode variants+ -- | Use these on hot paths with wide rows. 'FromRowFast' throws+ -- exceptions instead of returning 'Either', eliminating per-column+ -- allocation overhead. Derive via @Generic@ like 'FromRow'.+ , FromRowFast (..)+ , DecodeColumnFast (..)+ , fetchAllFast+ , fetchOneFast+ , execute+ , executeReturning+ , executeReturningMany+ , executeMany+ , executeBatch+ , fetchBatchOne+ , fetchBatchAll++ -- * Raw (unchecked) queries+ -- | Escape hatch for dynamic SQL or queries that don't fit the+ -- 'Statement' model. Uses the extended query protocol with+ -- binary-encoded parameters but no compile-time type checking.+ -- Parameters are pre-encoded 'ByteString' values with explicit OIDs.+ , rawFetchAll+ , rawFetchOne+ , rawExecute++ -- * Named parameters+ -- | Optional record-based parameter passing for queries using+ -- @:name@ syntax in SQL files. Derive 'ToNamedParams' via 'Generic'+ -- to pass a record whose field names match the SQL parameter names.+ , ToNamedParams (..)+ , NamedStatement+ , mkStatementNamed++ -- * Dynamic SQL+ -- | Build SQL at runtime for queries with dynamic WHERE clauses,+ -- optional filters, or variable structure.+ --+ -- Import "Valiant.Dynamic" directly for the full API (@sql@, @param@,+ -- @runSnippet@, etc.). The names are intentionally not re-exported+ -- here to avoid shadowing common local variable names.++ -- * UNNEST batch fetch+ -- | Fetch multiple entities by ID in a single query using+ -- @WHERE id = ANY($1::type[])@. Faster than pipelining for+ -- the common \"load N entities by ID\" pattern.+ , fetchByIds++ -- * Connection+ -- | Manage connections to PostgreSQL. Use 'connectString' for simple+ -- usage or 'connect' with a 'ConnConfig' for full control. For+ -- production use, prefer 'Pool' over direct connections.+ , Connection+ , ConnConfig (..)+ , TlsMode (..)+ , defaultConnConfig+ , connect+ , connectString+ , close+ , withConnection+ , simpleQuery++ -- * Pool+ -- | Thread-safe connection pool with configurable size, idle reaping,+ -- max lifetime, and health checking. Use 'withResource' to acquire+ -- and automatically release connections.+ , Pool+ , PoolConfig (..)+ , RecyclingMethod (..)+ , QueueMode (..)+ , defaultPoolConfig+ , newPool+ , closePool+ , drainPool+ , withResource+ , withResourceTimeout+ , newPoolFromString+ , PoolStats (..)+ , poolStats+ , poolIsAlive+ , resize+ , retain+ , setPostCreateHook+ , setOnAcquireHook+ , setPreReleaseHook+ , PoolLogger+ , nullPoolLogger+ -- | For structured pool event observations, import+ -- "PgWire.Pool.Observation" directly. Events include+ -- 'ConnectionCreated', 'ConnectionDestroyed', 'AcquireTimeout', etc.++ -- * Transactions+ -- | Run actions inside a database transaction. If an exception is+ -- thrown, the transaction is automatically rolled back.+ , Transaction (..)+ , IsolationLevel (..)+ , TransactionMode (..)+ , defaultTransactionMode+ , withTransaction+ , withTransaction_+ , withTransactionLevel+ , withTransactionMode+ , withTransactionConn+ , withTransactionLevelConn+ , withTransactionModeConn+ , withReadOnlyTransaction+ , withDeferrableTransaction+ , withTransactionRetry+ , withTransactionRetryIf+ , withSavepoint++ -- * Valiant monad (optional convenience)+ -- | For the @ReaderT Pool IO@ monad with lifted @*M@ functions,+ -- install @valiant-mtl@ and import "Valiant.Monad".+ -- For generic MTL-style operations, import "Valiant.Mtl".++ -- * Large objects+ -- | Streaming interface for binary data too large for a single column.+ -- All operations must occur within a transaction.+ , LoFd (..)+ , LoMode (..)+ , loCreate+ , loUnlink+ , loOpen+ , loClose+ , withLargeObject+ , loRead+ , loWrite+ , loSeek+ , loTell+ , loTruncate+ , loImport+ , loExport++ -- * Advisory locks+ -- | PostgreSQL advisory locks for distributed coordination.+ -- Session-scoped locks require explicit release; transaction-scoped+ -- locks are released automatically on COMMIT\/ROLLBACK.+ , withAdvisoryLock+ , withAdvisoryLockTry+ , advisoryLock+ , advisoryUnlock+ , withAdvisoryLockTx+ , withAdvisoryLockTxTry+ , advisoryLockTx++ -- * Cancellation+ -- | Cancel in-flight queries. 'cancelQuery' opens a separate TCP+ -- connection and sends a CancelRequest to the server.+ -- 'withQueryTimeout' wraps any action with a deadline.+ , cancelQuery+ , withQueryTimeout++ -- * Streaming fold (constant memory)+ -- | Process large result sets without buffering, no cursor needed.+ , RowFold (..)+ , executeWithFold++ -- * Pipelined reads+ -- | Execute multiple independent queries in a single round-trip+ -- using the 'Applicative' interface. Eliminates N+1 query overhead.+ , Pipeline+ , pipeFetchOne+ , pipeFetchAll+ , pipeFetchScalar+ , pipeExecute+ , runPipeline++ -- * Streaming (cursors)+ -- | Stream large result sets using server-side cursors, fetching+ -- rows in batches without loading everything into memory.+ -- Must be used inside a transaction.+ , CursorState+ , withCursor+ , fetchBatch+ , fetchAllCursor++ -- * LISTEN\/NOTIFY+ -- | Subscribe to PostgreSQL asynchronous notification channels.+ , Notification (..)+ , listen+ , unlisten+ , waitForNotification+ , waitForNotificationTimeout++ -- * COPY+ -- | Bulk data import\/export using the PostgreSQL COPY protocol.+ , CopyResult (..)+ , copyIn+ , copyInBinary+ , copyOut++ -- * Logging+ -- | Hooks for instrumenting query timing and connection events.+ , LogEvent (..)+ , LogLevel (..)+ , Logger+ , nullLogger+ , stderrLogger+ , poolLoggerFromLogger++ -- * Row decoding+ -- | Decode result rows into Haskell types. Instances are provided+ -- for tuples up to 6 elements, 'Maybe' for nullable columns, and+ -- @()@ for commands that return no rows.+ , FromRow (..)+ , FromRowStrict (..)+ , DecodeColumn (..)++ -- * Parameter encoding+ -- | Encode query parameters. Instances are provided for single+ -- values, 'Maybe' for nullable parameters, @()@ for no parameters,+ -- and tuples up to 6 elements.+ , ToParams (..)+ , EncodeField (..)++ -- * Error inspection+ -- | Structured error inspection, SQLSTATE access, and constraint+ -- violation helpers. See "Valiant.Error" for the full API.+ , ConstraintViolation (..)+ , constraintViolation+ , catchConstraintViolation+ , sqlState+ , pgErrorOf+ , isUniqueViolation+ , isForeignKeyViolation+ , isSerializationError+ , isDeadlockError++ -- * Binary codec type classes+ -- | Low-level encode\/decode classes for individual PostgreSQL types.+ -- Derive these for newtypes using @GeneralizedNewtypeDeriving@ or+ -- @DerivingVia@:+ --+ -- @+ -- newtype UserId = UserId Int32+ -- deriving newtype (PgEncode, PgDecode)+ -- @+ , PgEncode (..)+ , PgDecode (..)++ -- ** Refinement helpers+ -- | Validate decoded values against extra invariants without writing a+ -- full 'PgDecode' instance. Useful for newtypes whose binary+ -- representation matches a primitive but whose values must satisfy a+ -- predicate (e.g. non-empty 'Text', positive 'Int32', enum strings).+ --+ -- @+ -- newtype Email = Email Text+ --+ -- instance PgDecode Email where+ -- pgDecode = refineWith \"email\" $ \\t ->+ -- if Text.elem \'@\' t+ -- then Right (Email t)+ -- else Left (\"missing @ in \" <> show t)+ -- @+ , refine+ , refineWith++ -- * Binary types+ -- | PostgreSQL types that don't have a standard Haskell equivalent.+ , PgInet (..)+ , ipv4+ , ipv4Host+ , ipv6+ , ipv6Host+ , inetToText+ , PgMacAddr (..)+ , macAddr+ , macAddrToText+ , PgHStore (..)+ , hstoreFromList+ , hstoreToList+ , Unbounded (..)+ , finite+ , fromFinite+ , PgPoint (..)+ , PgInterval (..)+ , PgRange (..)+ , RangeBound (..)+ , CompositeField (..)++ -- * Re-exports+ -- | 'Generic' is re-exported for convenience so that users can derive+ -- 'FromRow' and 'ToNamedParams' without an extra import.+ , Generic+ ) where++import GHC.Generics (Generic)+import Valiant.Advisory (advisoryLock, advisoryLockTx, advisoryUnlock, withAdvisoryLock, withAdvisoryLockTry, withAdvisoryLockTx, withAdvisoryLockTxTry)+import Valiant.Batch (fetchByIds)+import Valiant.Binary.Composite (CompositeField (..))+import Valiant.Binary.Decode (refine, refineWith)+import Valiant.Binary.HStore (PgHStore (..), hstoreFromList, hstoreToList)+import Valiant.Binary.Inet (PgInet (..), ipv4, ipv4Host, ipv6, ipv6Host, inetToText)+import Valiant.Binary.MacAddr (PgMacAddr (..), macAddr, macAddrToText)+import Valiant.Binary.Point (PgPoint (..))+import Valiant.Binary.Unbounded (Unbounded (..), finite, fromFinite)+import Valiant.Binary.Interval (PgInterval (..))+import Valiant.Binary.Range (PgRange (..), RangeBound (..))+import Valiant.Copy (CopyResult (..), copyIn, copyInBinary, copyOut)+-- Empty import: Valiant.Dynamic is intentionally not re-exported (see the+-- "Dynamic SQL" haddock note above). The bare import ensures haddock can+-- resolve the "Valiant.Dynamic" cross-reference link.+import Valiant.Dynamic ()+import Valiant.Error (ConstraintViolation (..), catchConstraintViolation, constraintViolation, isDeadlockError, isForeignKeyViolation, isSerializationError, isUniqueViolation, pgErrorOf, sqlState)+import Valiant.Execute (execute, executeBatch, executeMany, executeReturning, executeReturningMany, fetchAll, fetchAllFast, fetchAllUnboxed, fetchAllVec, fetchAllWith, fetchBatchAll, fetchBatchOne, fetchExists, fetchFirst, fetchOne, fetchOneFast, fetchOneOr, fetchOneOrThrow, fetchScalar, forEach, rawExecute, rawFetchAll, rawFetchOne)+import Valiant.FromRowFast (FromRowFast (..), DecodeColumnFast (..))+import Valiant.Fold (RowFold (..), executeWithFold)+import Valiant.LargeObject (LoFd (..), LoMode (..), loClose, loCreate, loExport, loImport, loOpen, loRead, loSeek, loTell, loTruncate, loUnlink, loWrite, withLargeObject)+import Valiant.FromRow (DecodeColumn (..), FromRow (..))+import Valiant.FromRowStrict (FromRowStrict (..))+import Valiant.Logging (LogEvent (..), LogLevel (..), Logger, nullLogger, poolLoggerFromLogger, stderrLogger)+import Valiant.NamedParams (NamedStatement, ToNamedParams (..), mkStatementNamed)+import Valiant.Notify (Notification (..), listen, unlisten, waitForNotification, waitForNotificationTimeout)+import Valiant.Pipeline (Pipeline, pipeFetchOne, pipeFetchAll, pipeFetchScalar, pipeExecute, runPipeline)+import Valiant.Statement (Statement (..), mkStatement, query, queryFile, queryFileAs)+import Valiant.Streaming (CursorState, fetchAllCursor, fetchBatch, withCursor)+import Valiant.ToParams (EncodeField (..), ToParams (..))+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Valiant.Transaction (IsolationLevel (..), Transaction (..), TransactionMode (..), defaultTransactionMode, withDeferrableTransaction, withReadOnlyTransaction, withSavepoint, withTransaction, withTransactionConn, withTransactionLevel, withTransactionLevelConn, withTransactionMode, withTransactionModeConn, withTransactionRetry, withTransactionRetryIf, withTransaction_)+import PgWire.Cancel (cancelQuery, withQueryTimeout)+import PgWire.Connection (Connection, close, connect, connectString, simpleQuery, withConnection)+import PgWire.Connection.Config (ConnConfig (..), TlsMode (..), defaultConnConfig)+import PgWire.Pool (Pool, PoolStats (..), closePool, drainPool, newPool, poolIsAlive, poolStats, resize, retain, setPostCreateHook, setOnAcquireHook, setPreReleaseHook, withResource, withResourceTimeout)+import PgWire.Pool.Config (PoolConfig (..), PoolLogger, QueueMode (..), RecyclingMethod (..), defaultPoolConfig, nullPoolLogger)+-- Empty import: PgWire.Pool.Observation is not re-exported (events are an+-- opt-in observability surface, see haddock note above). The bare import+-- ensures haddock can resolve the cross-reference link.+import PgWire.Pool.Observation ()++import Data.ByteString (ByteString)++-- | Create a connection pool from a connection string with default settings.+--+-- Equivalent to @'newPool' 'defaultPoolConfig' { poolConnString = connStr }@.+--+-- @+-- pool <- newPoolFromString \"postgres:\/\/user:pass\@localhost:5432\/mydb\"+-- @+newPoolFromString :: ByteString -> IO Pool+newPoolFromString connStr = newPool defaultPoolConfig { poolConnString = connStr }
+ src/Valiant/Advisory.hs view
@@ -0,0 +1,139 @@+-- | PostgreSQL advisory lock helpers.+--+-- Advisory locks are application-level locks managed by PostgreSQL.+-- They are useful for distributed coordination, ensuring only one+-- process runs a particular operation at a time.+--+-- Transaction-scoped locks ('withAdvisoryLockTx') are released+-- automatically when the transaction ends. Session-scoped locks+-- ('withAdvisoryLock') must be explicitly released.+--+-- @+-- -- Transaction-scoped: lock released on COMMIT/ROLLBACK+-- withTransaction pool $ \\tx ->+-- withAdvisoryLockTx (txConn tx) 42 $+-- execute (txConn tx) exclusiveStmt params+--+-- -- Session-scoped with bracket+-- withResource pool $ \\conn ->+-- withAdvisoryLock conn 42 $+-- doExclusiveWork conn+-- @+module Valiant.Advisory+ ( -- * Session-scoped locks+ withAdvisoryLock+ , withAdvisoryLockTry+ , advisoryLock+ , advisoryUnlock+ -- * Transaction-scoped locks+ , withAdvisoryLockTx+ , withAdvisoryLockTxTry+ , advisoryLockTx+ ) where++import Control.Exception (bracket_, finally)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.Int (Int64)+import PgWire.Connection (Connection, simpleQuery, transactionStatus)+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend (TxStatus (..))++-- | Acquire a session-scoped advisory lock, run an action, then release.+-- Blocks until the lock is available.+--+-- @+-- withAdvisoryLock conn 12345 $ do+-- -- exclusive access guaranteed+-- ...+-- @+withAdvisoryLock :: Connection -> Int64 -> IO a -> IO a+withAdvisoryLock conn key =+ bracket_ (advisoryLock conn key) (advisoryUnlock conn key)++-- | Try to acquire a session-scoped advisory lock without blocking.+-- Returns 'Nothing' if the lock is already held.+--+-- @+-- result <- withAdvisoryLockTry conn 12345 doWork+-- case result of+-- Just val -> putStrLn \"Got the lock\"+-- Nothing -> putStrLn \"Lock is held by another session\"+-- @+withAdvisoryLockTry :: Connection -> Int64 -> IO a -> IO (Maybe a)+withAdvisoryLockTry conn key action = do+ acquired <- advisoryTryLock conn key+ if acquired+ then do+ result <- action `finally` advisoryUnlock conn key+ pure (Just result)+ else pure Nothing++-- | Acquire a session-scoped advisory lock. Blocks until available.+advisoryLock :: Connection -> Int64 -> IO ()+advisoryLock conn key = do+ _ <- simpleQuery conn ("SELECT pg_advisory_lock(" <> BS8.pack (show key) <> ")")+ pure ()++-- | Release a session-scoped advisory lock.+advisoryUnlock :: Connection -> Int64 -> IO ()+advisoryUnlock conn key = do+ _ <- simpleQuery conn ("SELECT pg_advisory_unlock(" <> BS8.pack (show key) <> ")")+ pure ()++-- | Try to acquire a session-scoped advisory lock without blocking.+-- Returns 'True' if acquired.+advisoryTryLock :: Connection -> Int64 -> IO Bool+advisoryTryLock conn key = do+ (rows, _) <- simpleQuery conn ("SELECT pg_try_advisory_lock(" <> BS8.pack (show key) <> ")")+ pure $ case rows of+ [[Just "t"]] -> True+ _ -> False++-- | Acquire a transaction-scoped advisory lock, run an action.+-- The lock is automatically released when the transaction ends.+-- Must be called within a transaction.+--+-- @+-- withTransaction pool $ \\tx ->+-- withAdvisoryLockTx (txConn tx) 42 $+-- execute (txConn tx) stmt params+-- @+withAdvisoryLockTx :: Connection -> Int64 -> IO a -> IO a+withAdvisoryLockTx conn key action = do+ requireTransaction conn "withAdvisoryLockTx"+ advisoryLockTx conn key+ action++-- | Try to acquire a transaction-scoped advisory lock without blocking.+-- Returns 'Nothing' if the lock is already held.+withAdvisoryLockTxTry :: Connection -> Int64 -> IO a -> IO (Maybe a)+withAdvisoryLockTxTry conn key action = do+ requireTransaction conn "withAdvisoryLockTxTry"+ acquired <- advisoryTryLockTx conn key+ if acquired+ then Just <$> action+ else pure Nothing++-- | Acquire a transaction-scoped advisory lock. Blocks until available.+-- Released automatically on COMMIT/ROLLBACK.+advisoryLockTx :: Connection -> Int64 -> IO ()+advisoryLockTx conn key = do+ _ <- simpleQuery conn ("SELECT pg_advisory_xact_lock(" <> BS8.pack (show key) <> ")")+ pure ()++-- | Try to acquire a transaction-scoped advisory lock without blocking.+advisoryTryLockTx :: Connection -> Int64 -> IO Bool+advisoryTryLockTx conn key = do+ (rows, _) <- simpleQuery conn ("SELECT pg_try_advisory_xact_lock(" <> BS8.pack (show key) <> ")")+ pure $ case rows of+ [[Just "t"]] -> True+ _ -> False++-- | Check that the connection is in a transaction. Throws if not.+requireTransaction :: Connection -> ByteString -> IO ()+requireTransaction conn label = do+ status <- transactionStatus conn+ case status of+ TxInTransaction -> pure ()+ _ -> throwPgWire (ProtocolError (label <> ": must be called within a transaction"))
+ src/Valiant/Batch.hs view
@@ -0,0 +1,140 @@+-- | UNNEST/ANY-based batch fetching for loading multiple entities by ID+-- in a single query.+--+-- Instead of N queries or N pipelined queries, this rewrites a+-- single-parameter lookup into a set-returning query using+-- @WHERE id = ANY($1::type[])@. PostgreSQL can optimize this into+-- a single index scan.+--+-- @+-- -- Fetch users 1, 2, 3, 42, 99 in one query:+-- users <- 'fetchByIds' conn+-- \"SELECT id, name, email FROM users WHERE id = ANY($1::int4[])\"+-- [23] -- element OID (int4)+-- [1, 2, 3, 42, 99]+-- -- users :: [(Int32, Text, Maybe Text)]+-- @+--+-- For the common pattern of loading entities by a list of IDs, this is+-- faster than pipelining because PostgreSQL handles one query plan+-- instead of N.+module Valiant.Batch+ ( fetchByIds+ ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.HashPSQ qualified as PSQ+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Word (Word32, Word64)+import PgWire.Async (Request (..), Response (..), ResponseCollector (..), submitRequest)+import PgWire.Binary.Types (PgEncode (..))+import PgWire.Connection (Connection (..))+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Frontend+import PgWire.Protocol.Oid (Oid (..))+import Valiant.Binary.Array (pgEncodeArray)+import Valiant.FromRow (FromRow (..))++-- | Fetch rows matching any of the given IDs in a single query.+--+-- The SQL must use @$1@ as an array parameter with @= ANY($1::type[])@.+-- The element OID specifies the array element type (e.g., 23 for int4).+--+-- @+-- users <- fetchByIds conn+-- \"SELECT id, name, email FROM users WHERE id = ANY($1::int4[])\"+-- 23 -- int4 element OID+-- [1, 2, 3, 42, 99]+-- @+fetchByIds+ :: (PgEncode a, FromRow r)+ => Connection+ -> ByteString+ -- ^ SQL with @$1@ as the array parameter+ -> Oid+ -- ^ Element type OID (e.g., 'PgWire.Protocol.Oid.oidInt4')+ -> [a]+ -- ^ IDs to fetch+ -> IO [r]+fetchByIds conn sql elemOid ids = do+ let arrayBytes = pgEncodeArray elemOid (V.fromList ids)+ arrayOid = arrayOidFor elemOid++ -- Prepare+ stmtName <- ensurePreparedRaw conn sql (V.singleton (unOid arrayOid))++ -- Bind + Execute + Sync via async channel+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery+ [ Bind "" stmtName+ (V.singleton BinaryFormat)+ (V.singleton (Just arrayBytes))+ (V.singleton BinaryFormat)+ , Execute "" 0+ , Sync+ ]+ CollectRows++ case resp of+ RespRows rows -> mapM (\row -> case fromRow row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> pure val) rows+ _ -> throwPgWire (ProtocolError "fetchByIds: unexpected response type")++-- Map element OID to array OID+arrayOidFor :: Oid -> Oid+arrayOidFor (Oid 16) = Oid 1000 -- bool[]+arrayOidFor (Oid 17) = Oid 1001 -- bytea[]+arrayOidFor (Oid 20) = Oid 1016 -- int8[]+arrayOidFor (Oid 21) = Oid 1005 -- int2[]+arrayOidFor (Oid 23) = Oid 1007 -- int4[]+arrayOidFor (Oid 25) = Oid 1009 -- text[]+arrayOidFor (Oid 1043) = Oid 1015 -- varchar[]+arrayOidFor (Oid 700) = Oid 1021 -- float4[]+arrayOidFor (Oid 701) = Oid 1022 -- float8[]+arrayOidFor (Oid 1114) = Oid 1115 -- timestamp[]+arrayOidFor (Oid 1184) = Oid 1185 -- timestamptz[]+arrayOidFor (Oid 1082) = Oid 1182 -- date[]+arrayOidFor (Oid 2950) = Oid 2951 -- uuid[]+arrayOidFor oid = oid -- fallback: use as-is++-- Internal helpers --------------------------------------------------------++-- | Bump the monotonic tick counter and return the new value.+nextTick :: Connection -> IO Word64+nextTick conn = atomicModifyIORef' (connStmtTick conn) (\n -> (n + 1, n + 1))++ensurePreparedRaw :: Connection -> ByteString -> Vector Word32 -> IO ByteString+ensurePreparedRaw conn sql oids = do+ cache <- readIORef (connStmtCache conn)+ case PSQ.lookup sql cache of+ Just (_prio, name) -> do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)+ pure name+ Nothing -> do+ evictIfNeeded conn cache+ counter <- atomicModifyIORef' (connStmtCounter conn) (\n -> (n + 1, n))+ let name = "s" <> BS8.pack (show counter)+ resp <- submitRequest (connAsync conn) $ ReqPrepare (Parse name sql oids)+ case resp of+ RespParsed -> do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)+ pure name+ _ -> throwPgWire (ProtocolError "ensurePreparedRaw: unexpected response type")++-- | Evict the least recently used statement if the cache is full.+evictIfNeeded :: Connection -> PSQ.HashPSQ ByteString Word64 ByteString -> IO ()+evictIfNeeded conn cache+ | PSQ.size cache < connMaxPreparedStatements conn = pure ()+ | otherwise = case PSQ.findMin cache of+ Nothing -> pure ()+ Just (oldSql, _prio, oldName) -> do+ resp <- submitRequest (connAsync conn) $ ReqClose (Close DescribeStatement oldName)+ case resp of+ RespClosed -> pure ()+ _ -> pure () -- best effort+ modifyIORef' (connStmtCache conn) (PSQ.delete oldSql)
+ src/Valiant/Binary/Array.hs view
@@ -0,0 +1,108 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Valiant.Binary.Array+ ( pgEncodeArray+ , pgDecodeArray+ ) where++import Data.Bits (shiftL, (.|.))+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 (Int32)+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Word (Word32)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (Oid (..))++-- | Encode a Vector as a PostgreSQL one-dimensional array in binary format.+--+-- PG binary array format:+-- 4 bytes: number of dimensions (1 for 1-D)+-- 4 bytes: has-null flag (0 or 1)+-- 4 bytes: element OID+-- -- per dimension:+-- 4 bytes: dimension size (number of elements)+-- 4 bytes: lower bound (1 for standard arrays)+-- -- per element:+-- 4 bytes: element length (-1 for NULL)+-- N bytes: element data+pgEncodeArray :: (PgEncode a) => Oid -> Vector a -> ByteString+pgEncodeArray (Oid elemOid) elems =+ LBS.toStrict . B.toLazyByteString $+ B.int32BE 1 -- 1 dimension+ <> B.int32BE 0 -- no nulls+ <> B.word32BE elemOid -- element OID+ <> B.int32BE (fromIntegral (V.length elems)) -- dimension size+ <> B.int32BE 1 -- lower bound (1-based)+ <> V.foldl' (\acc a -> acc <> encodeElement a) mempty elems+ where+ encodeElement a =+ let bs = pgEncode a+ in B.int32BE (fromIntegral (BS.length bs)) <> B.byteString bs++-- | Decode a PostgreSQL one-dimensional binary array to a @Vector@.+--+-- Limitations:+--+-- * Only one-dimensional arrays are supported. Multi-dimensional arrays+-- produce an error.+-- * NULL elements are not supported — PostgreSQL arrays containing NULLs+-- will fail with @\"NULL elements not supported in Vector decode\"@.+-- Use @Vector (Maybe a)@ with a custom decoder if you need nullable+-- array elements.+pgDecodeArray :: (PgDecode a) => ByteString -> Either String (Vector a)+pgDecodeArray bs+ | BS.length bs < 12 = Left "array: header too short (need at least 12 bytes)"+ | otherwise = do+ ndim <- decodeInt32At bs 0+ _hasNull <- decodeInt32At bs 4+ _elemOid <- decodeWord32At bs 8+ case ndim of+ 0 -> Right V.empty+ 1 -> do+ when (BS.length bs < 20) $ Left "array: 1-D header too short (need at least 20 bytes)"+ nelems <- decodeInt32At bs 12+ _lbound <- decodeInt32At bs 16+ decodeElements bs 20 (fromIntegral nelems)+ _ -> Left $ "array: only 1-D arrays supported, got " <> show ndim <> " dimensions"++decodeElements :: (PgDecode a) => ByteString -> Int -> Int -> Either String (Vector a)+decodeElements bs offset count = go offset count []+ where+ go _off 0 acc = Right (V.fromList (reverse acc))+ go off n acc+ | off + 4 > BS.length bs = Left "array: unexpected end of data reading element length"+ | otherwise = do+ elemLen <- decodeInt32At bs off+ if elemLen == -1+ then Left "array: NULL elements not supported in Vector decode"+ else do+ let dataOff = off + 4+ dataEnd = dataOff + fromIntegral elemLen+ when (dataEnd > BS.length bs) $+ Left "array: unexpected end of data reading element"+ let elemBs = BS.take (fromIntegral elemLen) (BS.drop dataOff bs)+ val <- pgDecode elemBs+ go dataEnd (n - 1) (val : acc)++-- Helpers -------------------------------------------------------------------++decodeInt32At :: ByteString -> Int -> Either String Int32+decodeInt32At bs off+ | off + 4 > BS.length bs = Left "array: unexpected end of data"+ | otherwise =+ 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 Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)++decodeWord32At :: ByteString -> Int -> Either String Word32+decodeWord32At bs off = fromIntegral <$> decodeInt32At bs off++when :: Bool -> Either String () -> Either String ()+when True e = e+when False _ = Right ()
+ src/Valiant/Binary/Composite.hs view
@@ -0,0 +1,90 @@+-- | Binary encode/decode for PostgreSQL composite (row) types.+--+-- PG binary format for composite values:+-- 4 bytes: number of fields (Int32)+-- per field:+-- 4 bytes: field type OID (Word32)+-- 4 bytes: field data length (-1 for NULL)+-- N bytes: field data+module Valiant.Binary.Composite+ ( pgEncodeComposite+ , pgDecodeComposite+ , CompositeField (..)+ ) where++import Data.Bits (shiftL, (.|.))+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 (Int32)+import Data.Word (Word32)++-- | A single field in a composite value.+data CompositeField = CompositeField+ { cfOid :: Word32+ , cfValue :: Maybe ByteString+ -- ^ Nothing for NULL+ }+ deriving stock (Show, Eq)++-- | Encode a list of composite fields to PG binary format.+pgEncodeComposite :: [CompositeField] -> ByteString+pgEncodeComposite fields =+ LBS.toStrict . B.toLazyByteString $+ B.int32BE (fromIntegral (length fields))+ <> foldMap encodeField fields+ where+ encodeField (CompositeField oid Nothing) =+ B.word32BE oid <> B.int32BE (-1)+ encodeField (CompositeField oid (Just bs)) =+ B.word32BE oid+ <> B.int32BE (fromIntegral (BS.length bs))+ <> B.byteString bs++-- | Decode PG binary composite format to a list of fields.+pgDecodeComposite :: ByteString -> Either String [CompositeField]+pgDecodeComposite bs+ | BS.length bs < 4 = Left "composite: header too short"+ | otherwise = do+ nfields <- decodeInt32At bs 0+ decodeFields bs 4 (fromIntegral nfields)++decodeFields :: ByteString -> Int -> Int -> Either String [CompositeField]+decodeFields _ _ 0 = Right []+decodeFields bs off n+ | off + 8 > BS.length bs = Left "composite: unexpected end reading field header"+ | otherwise = do+ oid <- decodeWord32At bs off+ len <- decodeInt32At bs (off + 4)+ if len == -1+ then do+ rest <- decodeFields bs (off + 8) (n - 1)+ Right (CompositeField oid Nothing : rest)+ else do+ let dataOff = off + 8+ dataEnd = dataOff + fromIntegral len+ when (dataEnd > BS.length bs) $+ Left "composite: unexpected end reading field data"+ let val = BS.take (fromIntegral len) (BS.drop dataOff bs)+ rest <- decodeFields bs dataEnd (n - 1)+ Right (CompositeField oid (Just val) : rest)++-- Helpers -------------------------------------------------------------------++decodeInt32At :: ByteString -> Int -> Either String Int32+decodeInt32At bs off+ | off + 4 > BS.length bs = Left "composite: int32 out of bounds"+ | otherwise =+ 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 Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)++decodeWord32At :: ByteString -> Int -> Either String Word32+decodeWord32At bs off = fromIntegral <$> decodeInt32At bs off++when :: Bool -> Either String () -> Either String ()+when True e = e+when False _ = Right ()
+ src/Valiant/Binary/Decode.hs view
@@ -0,0 +1,203 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Valiant.Binary.Decode+ ( refine+ , refineWith+ ) 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.Text (Text)+import Data.Text.Encoding qualified as TE+import Data.Time+ ( Day+ , LocalTime (..)+ , TimeOfDay+ , TimeZone (..)+ , UTCTime (..)+ , ZonedTime (..)+ , fromGregorian+ , minutesToTimeZone+ , picosecondsToDiffTime+ , timeToTimeOfDay+ , utcToZonedTime+ )+import Data.Time.Calendar (addDays)+import Data.Word (Word32, Word64)+import GHC.Float (castWord32ToFloat, castWord64ToDouble)+import PgWire.Binary.Types (PgDecode (..))++-- PG epoch offset from Unix epoch in seconds+pgEpochOffsetSeconds :: Int64+pgEpochOffsetSeconds = 946684800+{-# INLINE pgEpochOffsetSeconds #-}++pgEpochDay :: Day+pgEpochDay = fromGregorian 2000 1 1+{-# INLINE pgEpochDay #-}++-- Refinement --------------------------------------------------------------++-- | Apply a refinement check after binary decode. Useful for types whose+-- binary representation is the same as a primitive but which carry+-- additional invariants (e.g. validated email, UUID-formatted Text,+-- positive Int32, enum strings).+refine :: PgDecode a => (a -> Either String b) -> ByteString -> Either String b+refine check bs = pgDecode bs >>= check+{-# INLINE refine #-}++-- | Like 'refine' but prepends a context prefix to error messages.+refineWith :: PgDecode a => String -> (a -> Either String b) -> ByteString -> Either String b+refineWith ctx check bs = case pgDecode bs of+ Left e -> Left (ctx <> ": " <> e)+ Right a -> case check a of+ Left e -> Left (ctx <> ": " <> e)+ Right b -> Right b+{-# INLINE refineWith #-}++-- Helpers -----------------------------------------------------------------++decodeInt16BE :: ByteString -> Either String Int16+decodeInt16BE bs+ | BS.length bs /= 2 = Left $ "int16: expected 2 bytes, got " <> show (BS.length bs)+ | otherwise =+ let b0 = fromIntegral (BU.unsafeIndex bs 0) :: Int16+ b1 = fromIntegral (BU.unsafeIndex bs 1) :: Int16+ in Right (b0 `shiftL` 8 .|. b1)+{-# INLINE decodeInt16BE #-}++decodeInt32BE :: ByteString -> Either String Int32+decodeInt32BE bs+ | BS.length bs /= 4 = Left $ "int32: expected 4 bytes, got " <> show (BS.length bs)+ | otherwise =+ let b0 = fromIntegral (BU.unsafeIndex bs 0) :: Int32+ b1 = fromIntegral (BU.unsafeIndex bs 1) :: Int32+ b2 = fromIntegral (BU.unsafeIndex bs 2) :: Int32+ b3 = fromIntegral (BU.unsafeIndex bs 3) :: Int32+ in Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)+{-# INLINE decodeInt32BE #-}++decodeInt64BE :: ByteString -> Either String Int64+decodeInt64BE bs+ | BS.length bs /= 8 = Left $ "int64: expected 8 bytes, got " <> show (BS.length bs)+ | otherwise =+ let !b0 = fromIntegral (BU.unsafeIndex bs 0) :: Int64+ !b1 = fromIntegral (BU.unsafeIndex bs 1) :: Int64+ !b2 = fromIntegral (BU.unsafeIndex bs 2) :: Int64+ !b3 = fromIntegral (BU.unsafeIndex bs 3) :: Int64+ !b4 = fromIntegral (BU.unsafeIndex bs 4) :: Int64+ !b5 = fromIntegral (BU.unsafeIndex bs 5) :: Int64+ !b6 = fromIntegral (BU.unsafeIndex bs 6) :: Int64+ !b7 = fromIntegral (BU.unsafeIndex bs 7) :: Int64+ in Right (b0 `shiftL` 56 .|. b1 `shiftL` 48 .|. b2 `shiftL` 40 .|. b3 `shiftL` 32+ .|. b4 `shiftL` 24 .|. b5 `shiftL` 16 .|. b6 `shiftL` 8 .|. b7)+{-# INLINE decodeInt64BE #-}++decodeWord32BE :: ByteString -> Either String Word32+decodeWord32BE bs = fromIntegral <$> decodeInt32BE bs+{-# INLINE decodeWord32BE #-}++decodeWord64BE :: ByteString -> Either String Word64+decodeWord64BE bs = fromIntegral <$> decodeInt64BE bs+{-# INLINE decodeWord64BE #-}++-- Instances ---------------------------------------------------------------++instance PgDecode Bool where+ pgDecode bs+ | BS.length bs /= 1 = Left "bool: expected 1 byte"+ | otherwise = Right (BU.unsafeIndex bs 0 /= 0)+ {-# INLINE pgDecode #-}++instance PgDecode Int16 where+ pgDecode = decodeInt16BE+ {-# INLINE pgDecode #-}++instance PgDecode Int32 where+ pgDecode = decodeInt32BE+ {-# INLINE pgDecode #-}++instance PgDecode Int64 where+ pgDecode = decodeInt64BE+ {-# INLINE pgDecode #-}++instance PgDecode Float where+ pgDecode bs = castWord32ToFloat <$> decodeWord32BE bs+ {-# INLINE pgDecode #-}++instance PgDecode Double where+ pgDecode bs = castWord64ToDouble <$> decodeWord64BE bs+ {-# INLINE pgDecode #-}++instance PgDecode Text where+ -- PostgreSQL guarantees valid UTF-8 for text/varchar columns, so we+ -- use decodeUtf8 (no validation) instead of decodeUtf8' (validates+ -- every byte). This eliminates ~20-30ns per text column on ASCII data.+ pgDecode bs = Right (TE.decodeUtf8 bs)+ {-# INLINE pgDecode #-}++instance PgDecode ByteString where+ pgDecode = Right+ {-# INLINE pgDecode #-}++instance PgDecode UTCTime where+ pgDecode bs = do+ pgMicros <- decodeInt64BE bs+ let totalMicros = pgMicros + pgEpochOffsetSeconds * 1000000+ (days, remainMicros) = totalMicros `divMod` 86400000000+ unixEpoch = fromGregorian 1970 1 1+ day = addDays (fromIntegral days) unixEpoch+ picos = fromIntegral remainMicros * 1000000+ Right (UTCTime day (picosecondsToDiffTime picos))+ {-# INLINE pgDecode #-}++instance PgDecode Day where+ pgDecode bs = do+ days <- decodeInt32BE bs+ Right (addDays (fromIntegral days) pgEpochDay)+ {-# INLINE pgDecode #-}++instance PgDecode TimeOfDay where+ pgDecode bs = do+ micros <- decodeInt64BE bs+ let picos = fromIntegral micros * 1000000+ Right (timeToTimeOfDay (picosecondsToDiffTime picos))+ {-# INLINE pgDecode #-}++instance PgDecode LocalTime where+ pgDecode bs = do+ totalMicros <- decodeInt64BE bs+ let (days, remainMicros) = totalMicros `divMod` 86400000000+ day = addDays (fromIntegral days) pgEpochDay+ picos = fromIntegral remainMicros * 1000000+ tod = timeToTimeOfDay (picosecondsToDiffTime picos)+ Right (LocalTime day tod)+ {-# INLINE pgDecode #-}++-- | Decode @timestamptz@ as 'ZonedTime'. PostgreSQL stores timestamptz+-- as UTC microseconds (same as UTCTime). We decode to UTC and tag with+-- the UTC timezone. To get the original timezone, use application-level+-- conversion.+instance PgDecode ZonedTime where+ pgDecode bs = do+ utc <- pgDecode bs+ Right (utcToZonedTime (minutesToTimeZone 0) utc)+ {-# INLINE pgDecode #-}++-- | Decode @timetz@ as @(TimeOfDay, TimeZone)@. Binary format:+-- 8 bytes (microseconds) + 4 bytes (UTC offset in seconds, negated).+instance PgDecode (TimeOfDay, TimeZone) where+ pgDecode bs+ | BS.length bs /= 12 = Left $ "timetz: expected 12 bytes, got " <> show (BS.length bs)+ | otherwise = do+ micros <- decodeInt64BE (BS.take 8 bs)+ offsetRaw <- decodeInt32BE (BS.drop 8 bs)+ let picos = fromIntegral micros * 1000000+ tod = timeToTimeOfDay (picosecondsToDiffTime picos)+ -- PG stores offset as seconds west of UTC (negated)+ tzMins = negate (fromIntegral offsetRaw) `div` 60 :: Int+ Right (tod, minutesToTimeZone tzMins)+ {-# INLINE pgDecode #-}
+ src/Valiant/Binary/Encode.hs view
@@ -0,0 +1,221 @@+{-# OPTIONS_GHC -Wno-orphans #-}++module Valiant.Binary.Encode+ () where++import Data.Bits (unsafeShiftR)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Internal (unsafeCreate)+import Data.Int (Int16, Int32, Int64)+import Data.Text (Text)+import Data.Text.Encoding qualified as TE+import Data.Time+ ( Day+ , LocalTime (..)+ , TimeOfDay (..)+ , TimeZone (..)+ , UTCTime (..)+ , ZonedTime (..)+ , diffTimeToPicoseconds+ , fromGregorian+ , timeOfDayToTime+ , toModifiedJulianDay+ , zonedTimeToUTC+ )+import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)+import Data.Word (Word8, Word32, Word64)+import Foreign.Ptr (Ptr)+import Foreign.Storable (pokeByteOff)+import GHC.Float (castFloatToWord32, castDoubleToWord64)+import PgWire.Binary.Types (PgEncode (..))+import PgWire.Protocol.Oid++-- PG epoch: 2000-01-01 00:00:00 UTC+pgEpochOffsetMicros :: Int64+pgEpochOffsetMicros = 946684800000000+{-# INLINE pgEpochOffsetMicros #-}++pgEpochDay :: Integer+pgEpochDay = toModifiedJulianDay (fromGregorian 2000 1 1)+{-# INLINE pgEpochDay #-}++-- Direct byte writes --------------------------------------------------------+-- These avoid Builder → LazyByteString → toStrict overhead for fixed-size+-- types. A single unsafeCreate + poke is ~3x faster than the Builder path+-- for 2-8 byte values.++int16BE :: Int16 -> ByteString+int16BE n = unsafeCreate 2 $ \p -> pokeInt16BE p 0 n+{-# INLINE int16BE #-}++int32BE :: Int32 -> ByteString+int32BE n = unsafeCreate 4 $ \p -> pokeInt32BE p 0 n+{-# INLINE int32BE #-}++int64BE :: Int64 -> ByteString+int64BE n = unsafeCreate 8 $ \p -> pokeInt64BE p 0 n+{-# INLINE int64BE #-}++floatBE :: Float -> ByteString+floatBE f = let !w = castFloatToWord32 f in unsafeCreate 4 $ \p -> pokeWord32BE p 0 w+{-# INLINE floatBE #-}++doubleBE :: Double -> ByteString+doubleBE d = let !w = castDoubleToWord64 d in unsafeCreate 8 $ \p -> pokeWord64BE p 0 w+{-# INLINE doubleBE #-}++-- Poke helpers — write big-endian at an offset into a Ptr+pokeInt16BE :: Ptr Word8 -> Int -> Int16 -> IO ()+pokeInt16BE p off n = do+ pokeByteOff p off (fromIntegral (n `unsafeShiftR` 8) :: Word8)+ pokeByteOff p (off + 1) (fromIntegral n :: Word8)+{-# INLINE pokeInt16BE #-}++pokeInt32BE :: Ptr Word8 -> Int -> Int32 -> IO ()+pokeInt32BE p off n = do+ pokeByteOff p off (fromIntegral (n `unsafeShiftR` 24) :: Word8)+ pokeByteOff p (off + 1) (fromIntegral (n `unsafeShiftR` 16) :: Word8)+ pokeByteOff p (off + 2) (fromIntegral (n `unsafeShiftR` 8) :: Word8)+ pokeByteOff p (off + 3) (fromIntegral n :: Word8)+{-# INLINE pokeInt32BE #-}++pokeInt64BE :: Ptr Word8 -> Int -> Int64 -> IO ()+pokeInt64BE p off n = do+ pokeByteOff p off (fromIntegral (n `unsafeShiftR` 56) :: Word8)+ pokeByteOff p (off + 1) (fromIntegral (n `unsafeShiftR` 48) :: Word8)+ pokeByteOff p (off + 2) (fromIntegral (n `unsafeShiftR` 40) :: Word8)+ pokeByteOff p (off + 3) (fromIntegral (n `unsafeShiftR` 32) :: Word8)+ pokeByteOff p (off + 4) (fromIntegral (n `unsafeShiftR` 24) :: Word8)+ pokeByteOff p (off + 5) (fromIntegral (n `unsafeShiftR` 16) :: Word8)+ pokeByteOff p (off + 6) (fromIntegral (n `unsafeShiftR` 8) :: Word8)+ pokeByteOff p (off + 7) (fromIntegral n :: Word8)+{-# INLINE pokeInt64BE #-}++pokeWord32BE :: Ptr Word8 -> Int -> Word32 -> IO ()+pokeWord32BE p off n = do+ pokeByteOff p off (fromIntegral (n `unsafeShiftR` 24) :: Word8)+ pokeByteOff p (off + 1) (fromIntegral (n `unsafeShiftR` 16) :: Word8)+ pokeByteOff p (off + 2) (fromIntegral (n `unsafeShiftR` 8) :: Word8)+ pokeByteOff p (off + 3) (fromIntegral n :: Word8)+{-# INLINE pokeWord32BE #-}++pokeWord64BE :: Ptr Word8 -> Int -> Word64 -> IO ()+pokeWord64BE p off n = do+ pokeByteOff p off (fromIntegral (n `unsafeShiftR` 56) :: Word8)+ pokeByteOff p (off + 1) (fromIntegral (n `unsafeShiftR` 48) :: Word8)+ pokeByteOff p (off + 2) (fromIntegral (n `unsafeShiftR` 40) :: Word8)+ pokeByteOff p (off + 3) (fromIntegral (n `unsafeShiftR` 32) :: Word8)+ pokeByteOff p (off + 4) (fromIntegral (n `unsafeShiftR` 24) :: Word8)+ pokeByteOff p (off + 5) (fromIntegral (n `unsafeShiftR` 16) :: Word8)+ pokeByteOff p (off + 6) (fromIntegral (n `unsafeShiftR` 8) :: Word8)+ pokeByteOff p (off + 7) (fromIntegral n :: Word8)+{-# INLINE pokeWord64BE #-}++-- Instances ---------------------------------------------------------------++instance PgEncode Bool where+ pgEncode True = BS.singleton 1+ pgEncode False = BS.singleton 0+ {-# INLINE pgEncode #-}+ pgOid _ = oidBool+ {-# INLINE pgOid #-}++instance PgEncode Int16 where+ pgEncode = int16BE+ {-# INLINE pgEncode #-}+ pgOid _ = oidInt2+ {-# INLINE pgOid #-}++instance PgEncode Int32 where+ pgEncode = int32BE+ {-# INLINE pgEncode #-}+ pgOid _ = oidInt4+ {-# INLINE pgOid #-}++instance PgEncode Int64 where+ pgEncode = int64BE+ {-# INLINE pgEncode #-}+ pgOid _ = oidInt8+ {-# INLINE pgOid #-}++instance PgEncode Float where+ pgEncode = floatBE+ {-# INLINE pgEncode #-}+ pgOid _ = oidFloat4+ {-# INLINE pgOid #-}++instance PgEncode Double where+ pgEncode = doubleBE+ {-# INLINE pgEncode #-}+ pgOid _ = oidFloat8+ {-# INLINE pgOid #-}++instance PgEncode Text where+ pgEncode = TE.encodeUtf8+ {-# INLINE pgEncode #-}+ pgOid _ = oidText+ {-# INLINE pgOid #-}++instance PgEncode ByteString where+ pgEncode = id+ {-# INLINE pgEncode #-}+ pgOid _ = oidBytea+ {-# INLINE pgOid #-}++instance PgEncode UTCTime where+ pgEncode t =+ let !unixMicros = round (utcTimeToPOSIXSeconds t * 1000000) :: Int64+ !pgMicros = unixMicros - pgEpochOffsetMicros+ in int64BE pgMicros+ {-# INLINE pgEncode #-}+ pgOid _ = oidTimestamptz+ {-# INLINE pgOid #-}++instance PgEncode Day where+ pgEncode d =+ let !daysSincePgEpoch = toModifiedJulianDay d - pgEpochDay+ in int32BE (fromIntegral daysSincePgEpoch)+ {-# INLINE pgEncode #-}+ pgOid _ = oidDate+ {-# INLINE pgOid #-}++instance PgEncode TimeOfDay where+ pgEncode tod =+ let !picos = diffTimeToPicoseconds (timeOfDayToTime tod)+ !micros = picos `div` 1000000+ in int64BE (fromIntegral micros)+ {-# INLINE pgEncode #-}+ pgOid _ = oidTime+ {-# INLINE pgOid #-}++instance PgEncode LocalTime where+ pgEncode (LocalTime d tod) =+ let !daysSincePgEpoch = toModifiedJulianDay d - pgEpochDay+ !picos = diffTimeToPicoseconds (timeOfDayToTime tod)+ !microsSinceMidnight = picos `div` 1000000+ !totalMicros = fromIntegral daysSincePgEpoch * 86400000000 + fromIntegral microsSinceMidnight+ in int64BE totalMicros+ {-# INLINE pgEncode #-}+ pgOid _ = oidTimestamp+ {-# INLINE pgOid #-}++instance PgEncode ZonedTime where+ pgEncode = pgEncode . zonedTimeToUTC+ {-# INLINE pgEncode #-}+ pgOid _ = oidTimestamptz+ {-# INLINE pgOid #-}++-- | @timetz@ binary format: 8 bytes (microseconds) + 4 bytes (UTC offset in seconds, negated).+-- PostgreSQL stores the offset as seconds *west* of UTC (negated from the usual convention).+instance PgEncode (TimeOfDay, TimeZone) where+ pgEncode (tod, tz) =+ let !picos = diffTimeToPicoseconds (timeOfDayToTime tod)+ !micros = picos `div` 1000000+ !offsetSecs = fromIntegral (negate (timeZoneMinutes tz * 60)) :: Int32+ in unsafeCreate 12 $ \p -> do+ pokeInt64BE p 0 (fromIntegral micros)+ pokeInt32BE p 8 offsetSecs+ {-# INLINE pgEncode #-}+ pgOid _ = oidTimetz+ {-# INLINE pgOid #-}
+ src/Valiant/Binary/Enum.hs view
@@ -0,0 +1,44 @@+-- | Binary encode/decode for PostgreSQL enum types.+--+-- PG sends enum values as text (UTF-8 label) in binary format,+-- identical to the text encoding. This module provides the 'PgEnum'+-- type class for Haskell sum types that map to PG enums.+--+-- @+-- data Mood = Happy | Sad | Angry+-- deriving stock (Show, Eq, Bounded, Enum)+--+-- instance PgEnum Mood where+-- pgEnumToText Happy = "happy"+-- pgEnumToText Sad = "sad"+-- pgEnumToText Angry = "angry"+-- pgEnumFromText "happy" = Right Happy+-- pgEnumFromText "sad" = Right Sad+-- pgEnumFromText "angry" = Right Angry+-- pgEnumFromText t = Left ("unknown mood: " <> T.unpack t)+-- @+module Valiant.Binary.Enum+ ( PgEnum (..)+ , pgEncodeEnum+ , pgDecodeEnum+ ) where++import Data.ByteString (ByteString)+import Data.Text (Text)+import Data.Text.Encoding qualified as TE++-- | Type class for Haskell types that correspond to PostgreSQL enums.+-- PG encodes enums as their text label in both text and binary format.+class PgEnum a where+ pgEnumToText :: a -> Text+ pgEnumFromText :: Text -> Either String a++-- | Encode an enum value to its PG binary representation (UTF-8 label).+pgEncodeEnum :: (PgEnum a) => a -> ByteString+pgEncodeEnum = TE.encodeUtf8 . pgEnumToText+{-# INLINE pgEncodeEnum #-}++-- | Decode a PG binary enum value (UTF-8 label) to a Haskell value.+pgDecodeEnum :: (PgEnum a) => ByteString -> Either String a+pgDecodeEnum = pgEnumFromText . TE.decodeUtf8+{-# INLINE pgDecodeEnum #-}
+ src/Valiant/Binary/HStore.hs view
@@ -0,0 +1,130 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary codec for the PostgreSQL @hstore@ extension type.+--+-- PG binary format for hstore:+-- 4 bytes: number of pairs (Int32)+-- per pair:+-- 4 bytes: key length (Int32)+-- N bytes: key data (UTF-8)+-- 4 bytes: value length (Int32, -1 for NULL)+-- N bytes: value data (UTF-8)+--+-- The extension must be enabled: @CREATE EXTENSION IF NOT EXISTS hstore@.+-- hstore OID is not fixed — it depends on installation. The OID is+-- typically resolved dynamically via @pg_type@ during @valiant prepare@.+module Valiant.Binary.HStore+ ( PgHStore (..)+ , hstoreFromList+ , hstoreToList+ ) where++import Control.DeepSeq (NFData)+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.Bifunctor (second)+import Data.Bits (shiftL, (.|.))+import Data.Int (Int32)+import Data.Map.Strict (Map)+import Data.Map.Strict qualified as Map+import Data.Text (Text)+import Data.Text.Encoding qualified as T+import GHC.Generics (Generic)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (Oid (..))++-- | A PostgreSQL @hstore@ value — a set of key/value string pairs.+-- Values are nullable (a key can map to SQL NULL).+newtype PgHStore = PgHStore+ { unHStore :: Map Text (Maybe Text)+ -- ^ The underlying map. 'Nothing' values represent SQL NULL.+ }+ deriving stock (Show, Eq, Generic)++instance NFData PgHStore++-- | Construct a 'PgHStore' from a list of key/value pairs.+-- All values are non-NULL.+--+-- @+-- hstoreFromList [(\"color\", \"red\"), (\"size\", \"large\")]+-- @+hstoreFromList :: [(Text, Text)] -> PgHStore+hstoreFromList = PgHStore . Map.fromList . map (second Just)++-- | Extract non-NULL key/value pairs as a list.+hstoreToList :: PgHStore -> [(Text, Text)]+hstoreToList (PgHStore m) =+ [(k, v) | (k, Just v) <- Map.toList m]++-- | hstore extension OID is installation-dependent. We use 0 here as a+-- placeholder; the actual OID is resolved at prepare time.+oidHStore :: Oid+oidHStore = Oid 0++instance PgEncode PgHStore where+ pgEncode (PgHStore m) =+ let pairs = Map.toList m+ nPairs = fromIntegral (length pairs) :: Int32+ in LBS.toStrict . B.toLazyByteString $+ B.int32BE nPairs <> foldMap encodePair pairs+ where+ encodePair (k, mv) =+ let kbs = T.encodeUtf8 k+ in B.int32BE (fromIntegral (BS.length kbs))+ <> B.byteString kbs+ <> case mv of+ Nothing -> B.int32BE (-1)+ Just v ->+ let vbs = T.encodeUtf8 v+ in B.int32BE (fromIntegral (BS.length vbs))+ <> B.byteString vbs+ {-# INLINE pgEncode #-}+ pgOid _ = oidHStore+ {-# INLINE pgOid #-}++instance PgDecode PgHStore where+ pgDecode bs+ | BS.length bs < 4 = Left "hstore: too short for pair count"+ | otherwise = do+ nPairs <- decodeInt32At bs 0+ decodePairs bs 4 (fromIntegral nPairs :: Int) []+ where+ decodePairs _ _ 0 !acc = Right (PgHStore (Map.fromList (reverse acc)))+ decodePairs b off n !acc+ | off + 4 > BS.length b = Left "hstore: unexpected end reading key length"+ | otherwise = do+ kLen <- decodeInt32At b off+ let kOff = off + 4+ kEnd = kOff + fromIntegral kLen+ if kEnd > BS.length b+ then Left "hstore: unexpected end reading key data"+ else do+ let key = T.decodeUtf8 (BS.take (fromIntegral kLen) (BS.drop kOff b))+ if kEnd + 4 > BS.length b+ then Left "hstore: unexpected end reading value length"+ else do+ vLen <- decodeInt32At b kEnd+ if vLen == -1+ then decodePairs b (kEnd + 4) (n - 1) ((key, Nothing) : acc)+ else do+ let vOff = kEnd + 4+ vEnd = vOff + fromIntegral vLen+ if vEnd > BS.length b+ then Left "hstore: unexpected end reading value data"+ else do+ let val = T.decodeUtf8 (BS.take (fromIntegral vLen) (BS.drop vOff b))+ decodePairs b vEnd (n - 1) ((key, Just val) : acc)+ {-# INLINE pgDecode #-}++decodeInt32At :: ByteString -> Int -> Either String Int32+decodeInt32At bs off+ | off + 4 > BS.length bs = Left "hstore: int32 out of bounds"+ | otherwise =+ 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 Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)
+ src/Valiant/Binary/Inet.hs view
@@ -0,0 +1,198 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary codec for PostgreSQL @inet@ and @cidr@ types.+--+-- PG binary format (both inet and cidr):+--+-- 1 byte: address family (2 = IPv4, 3 = IPv6)+-- 1 byte: prefix length (netmask bits: 0–32 for v4, 0–128 for v6)+-- 1 byte: is_cidr flag (0 = inet, 1 = cidr)+-- 1 byte: address length (4 for IPv4, 16 for IPv6)+-- N bytes: address bytes (network byte order)+--+-- Total: 8 bytes for IPv4, 20 bytes for IPv6.+module Valiant.Binary.Inet+ ( PgInet (..)+ , ipv4+ , ipv4Host+ , ipv6+ , ipv6Host+ , inetToText+ ) where++import Control.DeepSeq (NFData)+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.List (intercalate)+import Data.Text (Text)+import Data.Text qualified as T+import Data.Word (Word8, Word16)+import GHC.Generics (Generic)+import Numeric (showHex)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidInet)++-- | A PostgreSQL @inet@ or @cidr@ value.+--+-- Represents an IPv4 or IPv6 address with a prefix length (CIDR mask).+-- The 'inetIsCidr' flag distinguishes between @inet@ (host + optional mask)+-- and @cidr@ (network address, host bits must be zero).+data PgInet = PgInet+ { inetAddress :: !ByteString+ -- ^ Raw address bytes: 4 bytes for IPv4, 16 bytes for IPv6.+ , inetPrefix :: {-# UNPACK #-} !Word8+ -- ^ Prefix length (0–32 for IPv4, 0–128 for IPv6).+ , inetIsCidr :: !Bool+ -- ^ 'True' for @cidr@, 'False' for @inet@.+ }+ deriving stock (Eq, Ord, Generic)++instance NFData PgInet++-- | Renders in CIDR notation: @\"192.168.1.0/24\"@ or @\"::1/128\"@.+instance Show PgInet where+ show = T.unpack . inetToText++-- | Construct an IPv4 @inet@ value with a prefix length.+--+-- @ipv4 192 168 1 0 24@ represents @192.168.1.0\/24@.+ipv4 :: Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> PgInet+ipv4 a b c d prefix =+ PgInet+ { inetAddress = BS.pack [a, b, c, d]+ , inetPrefix = prefix+ , inetIsCidr = False+ }++-- | Construct an IPv4 @inet@ host (prefix = 32).+ipv4Host :: Word8 -> Word8 -> Word8 -> Word8 -> PgInet+ipv4Host a b c d = ipv4 a b c d 32++-- | Construct an IPv6 @inet@ value from raw address bytes with a prefix+-- length. The 'ByteString' must be exactly 16 bytes.+ipv6 :: ByteString -> Word8 -> PgInet+ipv6 addr prefix =+ PgInet+ { inetAddress = addr+ , inetPrefix = prefix+ , inetIsCidr = False+ }++-- | Construct an IPv6 @inet@ host (prefix = 128) from raw address bytes.+-- The 'ByteString' must be exactly 16 bytes.+ipv6Host :: ByteString -> PgInet+ipv6Host addr = ipv6 addr 128++-- | Render a 'PgInet' as CIDR notation text.+--+-- >>> inetToText (ipv4 192 168 1 0 24)+-- "192.168.1.0/24"+-- >>> inetToText (ipv6Host (BS.pack [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1]))+-- "::1/128"+inetToText :: PgInet -> Text+inetToText (PgInet addr prefix _isCidr)+ | BS.length addr == 4 = renderIPv4 addr <> "/" <> T.pack (show prefix)+ | BS.length addr == 16 = renderIPv6 addr <> "/" <> T.pack (show prefix)+ | otherwise = "<invalid inet>"++renderIPv4 :: ByteString -> Text+renderIPv4 bs =+ T.pack $+ intercalate "." $+ map (show . BS.index bs) [0 .. 3]++renderIPv6 :: ByteString -> Text+renderIPv6 bs =+ let groups :: [Word16]+ groups =+ [ fromIntegral (BS.index bs i) * 256 + fromIntegral (BS.index bs (i + 1))+ | i <- [0, 2 .. 14]+ ]+ hexGroups = map (\g -> T.pack (showHex g "")) groups+ in compressIPv6 hexGroups++-- | Compress an IPv6 address by replacing the longest run of zero groups+-- with "::".+compressIPv6 :: [Text] -> Text+compressIPv6 groups =+ let isZero g = g == "0"+ -- Find runs of zeros+ indexed = zip [0 :: Int ..] (map isZero groups)+ runs = findZeroRuns indexed+ in case runs of+ [] -> T.intercalate ":" groups+ _ ->+ let (start, len) = longestRun runs+ before = take start groups+ after = drop (start + len) groups+ in case (null before, null after) of+ (True, True) -> "::"+ (True, False) -> "::" <> T.intercalate ":" after+ (False, True) -> T.intercalate ":" before <> "::"+ (False, False) -> T.intercalate ":" before <> "::" <> T.intercalate ":" after++findZeroRuns :: [(Int, Bool)] -> [(Int, Int)]+findZeroRuns [] = []+findZeroRuns ((i, True) : rest) =+ let runLen = 1 + length (takeWhile snd rest)+ remaining = drop (runLen - 1) rest+ in (i, runLen) : findZeroRuns remaining+findZeroRuns (_ : rest) = findZeroRuns rest++longestRun :: [(Int, Int)] -> (Int, Int)+longestRun = foldl1 (\a@(_, la) b@(_, lb) -> if lb > la then b else a)++-- Codec -------------------------------------------------------------------++-- PG address family codes+pgAfInet :: Word8+pgAfInet = 2++pgAfInet6 :: Word8+pgAfInet6 = 3++instance PgEncode PgInet where+ pgEncode (PgInet addr prefix isCidr) =+ let addrLen = BS.length addr+ family+ | addrLen == 4 = pgAfInet+ | otherwise = pgAfInet6+ cidrFlag :: Word8+ cidrFlag = if isCidr then 1 else 0+ in LBS.toStrict . B.toLazyByteString $+ B.word8 family+ <> B.word8 prefix+ <> B.word8 cidrFlag+ <> B.word8 (fromIntegral addrLen)+ <> B.byteString addr+ {-# INLINE pgEncode #-}+ pgOid _ = oidInet+ {-# INLINE pgOid #-}++instance PgDecode PgInet where+ pgDecode bs+ | BS.length bs < 4 = Left $ "inet: expected at least 4 header bytes, got " <> show (BS.length bs)+ | otherwise =+ let family = BS.index bs 0+ prefix = BS.index bs 1+ cidrFlag = BS.index bs 2+ addrLen = fromIntegral (BS.index bs 3) :: Int+ expectedFamily+ | addrLen == 4 = pgAfInet+ | addrLen == 16 = pgAfInet6+ | otherwise = 0+ in if BS.length bs /= 4 + addrLen+ then Left $ "inet: expected " <> show (4 + addrLen) <> " bytes, got " <> show (BS.length bs)+ else+ if family /= expectedFamily+ then Left $ "inet: address family mismatch: " <> show family <> " for " <> show addrLen <> " byte address"+ else+ Right+ PgInet+ { inetAddress = BS.drop 4 bs+ , inetPrefix = prefix+ , inetIsCidr = cidrFlag /= 0+ }+ {-# INLINE pgDecode #-}
+ src/Valiant/Binary/Interval.hs view
@@ -0,0 +1,71 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary encode/decode for PostgreSQL @interval@ type.+--+-- PG stores interval as:+-- 8 bytes: microseconds (Int64)+-- 4 bytes: days (Int32)+-- 4 bytes: months (Int32)+-- Total: 16 bytes, big-endian.+module Valiant.Binary.Interval+ ( PgInterval (..)+ ) where++import Control.DeepSeq (NFData)+import Data.Bits (shiftL, (.|.))+import Data.ByteString (ByteString)+import GHC.Generics (Generic)+import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as B+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int32, Int64)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidInterval)++-- | A PostgreSQL interval value.+data PgInterval = PgInterval+ { intervalMicroseconds :: {-# UNPACK #-} !Int64+ , intervalDays :: {-# UNPACK #-} !Int32+ , intervalMonths :: {-# UNPACK #-} !Int32+ }+ deriving stock (Show, Eq, Ord, Generic)++instance NFData PgInterval++instance PgEncode PgInterval where+ pgEncode (PgInterval micros days months) =+ LBS.toStrict . B.toLazyByteString $+ B.int64BE micros+ <> B.int32BE days+ <> B.int32BE months+ pgOid _ = oidInterval++instance PgDecode PgInterval where+ pgDecode bs+ | BS.length bs /= 16 = Left $ "interval: expected 16 bytes, got " <> show (BS.length bs)+ | otherwise = do+ micros <- decodeInt64BE bs 0+ days <- decodeInt32BE bs 8+ months <- decodeInt32BE bs 12+ Right (PgInterval micros days months)++-- Helpers -------------------------------------------------------------------++decodeInt32BE :: ByteString -> Int -> Either String Int32+decodeInt32BE bs off+ | off + 4 > BS.length bs = Left "interval: int32 out of bounds"+ | otherwise =+ 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 Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)++decodeInt64BE :: ByteString -> Int -> Either String Int64+decodeInt64BE bs off+ | off + 8 > BS.length bs = Left "interval: int64 out of bounds"+ | otherwise =+ let go !acc !i+ | i >= 8 = acc+ | otherwise = go (acc `shiftL` 8 .|. fromIntegral (BS.index bs (off + i))) (i + 1)+ in Right (go 0 0)
+ src/Valiant/Binary/JSON.hs view
@@ -0,0 +1,44 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary encode/decode for PostgreSQL @json@ and @jsonb@ types.+--+-- JSON binary format: raw UTF-8 JSON text (same as text format).+-- JSONB binary format: version byte (0x01) followed by UTF-8 JSON text.+--+-- We provide instances for @aeson@'s 'Value' type. For json columns,+-- encoding is just the JSON bytes. For jsonb, PG prepends a version byte+-- on the wire, but the extended query protocol handles this transparently+-- when using binary format — the version byte is part of the binary+-- representation that PG sends/expects.+module Valiant.Binary.JSON+ () where++import Data.Aeson (Value, eitherDecodeStrict', encode)+import Data.ByteString qualified as BS+import Data.ByteString.Lazy qualified as LBS+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidJson)++-- | Encode a JSON Value. Works for both json and jsonb columns.+-- For jsonb, PG expects a version byte (0x01) prefix in binary format.+instance PgEncode Value where+ pgEncode = LBS.toStrict . encode+ {-# INLINE pgEncode #-}+ pgOid _ = oidJson+ {-# INLINE pgOid #-}++-- | Decode a JSON Value. Handles both json (raw JSON) and jsonb+-- (version byte + JSON) binary formats.+instance PgDecode Value where+ pgDecode bs+ -- jsonb binary format: first byte is version (0x01), rest is JSON+ | not (BS.null bs) && BS.index bs 0 == 1 =+ case eitherDecodeStrict' (BS.drop 1 bs) of+ Left err -> Left ("jsonb decode: " <> err)+ Right v -> Right v+ -- json binary format: raw JSON bytes+ | otherwise =+ case eitherDecodeStrict' bs of+ Left err -> Left ("json decode: " <> err)+ Right v -> Right v+ {-# INLINE pgDecode #-}
+ src/Valiant/Binary/MacAddr.hs view
@@ -0,0 +1,65 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary codec for PostgreSQL @macaddr@ and @macaddr8@ types.+--+-- PG binary format:+-- @macaddr@: 6 raw bytes (OID 829)+-- @macaddr8@: 8 raw bytes (OID 774)+module Valiant.Binary.MacAddr+ ( PgMacAddr (..)+ , macAddr+ , macAddrToText+ ) where++import Control.DeepSeq (NFData)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Text (Text)+import Data.Text qualified as T+import Data.Word (Word8)+import GHC.Generics (Generic)+import Numeric (showHex)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidMacAddr)++-- | A PostgreSQL @macaddr@ value (6 bytes) or @macaddr8@ value (8 bytes).+data PgMacAddr = PgMacAddr+ { macAddrBytes :: !ByteString+ -- ^ Raw MAC address bytes: 6 bytes for @macaddr@, 8 bytes for @macaddr8@.+ }+ deriving stock (Eq, Ord, Generic)++instance NFData PgMacAddr++instance Show PgMacAddr where+ show = T.unpack . macAddrToText++-- | Construct a @macaddr@ from 6 bytes.+macAddr :: Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> Word8 -> PgMacAddr+macAddr a b c d e f = PgMacAddr (BS.pack [a, b, c, d, e, f])++-- | Render a 'PgMacAddr' as colon-separated hex text.+--+-- >>> macAddrToText (macAddr 0x08 0x00 0x2b 0x01 0x02 0x03)+-- "08:00:2b:01:02:03"+macAddrToText :: PgMacAddr -> Text+macAddrToText (PgMacAddr bs) =+ T.intercalate ":" $+ map (\i -> let w = BS.index bs i in T.pack (pad2 (showHex w "")))+ [0 .. BS.length bs - 1]+ where+ pad2 [c] = ['0', c]+ pad2 s = s++instance PgEncode PgMacAddr where+ pgEncode (PgMacAddr bs) = bs+ {-# INLINE pgEncode #-}+ pgOid _ = oidMacAddr+ {-# INLINE pgOid #-}++instance PgDecode PgMacAddr where+ pgDecode bs+ | len == 6 || len == 8 = Right (PgMacAddr bs)+ | otherwise = Left $ "macaddr: expected 6 or 8 bytes, got " <> show len+ where len = BS.length bs+ {-# INLINE pgDecode #-}
+ src/Valiant/Binary/Point.hs view
@@ -0,0 +1,62 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary codec for PostgreSQL @point@ type (OID 600).+--+-- PG binary format: two float8 values (x, y), 16 bytes total.+module Valiant.Binary.Point+ ( PgPoint (..)+ ) where++import Control.DeepSeq (NFData)+import Data.Bits (shiftL, (.|.))+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.Word (Word64)+import GHC.Float (castWord64ToDouble, castDoubleToWord64)+import GHC.Generics (Generic)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (Oid (..))++-- | A PostgreSQL @point@ value — an (x, y) coordinate pair.+data PgPoint = PgPoint+ { pointX :: {-# UNPACK #-} !Double+ , pointY :: {-# UNPACK #-} !Double+ }+ deriving stock (Show, Eq, Ord, Generic)++instance NFData PgPoint++oidPoint :: Oid+oidPoint = Oid 600++instance PgEncode PgPoint where+ pgEncode (PgPoint x y) =+ LBS.toStrict . B.toLazyByteString $+ B.word64BE (castDoubleToWord64 x) <> B.word64BE (castDoubleToWord64 y)+ {-# INLINE pgEncode #-}+ pgOid _ = oidPoint+ {-# INLINE pgOid #-}++instance PgDecode PgPoint where+ pgDecode bs+ | BS.length bs /= 16 = Left $ "point: expected 16 bytes, got " <> show (BS.length bs)+ | otherwise =+ let x = castWord64ToDouble (decodeWord64BE bs 0)+ y = castWord64ToDouble (decodeWord64BE bs 8)+ in Right (PgPoint x y)+ {-# INLINE pgDecode #-}++decodeWord64BE :: ByteString -> Int -> Word64+decodeWord64BE bs off =+ let b0 = fromIntegral (BS.index bs off) :: Word64+ b1 = fromIntegral (BS.index bs (off + 1)) :: Word64+ b2 = fromIntegral (BS.index bs (off + 2)) :: Word64+ b3 = fromIntegral (BS.index bs (off + 3)) :: Word64+ b4 = fromIntegral (BS.index bs (off + 4)) :: Word64+ b5 = fromIntegral (BS.index bs (off + 5)) :: Word64+ b6 = fromIntegral (BS.index bs (off + 6)) :: Word64+ b7 = fromIntegral (BS.index bs (off + 7)) :: Word64+ in b0 `shiftL` 56 .|. b1 `shiftL` 48 .|. b2 `shiftL` 40 .|. b3 `shiftL` 32+ .|. b4 `shiftL` 24 .|. b5 `shiftL` 16 .|. b6 `shiftL` 8 .|. b7
+ src/Valiant/Binary/Range.hs view
@@ -0,0 +1,135 @@+-- | Binary encode/decode for PostgreSQL range types.+--+-- PG binary format for ranges:+-- 1 byte: flags+-- 0x01 = empty+-- 0x02 = has lower bound+-- 0x04 = has upper bound+-- 0x08 = lower bound inclusive+-- 0x10 = upper bound inclusive+-- if has lower bound:+-- 4 bytes: lower bound data length+-- N bytes: lower bound data+-- if has upper bound:+-- 4 bytes: upper bound data length+-- N bytes: upper bound data+module Valiant.Binary.Range+ ( PgRange (..)+ , RangeBound (..)+ , pgEncodeRange+ , pgDecodeRange+ ) where++import Data.Bits (shiftL, testBit, (.|.))+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 (Int32)+import Data.Word (Word8)++-- | A bound of a range.+data RangeBound a+ = Inclusive a+ | Exclusive a+ deriving stock (Show, Eq, Ord)++-- | A PostgreSQL range value.+data PgRange a+ = EmptyRange+ | PgRange (Maybe (RangeBound a)) (Maybe (RangeBound a))+ -- ^ Lower bound, upper bound. Nothing means unbounded.+ deriving stock (Show, Eq)++-- Flags+flagEmpty, flagHasLower, flagHasUpper, flagLowerInc, flagUpperInc :: Word8+flagEmpty = 0x01+flagHasLower = 0x02+flagHasUpper = 0x04+flagLowerInc = 0x08+flagUpperInc = 0x10++-- | Encode a range value given an element encoder.+pgEncodeRange :: (a -> ByteString) -> PgRange a -> ByteString+pgEncodeRange _ EmptyRange =+ BS.singleton flagEmpty+pgEncodeRange enc (PgRange mLower mUpper) =+ LBS.toStrict . B.toLazyByteString $+ B.word8 flags+ <> encodeBound enc mLower+ <> encodeBound enc mUpper+ where+ flags =+ (if hasBound mLower then flagHasLower else 0)+ .|. (if hasBound mUpper then flagHasUpper else 0)+ .|. (if isInclusive mLower then flagLowerInc else 0)+ .|. (if isInclusive mUpper then flagUpperInc else 0)++ hasBound Nothing = False+ hasBound (Just _) = True++ isInclusive (Just (Inclusive _)) = True+ isInclusive _ = False++encodeBound :: (a -> ByteString) -> Maybe (RangeBound a) -> B.Builder+encodeBound _ Nothing = mempty+encodeBound enc (Just bound) =+ let bs = enc (boundValue bound)+ in B.int32BE (fromIntegral (BS.length bs)) <> B.byteString bs++boundValue :: RangeBound a -> a+boundValue (Inclusive a) = a+boundValue (Exclusive a) = a++-- | Decode a range value given an element decoder.+pgDecodeRange :: (ByteString -> Either String a) -> ByteString -> Either String (PgRange a)+pgDecodeRange _ bs+ | BS.null bs = Left "range: empty input"+pgDecodeRange dec bs = do+ let flags = BS.index bs 0+ if testBit flags 0+ then Right EmptyRange+ else do+ let hasLower = testBit flags 1+ hasUpper = testBit flags 2+ lowerInc = testBit flags 3+ upperInc = testBit flags 4+ (mLower, off1) <-+ if hasLower+ then do+ (val, off) <- decodeBoundAt dec bs 1+ let bound = if lowerInc then Inclusive val else Exclusive val+ Right (Just bound, off)+ else Right (Nothing, 1)+ (mUpper, _) <-+ if hasUpper+ then do+ (val, off) <- decodeBoundAt dec bs off1+ let bound = if upperInc then Inclusive val else Exclusive val+ Right (Just bound, off)+ else Right (Nothing, off1)+ Right (PgRange mLower mUpper)++decodeBoundAt :: (ByteString -> Either String a) -> ByteString -> Int -> Either String (a, Int)+decodeBoundAt dec bs off+ | off + 4 > BS.length bs = Left "range: unexpected end reading bound length"+ | otherwise = do+ len <- decodeInt32At bs off+ let dataOff = off + 4+ dataEnd = dataOff + fromIntegral len+ if dataEnd > BS.length bs+ then Left "range: unexpected end reading bound data"+ else do+ let val = BS.take (fromIntegral len) (BS.drop dataOff bs)+ a <- dec val+ Right (a, dataEnd)++decodeInt32At :: ByteString -> Int -> Either String Int32+decodeInt32At bs off+ | off + 4 > BS.length bs = Left "range: int32 out of bounds"+ | otherwise =+ 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 Right (b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3)
+ src/Valiant/Binary/Scientific.hs view
@@ -0,0 +1,147 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary encode/decode for PostgreSQL @numeric@ type.+--+-- PG stores numeric in a base-10000 representation:+-- 2 bytes: ndigits (number of base-10000 digits)+-- 2 bytes: weight (exponent of first digit, in base-10000)+-- 2 bytes: sign (0x0000 = positive, 0x4000 = negative, 0xC000 = NaN)+-- 2 bytes: dscale (number of digits after decimal point)+-- ndigits * 2 bytes: base-10000 digits (0..9999 each, big-endian Int16)+module Valiant.Binary.Scientific+ () where++import Data.Bits (shiftL, (.|.))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as B+import Data.ByteString.Char8 qualified as BS8+import Data.ByteString.Lazy qualified as LBS+import Data.Int (Int16)+import Data.Scientific (Scientific)+import Data.Scientific qualified as Sci+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidNumeric)++-- Constants+signPositive, signNegative, signNaN :: Int16+signPositive = 0x0000+signNegative = 0x4000+signNaN = fromIntegral (0xC000 :: Int)++-- Encoding ------------------------------------------------------------------++instance PgEncode Scientific where+ pgEncode = encodeNumeric+ pgOid _ = oidNumeric++encodeNumeric :: Scientific -> ByteString+encodeNumeric sci+ | sci == 0 =+ LBS.toStrict . B.toLazyByteString $+ int16 0 <> int16 0 <> int16 signPositive <> int16 0+ | otherwise =+ let (isNeg, absSci) = if sci < 0 then (True, negate sci) else (False, sci)+ !sign = if isNeg then signNegative else signPositive++ c = Sci.coefficient absSci+ e = Sci.base10Exponent absSci++ -- Digit bytes of the absolute coefficient+ !coeffBs = BS8.pack (show (abs c))+ !coeffLen = BS.length coeffBs++ !digitsBeforeDecimal = max 0 (coeffLen + e)+ !digitsAfterDecimal = max 0 (negate e)++ -- Number of base-10000 groups+ !groupsBefore = (digitsBeforeDecimal + 3) `div` 4+ !padBefore = groupsBefore * 4 - digitsBeforeDecimal+ !groupsAfter = (digitsAfterDecimal + 3) `div` 4+ !padAfter = groupsAfter * 4 - digitsAfterDecimal++ -- Build padded digit bytes: [padBefore zeros] [intPart] [fracPart] [padAfter zeros]+ (!intPart, !fracPart)+ | e >= 0 = (coeffBs <> BS.replicate e 0x30, BS.empty) -- 0x30 = '0'+ | otherwise =+ let splitPos = coeffLen + e+ in if splitPos > 0+ then (BS.take splitPos coeffBs, BS.drop splitPos coeffBs)+ else (BS.empty, BS.replicate (negate splitPos) 0x30 <> coeffBs)++ !allDigits = BS.replicate padBefore 0x30 <> intPart <> fracPart <> BS.replicate padAfter 0x30++ -- Convert to base-10000 groups+ !groups = toBase10000Groups allDigits++ -- Strip trailing zero groups+ !groups' = reverse (dropWhile (== 0) (reverse groups))+ !ndigits' = length groups'+ !weight = fromIntegral (groupsBefore - 1) :: Int16+ !dscale = fromIntegral digitsAfterDecimal :: Int16++ in LBS.toStrict . B.toLazyByteString $+ int16 (fromIntegral ndigits')+ <> int16 weight+ <> int16 sign+ <> int16 dscale+ <> mconcat (map (int16 . fromIntegral) groups')++-- | Convert a ByteString of ASCII digit bytes to base-10000 groups.+-- Processes 4 bytes at a time.+toBase10000Groups :: ByteString -> [Int]+toBase10000Groups bs+ | BS.null bs = []+ | otherwise =+ let (!chunk, !rest) = BS.splitAt 4 bs+ !val = BS.foldl' (\acc w -> acc * 10 + fromIntegral (w - 0x30)) 0 chunk+ in val : toBase10000Groups rest++int16 :: Int16 -> B.Builder+int16 = B.int16BE+{-# INLINE int16 #-}++-- Decoding ------------------------------------------------------------------++instance PgDecode Scientific where+ pgDecode = decodeNumeric++decodeNumeric :: ByteString -> Either String Scientific+decodeNumeric bs+ | BS.length bs < 8 = Left "numeric: header too short (need 8 bytes)"+ | otherwise = do+ ndigits <- getInt16 bs 0+ weight <- getInt16 bs 2+ sign <- getInt16 bs 4+ _dscale <- getInt16 bs 6+ let expectedLen = 8 + fromIntegral ndigits * 2+ if BS.length bs < expectedLen+ then Left $ "numeric: payload too short (need " <> show expectedLen <> " bytes)"+ else do+ when (sign == signNaN) $ Left "numeric: NaN not representable as Scientific"+ digits <- mapM (\i -> getInt16 bs (8 + i * 2)) [0 .. fromIntegral ndigits - 1]+ let !isNeg = sign == signNegative+ !intVal = foldlStrict (\acc d -> acc * 10000 + fromIntegral d) (0 :: Integer) digits+ !expo = 4 * (fromIntegral weight - fromIntegral ndigits + 1)+ !sci = Sci.scientific (if isNeg then negate intVal else intVal) expo+ Right sci++-- | Strict left fold over a list (avoids importing Data.List just for foldl').+foldlStrict :: (b -> a -> b) -> b -> [a] -> b+foldlStrict _ !z [] = z+foldlStrict f !z (x : xs) = foldlStrict f (f z x) xs+{-# INLINE foldlStrict #-}++getInt16 :: ByteString -> Int -> Either String Int16+getInt16 bs off+ | off + 2 > BS.length bs = Left "numeric: getInt16 out of bounds"+ | otherwise =+ let hi = fromIntegral (BS.index bs off) :: Int16+ lo = fromIntegral (BS.index bs (off + 1)) :: Int16+ in Right (hi `shiftL` 8 .|. lo)+{-# INLINE getInt16 #-}++when :: Bool -> Either String () -> Either String ()+when True e = e+when False _ = Right ()+{-# INLINE when #-}
+ src/Valiant/Binary/UUID.hs view
@@ -0,0 +1,29 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Binary encode/decode for PostgreSQL @uuid@ type.+--+-- PG stores UUID as 16 bytes in network byte order, which is the same+-- as the RFC 4122 binary representation that uuid-types uses.+module Valiant.Binary.UUID+ () where++import Data.ByteString qualified as BS+import Data.ByteString.Lazy qualified as LBS+import Data.UUID.Types (UUID)+import Data.UUID.Types qualified as UUID+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidUuid)++instance PgEncode UUID where+ pgEncode = LBS.toStrict . UUID.toByteString+ {-# INLINE pgEncode #-}+ pgOid _ = oidUuid+ {-# INLINE pgOid #-}++instance PgDecode UUID where+ pgDecode bs+ | BS.length bs /= 16 = Left $ "uuid: expected 16 bytes, got " <> show (BS.length bs)+ | otherwise = case UUID.fromByteString (LBS.fromStrict bs) of+ Nothing -> Left "uuid: invalid bytes"+ Just u -> Right u+ {-# INLINE pgDecode #-}
+ src/Valiant/Binary/Unbounded.hs view
@@ -0,0 +1,167 @@+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Support for PostgreSQL @-infinity@ and @infinity@ timestamp/date values.+--+-- PostgreSQL allows @-infinity@ and @infinity@ as valid values for+-- @timestamp@, @timestamptz@, and @date@ columns. In binary format these+-- are represented as sentinel values:+--+-- * Timestamps: @INT64_MIN@ for @-infinity@, @INT64_MAX@ for @infinity@+-- * Dates: @INT32_MIN@ for @-infinity@, @INT32_MAX@ for @infinity@+--+-- The standard 'UTCTime', 'LocalTime', and 'Day' types cannot represent+-- infinity. This module provides 'Unbounded' as a wrapper:+--+-- @+-- findUser :: Statement Int32 (Int32, Text, Unbounded UTCTime)+-- findUser = queryFile \"users\/find_by_id.sql\"+-- @+module Valiant.Binary.Unbounded+ ( Unbounded (..)+ , finite+ , fromFinite+ ) where++import Control.DeepSeq (NFData (..))+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.Bits (shiftL, (.|.))+import Data.Int (Int32, Int64)+import Data.Time (Day, LocalTime, UTCTime)+import GHC.Generics (Generic)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidDate, oidTimestamp, oidTimestamptz)+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()++-- | A value that may be @-infinity@, a finite value, or @infinity@.+-- Used for PostgreSQL timestamp and date types that support infinity.+data Unbounded a+ = NegInfinity+ -- ^ @-infinity@+ | Finite !a+ -- ^ A regular finite value.+ | PosInfinity+ -- ^ @infinity@+ deriving stock (Show, Eq, Ord, Generic, Functor)++instance NFData a => NFData (Unbounded a)++-- | Wrap a value as 'Finite'.+finite :: a -> Unbounded a+finite = Finite+{-# INLINE finite #-}++-- | Extract the finite value, or 'Nothing' for infinity.+fromFinite :: Unbounded a -> Maybe a+fromFinite (Finite a) = Just a+fromFinite _ = Nothing+{-# INLINE fromFinite #-}++-- Sentinel values used by PostgreSQL binary format.+pgTimestampPosInf :: Int64+pgTimestampPosInf = 0x7FFFFFFFFFFFFFFF++pgTimestampNegInf :: Int64+pgTimestampNegInf = -0x8000000000000000++pgDatePosInf :: Int32+pgDatePosInf = 0x7FFFFFFF++pgDateNegInf :: Int32+pgDateNegInf = -0x80000000++------------------------------------------------------------------------+-- Unbounded UTCTime (timestamptz)+------------------------------------------------------------------------++instance PgEncode (Unbounded UTCTime) where+ pgEncode NegInfinity = encodeInt64BE pgTimestampNegInf+ pgEncode PosInfinity = encodeInt64BE pgTimestampPosInf+ pgEncode (Finite t) = pgEncode t+ {-# INLINE pgEncode #-}+ pgOid _ = oidTimestamptz+ {-# INLINE pgOid #-}++instance PgDecode (Unbounded UTCTime) where+ pgDecode bs = do+ micros <- decodeRawInt64 bs+ if micros == pgTimestampPosInf then Right PosInfinity+ else if micros == pgTimestampNegInf then Right NegInfinity+ else Finite <$> pgDecode bs+ {-# INLINE pgDecode #-}++------------------------------------------------------------------------+-- Unbounded LocalTime (timestamp)+------------------------------------------------------------------------++instance PgEncode (Unbounded LocalTime) where+ pgEncode NegInfinity = encodeInt64BE pgTimestampNegInf+ pgEncode PosInfinity = encodeInt64BE pgTimestampPosInf+ pgEncode (Finite t) = pgEncode t+ {-# INLINE pgEncode #-}+ pgOid _ = oidTimestamp+ {-# INLINE pgOid #-}++instance PgDecode (Unbounded LocalTime) where+ pgDecode bs = do+ micros <- decodeRawInt64 bs+ if micros == pgTimestampPosInf then Right PosInfinity+ else if micros == pgTimestampNegInf then Right NegInfinity+ else Finite <$> pgDecode bs+ {-# INLINE pgDecode #-}++------------------------------------------------------------------------+-- Unbounded Day (date)+------------------------------------------------------------------------++instance PgEncode (Unbounded Day) where+ pgEncode NegInfinity = encodeInt32BE pgDateNegInf+ pgEncode PosInfinity = encodeInt32BE pgDatePosInf+ pgEncode (Finite d) = pgEncode d+ {-# INLINE pgEncode #-}+ pgOid _ = oidDate+ {-# INLINE pgOid #-}++instance PgDecode (Unbounded Day) where+ pgDecode bs = do+ days <- decodeRawInt32 bs+ if days == pgDatePosInf then Right PosInfinity+ else if days == pgDateNegInf then Right NegInfinity+ else Finite <$> pgDecode bs+ {-# INLINE pgDecode #-}++------------------------------------------------------------------------+-- Helpers+------------------------------------------------------------------------++decodeRawInt64 :: ByteString -> Either String Int64+decodeRawInt64 bs+ | BS.length bs /= 8 = Left $ "Unbounded: expected 8 bytes, got " <> show (BS.length bs)+ | otherwise =+ let b0 = fromIntegral (BS.index bs 0) :: Int64+ b1 = fromIntegral (BS.index bs 1) :: Int64+ b2 = fromIntegral (BS.index bs 2) :: Int64+ b3 = fromIntegral (BS.index bs 3) :: Int64+ b4 = fromIntegral (BS.index bs 4) :: Int64+ b5 = fromIntegral (BS.index bs 5) :: Int64+ b6 = fromIntegral (BS.index bs 6) :: Int64+ b7 = fromIntegral (BS.index bs 7) :: Int64+ in Right $ b0 `shiftL` 56 .|. b1 `shiftL` 48 .|. b2 `shiftL` 40 .|. b3 `shiftL` 32+ .|. b4 `shiftL` 24 .|. b5 `shiftL` 16 .|. b6 `shiftL` 8 .|. b7++decodeRawInt32 :: ByteString -> Either String Int32+decodeRawInt32 bs+ | BS.length bs /= 4 = Left $ "Unbounded: expected 4 bytes, got " <> show (BS.length bs)+ | otherwise =+ 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 Right $ b0 `shiftL` 24 .|. b1 `shiftL` 16 .|. b2 `shiftL` 8 .|. b3++encodeInt64BE :: Int64 -> ByteString+encodeInt64BE = pgEncode++encodeInt32BE :: Int32 -> ByteString+encodeInt32BE = pgEncode
+ src/Valiant/Copy.hs view
@@ -0,0 +1,204 @@+{-# OPTIONS_GHC -fno-full-laziness #-}++-- | COPY protocol support for bulk data import/export.+--+-- COPY operations require exclusive access to the connection's wire,+-- so they use 'submitExclusive' to pause the async pipeline.+module Valiant.Copy+ ( copyIn+ , copyOut+ , copyInBinary+ , CopyResult (..)+ ) where++import Control.Exception (SomeException, catch, throwIO)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Builder qualified as B+import Data.ByteString.Char8 qualified as BS8+import Data.ByteString.Lazy qualified as LBS+import Data.IORef+import Data.Int (Int16, Int64)+import Data.Vector (Vector)+import Data.Vector qualified as V+import PgWire.Async (submitExclusive)+import PgWire.Connection (Connection (..))+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend+import PgWire.Protocol.Frontend+import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsg)++-- | Result of a COPY operation, containing the number of rows transferred.+data CopyResult = CopyResult+ { copyRows :: Int64+ -- ^ Number of rows copied, as reported by the server's @CommandComplete@ tag.+ }+ deriving stock (Show, Eq)++-- | Execute a @COPY ... FROM STDIN@ command, sending data in chunks.+--+-- Usage:+--+-- > copyIn conn "COPY users (name, email) FROM STDIN WITH (FORMAT csv)" $ \sendChunk -> do+-- > sendChunk "Alice,alice@example.com\n"+-- > sendChunk "Bob,bob@example.com\n"+copyIn :: Connection -> ByteString -> ((ByteString -> IO ()) -> IO ()) -> IO CopyResult+copyIn conn sql producer =+ submitExclusive (connAsync conn) $ \wc txRef -> do+ sendFrontendMsg wc (Query sql)+ waitCopyIn wc+ producer (sendFrontendMsg wc . CopyData)+ `catch` \(e :: SomeException) -> do+ -- Best-effort: send CopyFail and drain to ReadyForQuery so the+ -- connection is usable. If the wire is dead, ignore the cleanup+ -- failure and re-throw the original exception.+ (do sendFrontendMsg wc (CopyFail (BS8.pack (show e)))+ drainToReady wc txRef+ ) `catch` \(_ :: SomeException) -> pure ()+ throwIO e+ sendFrontendMsg wc CopyDone+ collectCopyResult wc txRef++-- | Execute a @COPY ... TO STDOUT@ command, receiving data in chunks+-- via a callback.+--+-- Usage:+--+-- > copyOut conn "COPY users TO STDOUT WITH (FORMAT csv)" $ \chunk -> do+-- > BS.putStr chunk+copyOut :: Connection -> ByteString -> (ByteString -> IO ()) -> IO CopyResult+copyOut conn sql consumer =+ submitExclusive (connAsync conn) $ \wc txRef -> do+ sendFrontendMsg wc (Query sql)+ waitCopyOut wc+ let loop = do+ msg <- recvBackendMsg wc+ case msg of+ CopyDataMsg chunk -> do+ consumer chunk+ loop+ CopyDoneMsg -> pure ()+ ErrorResponse err -> throwPgWire (QueryError err)+ other -> throwPgWire (ProtocolError ("Unexpected in COPY OUT: " <> BS8.pack (show other)))+ loop+ collectCopyResult wc txRef++-- | Execute a @COPY ... FROM STDIN WITH (FORMAT binary)@ command,+-- sending rows in PostgreSQL's binary COPY format.+--+-- Each row is provided as a @Vector (Maybe ByteString)@ where each+-- element is an already-encoded binary value (using 'pgEncode'), or+-- 'Nothing' for NULL.+--+-- @+-- copyInBinary conn+-- \"COPY users (id, name) FROM STDIN WITH (FORMAT binary)\"+-- 2 -- number of columns+-- $ \\sendRow -> do+-- sendRow (V.fromList [Just (pgEncode (1 :: Int32)), Just (pgEncode (\"Alice\" :: Text))])+-- sendRow (V.fromList [Just (pgEncode (2 :: Int32)), Just (pgEncode (\"Bob\" :: Text))])+-- @+copyInBinary+ :: Connection+ -> ByteString+ -- ^ COPY ... FROM STDIN WITH (FORMAT binary) statement+ -> Int16+ -- ^ Number of columns+ -> ((Vector (Maybe ByteString) -> IO ()) -> IO ())+ -- ^ Producer: call @sendRow@ for each row+ -> IO CopyResult+copyInBinary conn sql numCols producer =+ submitExclusive (connAsync conn) $ \wc txRef -> do+ sendFrontendMsg wc (Query sql)+ waitCopyIn wc+ sendFrontendMsg wc (CopyData binaryCopyHeader)+ let sendRow row = do+ let rowBytes = encodeBinaryRow numCols row+ sendFrontendMsg wc (CopyData rowBytes)+ producer sendRow+ `catch` \(e :: SomeException) -> do+ (do sendFrontendMsg wc (CopyFail (BS8.pack (show e)))+ drainToReady wc txRef+ ) `catch` \(_ :: SomeException) -> pure ()+ throwIO e+ sendFrontendMsg wc (CopyData binaryCopyTrailer)+ sendFrontendMsg wc CopyDone+ collectCopyResult wc txRef++-- | Binary COPY header:+-- 11-byte signature: "PGCOPY\n\377\r\n\0"+-- 4-byte flags: 0 (no OID inclusion)+-- 4-byte header extension area length: 0+binaryCopyHeader :: ByteString+binaryCopyHeader = LBS.toStrict . B.toLazyByteString $+ B.byteString "PGCOPY\n\xff\r\n\0"+ <> B.int32BE 0 -- flags+ <> B.int32BE 0 -- header extension length++-- | Binary COPY trailer: -1 as Int16+binaryCopyTrailer :: ByteString+binaryCopyTrailer = LBS.toStrict . B.toLazyByteString $+ B.int16BE (-1)++-- | Encode a single row in binary COPY format:+-- Int16 field count, then per field: Int32 length (-1 for NULL) + data+encodeBinaryRow :: Int16 -> Vector (Maybe ByteString) -> ByteString+encodeBinaryRow numCols row = LBS.toStrict . B.toLazyByteString $+ B.int16BE numCols+ <> V.foldl' (\acc mv -> acc <> encodeField mv) mempty row+ where+ encodeField Nothing = B.int32BE (-1)+ encodeField (Just bs) = B.int32BE (fromIntegral (BS.length bs)) <> B.byteString bs++-- Internal ------------------------------------------------------------------++waitCopyIn :: WireConn -> IO ()+waitCopyIn wc = do+ msg <- recvBackendMsg wc+ case msg of+ CopyInResponse _ _ -> pure ()+ ErrorResponse err -> throwPgWire (QueryError err)+ other -> throwPgWire (ProtocolError ("Expected CopyInResponse, got: " <> BS8.pack (show other)))++waitCopyOut :: WireConn -> IO ()+waitCopyOut wc = do+ msg <- recvBackendMsg wc+ case msg of+ CopyOutResponse _ _ -> pure ()+ ErrorResponse err -> throwPgWire (QueryError err)+ other -> throwPgWire (ProtocolError ("Expected CopyOutResponse, got: " <> BS8.pack (show other)))++collectCopyResult :: WireConn -> IORef TxStatus -> IO CopyResult+collectCopyResult wc txRef = go 0+ where+ go !n = do+ msg <- recvBackendMsg wc+ case msg of+ CommandComplete tag -> go (tagRows tag)+ ReadyForQuery status -> do+ writeIORef txRef status+ pure (CopyResult n)+ ErrorResponse err -> throwPgWire (QueryError err)+ NoticeResponse _ -> go n+ other -> throwPgWire (ProtocolError ("Unexpected in COPY result: " <> BS8.pack (show other)))++ tagRows (InsertTag r) = r+ tagRows (UpdateTag r) = r+ tagRows (DeleteTag r) = r+ tagRows (SelectTag r) = r+ tagRows (OtherTag t) =+ let parts = BS8.words t+ in case parts of+ [_, numStr] -> maybe 0 (fromIntegral . fst) (BS8.readInt numStr)+ _ -> 0++-- | Drain messages until ReadyForQuery, ignoring errors.+-- Used after CopyFail to ensure the connection is in a clean state.+drainToReady :: WireConn -> IORef TxStatus -> IO ()+drainToReady wc txRef = go+ where+ go = do+ msg <- recvBackendMsg wc+ case msg of+ ReadyForQuery status -> writeIORef txRef status+ _ -> go
+ src/Valiant/Dynamic.hs view
@@ -0,0 +1,140 @@+{-# LANGUAGE TupleSections #-}++-- | Dynamic SQL statement builder for queries that cannot be expressed+-- as static @.sql@ files.+--+-- 'Snippet' is a composable SQL fragment with embedded typed parameters.+-- Use it for dynamic WHERE clauses, optional filters, dynamic ORDER BY,+-- and other cases where the SQL shape varies at runtime.+--+-- @+-- import Valiant.Dynamic+--+-- searchUsers :: Maybe Text -> Maybe Bool -> Connection -> IO [(Int32, Text)]+-- searchUsers mName mActive conn = do+-- let base = sql "SELECT id, name FROM users WHERE true"+-- nameFilter = case mName of+-- Just n -> sql " AND name ILIKE " <> param n+-- Nothing -> mempty+-- activeFilter = case mActive of+-- Just a -> sql " AND is_active = " <> param a+-- Nothing -> mempty+-- query = base <> nameFilter <> activeFilter <> sql " ORDER BY name"+-- runSnippet conn query+-- @+module Valiant.Dynamic+ ( Snippet+ , sql+ , param+ , runSnippet+ , runSnippetOne+ , snippetToSQL+ ) where++import Data.ByteString (ByteString)+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.Word (Word32)+import Valiant.FromRow (FromRow (..))+import PgWire.Binary.Types (PgEncode (..))+import PgWire.Connection (Connection)+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Oid (Oid (..))+import qualified Valiant.Execute as Execute+import Data.Proxy (Proxy (..))++-- | A composable SQL fragment with embedded parameters.+-- 'Snippet' is a 'Monoid', so fragments can be combined with '<>'.+data Snippet = Snippet+ { snipBuilder :: Int -> (Builder, Int)+ -- ^ Given the next parameter index, returns (SQL builder, next index after).+ -- Parameter placeholders are $1, $2, etc.+ , snipParams :: [(Maybe ByteString, Word32)]+ -- ^ Accumulated parameters in reverse order: (encoded value, OID).+ }++instance Semigroup Snippet where+ Snippet b1 p1 <> Snippet b2 p2 = Snippet+ { snipBuilder = \n ->+ let (sql1, n1) = b1 n+ (sql2, n2) = b2 n1+ in (sql1 <> sql2, n2)+ , snipParams = p2 ++ p1+ }+ {-# INLINE (<>) #-}++instance Monoid Snippet where+ mempty = Snippet { snipBuilder = (mempty,), snipParams = [] }+ {-# INLINE mempty #-}++-- | A raw SQL fragment with no parameters.+--+-- @+-- sql "SELECT id, name FROM users WHERE "+-- @+sql :: ByteString -> Snippet+sql s = Snippet+ { snipBuilder = (B.byteString s,)+ , snipParams = []+ }+{-# INLINE sql #-}++-- | Embed a typed parameter. Generates a @$N@ placeholder and captures+-- the encoded value and OID.+--+-- @+-- sql "WHERE id = " <> param (42 :: Int32)+-- -- produces: "WHERE id = $1" with Int32 parameter+-- @+param :: forall a. (PgEncode a) => a -> Snippet+param a = Snippet+ { snipBuilder = \n ->+ let placeholder = B.byteString ("$" <> BS8.pack (show n))+ in (placeholder, n + 1)+ , snipParams = [(Just (pgEncode a), unOid (pgOid (Proxy @a)))]+ }+ where+ unOid (Oid o) = o+{-# INLINE param #-}++-- | Execute a 'Snippet' and return all result rows.+--+-- @+-- rows <- runSnippet conn (sql "SELECT 1 + " <> param (2 :: Int32))+-- @+runSnippet :: (FromRow r) => Connection -> Snippet -> IO [r]+runSnippet conn snip = do+ let (sqlBs, oids, vals) = materialize snip+ rawRows <- Execute.rawFetchAll conn sqlBs (map fromIntegral oids) vals+ mapM (\row -> case fromRow row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure val) rawRows++-- | Execute a 'Snippet' and return the first row, or 'Nothing'.+runSnippetOne :: (FromRow r) => Connection -> Snippet -> IO (Maybe r)+runSnippetOne conn snip = do+ let (sqlBs, oids, vals) = materialize snip+ mRow <- Execute.rawFetchOne conn sqlBs (map fromIntegral oids) vals+ case mRow of+ Nothing -> pure Nothing+ Just row -> case fromRow row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure (Just val)++-- | Render a 'Snippet' to its SQL text (for debugging/logging).+snippetToSQL :: Snippet -> ByteString+snippetToSQL snip =+ let (builder, _) = snipBuilder snip 1+ in LBS.toStrict (B.toLazyByteString builder)++-- Internal: materialize a snippet into SQL + OIDs + values.+materialize :: Snippet -> (ByteString, [Word32], [Maybe ByteString])+materialize snip =+ let (builder, _) = snipBuilder snip 1+ sqlBs = LBS.toStrict (B.toLazyByteString builder)+ paramsOrdered = reverse (snipParams snip)+ oids = map snd paramsOrdered+ vals = map fst paramsOrdered+ in (sqlBs, oids, vals)
+ src/Valiant/Error.hs view
@@ -0,0 +1,191 @@+-- | Structured error inspection and constraint violation helpers.+--+-- PostgreSQL errors carry a SQLSTATE code (5 ASCII characters) that+-- categorifies the error. This module provides helpers for extracting+-- the SQLSTATE, identifying constraint violations, and catching+-- specific error classes.+--+-- See <https://www.postgresql.org/docs/current/errcodes-appendix.html>+-- for the full list of SQLSTATE codes.+--+-- @+-- result <- try (execute conn insertStmt params)+-- case result of+-- Left err+-- | Just (UniqueViolation name) <- constraintViolation err ->+-- handleDuplicate name+-- | isSerializationError err ->+-- retry+-- Right n -> pure n+-- @+module Valiant.Error+ ( -- * SQLSTATE access+ sqlState+ , pgErrorOf++ -- * Constraint violations+ , ConstraintViolation (..)+ , constraintViolation+ , catchConstraintViolation++ -- * Error predicates+ , isUniqueViolation+ , isForeignKeyViolation+ , isNotNullViolation+ , isCheckViolation+ , isExclusionViolation+ , isSerializationError+ , isDeadlockError+ , isNoActiveTransactionError+ , isFailedTransactionError+ , isInvalidTextRepresentation+ , isUndefinedTable+ , isUndefinedColumn+ , isSyntaxError+ ) where++import Control.Exception (catch, throwIO)+import Data.Maybe (fromMaybe)+import Data.ByteString (ByteString)+import PgWire.Error (PgWireError (..))+import PgWire.Protocol.Backend (PgError (..))++-- | Extract the SQLSTATE code from a 'PgWireError', if it wraps a+-- server-side 'QueryError'.+--+-- Returns 'Nothing' for non-query errors (connection errors, decode+-- errors, pool errors, etc.).+sqlState :: PgWireError -> Maybe ByteString+sqlState (QueryError err) = Just (pgCode err)+sqlState _ = Nothing++-- | Extract a field from the underlying 'PgError', if present.+--+-- @+-- mDetail <- pgErrorOf err pgDetail+-- mTable <- pgErrorOf err pgTable+-- @+pgErrorOf :: PgWireError -> (PgError -> a) -> Maybe a+pgErrorOf (QueryError err) f = Just (f err)+pgErrorOf _ _ = Nothing++------------------------------------------------------------------------+-- Constraint violations+------------------------------------------------------------------------++-- | A structured constraint violation, extracted from the SQLSTATE code+-- and constraint name of a PostgreSQL error.+--+-- The 'ByteString' payload is the constraint name (from the @C@ error+-- field), which may be empty if PostgreSQL did not report it.+data ConstraintViolation+ = NotNullViolation !ByteString+ -- ^ @23502@ — a NOT NULL constraint was violated.+ | ForeignKeyViolation !ByteString+ -- ^ @23503@ — a foreign key constraint was violated.+ | UniqueViolation !ByteString+ -- ^ @23505@ — a unique constraint or unique index was violated.+ | CheckViolation !ByteString+ -- ^ @23514@ — a CHECK constraint was violated.+ | ExclusionViolation !ByteString+ -- ^ @23P01@ — an exclusion constraint was violated.+ deriving stock (Show, Eq, Ord)++-- | Extract a 'ConstraintViolation' from a 'PgWireError', if the error+-- is a constraint violation (SQLSTATE class 23).+--+-- Returns 'Nothing' for non-query errors or query errors that are not+-- constraint violations.+constraintViolation :: PgWireError -> Maybe ConstraintViolation+constraintViolation (QueryError err) =+ let code = pgCode err+ name = fromMaybe "" (pgConstraint err)+ in case code of+ "23502" -> Just (NotNullViolation name)+ "23503" -> Just (ForeignKeyViolation name)+ "23505" -> Just (UniqueViolation name)+ "23514" -> Just (CheckViolation name)+ "23P01" -> Just (ExclusionViolation name)+ _ -> Nothing+constraintViolation _ = Nothing++-- | Catch a constraint violation and handle it with a callback.+-- Non-constraint errors are re-thrown.+--+-- @+-- catchConstraintViolation+-- (\\err cv -> case cv of+-- UniqueViolation _ -> pure fallbackValue+-- _ -> throwIO err)+-- (execute conn insertStmt params)+-- @+catchConstraintViolation+ :: (PgWireError -> ConstraintViolation -> IO a)+ -> IO a+ -> IO a+catchConstraintViolation handler action =+ action `catch` \err ->+ case constraintViolation err of+ Just cv -> handler err cv+ Nothing -> throwIO err++------------------------------------------------------------------------+-- Error predicates+------------------------------------------------------------------------++-- | @23505@ — unique constraint or unique index violation.+isUniqueViolation :: PgWireError -> Bool+isUniqueViolation = hasState "23505"++-- | @23503@ — foreign key constraint violation.+isForeignKeyViolation :: PgWireError -> Bool+isForeignKeyViolation = hasState "23503"++-- | @23502@ — NOT NULL constraint violation.+isNotNullViolation :: PgWireError -> Bool+isNotNullViolation = hasState "23502"++-- | @23514@ — CHECK constraint violation.+isCheckViolation :: PgWireError -> Bool+isCheckViolation = hasState "23514"++-- | @23P01@ — exclusion constraint violation.+isExclusionViolation :: PgWireError -> Bool+isExclusionViolation = hasState "23P01"++-- | @40001@ — serialization failure. Retry the transaction.+isSerializationError :: PgWireError -> Bool+isSerializationError = hasState "40001"++-- | @40P01@ — deadlock detected. Retry the transaction.+isDeadlockError :: PgWireError -> Bool+isDeadlockError = hasState "40P01"++-- | @25P01@ — no active transaction (e.g., COMMIT outside a transaction).+isNoActiveTransactionError :: PgWireError -> Bool+isNoActiveTransactionError = hasState "25P01"++-- | @25P02@ — current transaction is aborted, commands ignored until+-- end of transaction block.+isFailedTransactionError :: PgWireError -> Bool+isFailedTransactionError = hasState "25P02"++-- | @22P02@ — invalid input syntax for type (e.g., passing @\"abc\"@ for an integer).+isInvalidTextRepresentation :: PgWireError -> Bool+isInvalidTextRepresentation = hasState "22P02"++-- | @42P01@ — undefined table.+isUndefinedTable :: PgWireError -> Bool+isUndefinedTable = hasState "42P01"++-- | @42703@ — undefined column.+isUndefinedColumn :: PgWireError -> Bool+isUndefinedColumn = hasState "42703"++-- | @42601@ — syntax error in SQL.+isSyntaxError :: PgWireError -> Bool+isSyntaxError = hasState "42601"++-- Internal: check if an error has a specific SQLSTATE code.+hasState :: ByteString -> PgWireError -> Bool+hasState expected err = sqlState err == Just expected
+ src/Valiant/Execute.hs view
@@ -0,0 +1,817 @@+{-# LANGUAGE LambdaCase #-}+{-# OPTIONS_GHC -fno-full-laziness #-}++-- | Query execution functions using the PostgreSQL extended query protocol.+--+-- All functions use binary format for both parameters and results,+-- prepared statement caching, and message coalescing for minimal+-- round-trip overhead.+--+-- When a statement hasn't been prepared yet, Parse is coalesced with+-- Bind+Execute+Sync into a single round-trip (the asyncpg technique).+-- Subsequent executions skip Parse entirely.+--+-- For large batches (>256 items), 'executeBatch' streams Bind+Execute+-- pairs in chunks via exclusive mode to bound memory usage.+module Valiant.Execute+ ( -- * Queries+ fetchOne+ , fetchAll+ , fetchAllVec+ , fetchAllUnboxed+ , fetchAllWith+ , fetchScalar+ , fetchOneOrThrow+ , fetchOneOr+ , fetchFirst+ , fetchExists+ , forEach+ -- * Fast decode variants+ , fetchAllFast+ , fetchOneFast+ -- * Commands+ , execute+ , executeReturning+ , executeReturningMany+ , executeMany+ -- * Raw (unchecked) queries+ , rawFetchAll+ , rawFetchOne+ , rawExecute+ -- * Pipelined batch execution+ , executeBatch+ -- * Pipelined batch reads+ , fetchBatchOne+ , fetchBatchAll+ -- * Internal+ -- | This export is only visible to other modules within @valiant@+ -- (Pipeline, Fold). It is not part of the stable public API; do not+ -- rely on it from downstream code.+ , ensurePrepared+ ) where++import Control.Monad (when)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Int (Int64)+import Data.HashPSQ (HashPSQ)+import Data.HashPSQ qualified as PSQ+import Data.Maybe (fromMaybe)+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Vector.Unboxed qualified as U+import Data.Word (Word32, Word64)+import PgWire.Async (Request (..), Response (..), ResponseCollector (..), submitRequest, submitExclusive)+import PgWire.Connection (Connection (..))+import PgWire.Protocol.Oid qualified as Oid+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend+import PgWire.Protocol.Frontend+import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsg, sendFrontendMsgs)+import Valiant.FromRowFast (FromRowFast (..))+import Valiant.Statement (Statement (..))++------------------------------------------------------------------------+-- Shared constants (avoid per-call allocation)+------------------------------------------------------------------------++-- | @[BinaryFormat]@ — used for both parameter and result format codes.+binaryFmtVec :: Vector FormatCode+binaryFmtVec = V.singleton BinaryFormat+{-# NOINLINE binaryFmtVec #-}++------------------------------------------------------------------------+-- Queries+------------------------------------------------------------------------++-- | Fetch zero or one row from a typed 'Statement'.+-- Returns 'Nothing' if the query produces no results; if multiple rows are+-- returned, only the first is used (remaining rows are discarded at the+-- wire level without decoding).+--+-- @+-- mUser <- fetchOne conn findById 42+-- @+fetchOne :: Connection -> Statement p r -> p -> IO (Maybe r)+fetchOne conn stmt params = do+ mRow <- fetchFirstRowRaw conn stmt params+ case mRow of+ Nothing -> pure Nothing+ Just row -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure (Just val)++-- | Fetch all result rows as a list.+-- Decodes every row eagerly; for large result sets consider streaming instead.+--+-- @+-- users <- fetchAll conn listAllUsers ()+-- @+fetchAll :: Connection -> Statement p r -> p -> IO [r]+fetchAll conn stmt params = do+ rows <- fetchRowsRaw conn stmt params+ decodeRows (stmtDecode stmt) rows++-- | Fetch exactly one row and decode it as a scalar value.+-- Throws 'DecodeError' if the query returns zero rows.+--+-- @+-- n <- fetchScalar conn countUsers ()+-- @+fetchScalar :: Connection -> Statement p r -> p -> IO r+fetchScalar conn stmt params = do+ mRow <- fetchFirstRowRaw conn stmt params+ case mRow of+ Nothing -> throwPgWire (DecodeError "fetchScalar: query returned no rows")+ Just row -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure val++-- | Like 'fetchOne' but throws 'DecodeError' if no rows are returned.+-- Useful when you know the row must exist (e.g., fetching by primary key+-- after a successful INSERT...RETURNING).+--+-- @+-- user <- fetchOneOrThrow conn findById userId+-- @+fetchOneOrThrow :: Connection -> Statement p r -> p -> IO r+fetchOneOrThrow conn stmt params = do+ mResult <- fetchOne conn stmt params+ case mResult of+ Nothing -> throwPgWire (DecodeError "fetchOneOrThrow: query returned no rows")+ Just val -> pure val++-- | Check whether a query returns any rows. Useful for @EXISTS@-style queries.+-- More efficient than 'fetchAll' since it collects at most one row at the+-- wire level and does not decode any row data.+--+-- @+-- exists <- fetchExists conn userExistsById 42+-- @+fetchExists :: Connection -> Statement p r -> p -> IO Bool+fetchExists conn stmt params = do+ mRow <- fetchFirstRowRaw conn stmt params+ pure (case mRow of { Nothing -> False; Just _ -> True })++-- | Fetch all result rows as a 'Vector'. Uses a dedicated wire-level+-- collector that builds a 'Vector' directly via a growable mutable+-- buffer, avoiding the intermediate list that 'fetchAll' builds.+--+-- @+-- users <- fetchAllVec conn listAllUsers ()+-- @+fetchAllVec :: Connection -> Statement p r -> p -> IO (Vector r)+fetchAllVec conn stmt params = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encodedParams = stmtEncode stmt params+ bindExec =+ [ Bind "" name binaryFmtVec encodedParams binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExec+ else bindExec+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectRowsVec+ case resp of+ RespRowsVec rawVec -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ V.mapM (\row -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> pure val) rawVec+ _ -> throwPgWire (ProtocolError "fetchAllVec: unexpected response type")++-- | Like 'fetchAllVec' but stores results in an unboxed 'Data.Vector.Unboxed.Vector'.+-- Eliminates per-element pointer indirection for fixed-size primitive+-- result types, reducing memory by 2-4x for queries returning large+-- numbers of @Int32@, @Int64@, @Double@, scalar pairs, etc.+--+-- The result type must have an 'U.Unbox' instance. Common cases like+-- 'Int32', 'Int64', 'Double', 'Bool', and tuples of those work out of+-- the box; for newtype wrappers use @-XDeriveAnyClass@\/@-XDerivingVia@+-- with @Data.Vector.Unboxed.Deriving@.+--+-- @+-- ids :: U.Vector Int32 <- fetchAllUnboxed conn allUserIds ()+-- @+fetchAllUnboxed+ :: U.Unbox r => Connection -> Statement p r -> p -> IO (U.Vector r)+fetchAllUnboxed conn stmt params = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encodedParams = stmtEncode stmt params+ bindExec =+ [ Bind "" name binaryFmtVec encodedParams binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExec+ else bindExec+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectRowsVec+ case resp of+ RespRowsVec rawVec -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ let decode = stmtDecode stmt+ U.generateM (V.length rawVec) $ \i ->+ case decode (V.unsafeIndex rawVec i) of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> pure val+ _ -> throwPgWire (ProtocolError "fetchAllUnboxed: unexpected response type")++-- | Like 'fetchAll' but applies a transformation to each decoded row.+-- Useful for mapping database rows to domain types without an intermediate list.+--+-- @+-- names <- fetchAllWith conn listUsers () userName+-- @+fetchAllWith :: Connection -> Statement p r -> p -> (r -> a) -> IO [a]+fetchAllWith conn stmt params f = do+ rows <- fetchAll conn stmt params+ pure (map f rows)++-- | Like 'fetchOne' but returns a default value instead of 'Nothing'+-- when no rows are returned.+--+-- @+-- count <- fetchOneOr conn countByStatus "active" 0+-- @+fetchOneOr :: Connection -> Statement p r -> p -> r -> IO r+fetchOneOr conn stmt params def = do+ mResult <- fetchOne conn stmt params+ pure (fromMaybe def mResult)++-- | Fetch the first row from a query that may return multiple rows.+-- Equivalent to 'fetchOne' but communicates intent more clearly when the+-- query is known to return multiple rows and you want only the first.+--+-- @+-- newest <- fetchFirst conn listRecentUsers ()+-- @+fetchFirst :: Connection -> Statement p r -> p -> IO (Maybe r)+fetchFirst = fetchOne++-- | Execute a callback for each result row. Rows are decoded and+-- dispatched one at a time directly from the wire — no intermediate+-- list is ever built. Uses exclusive wire access for true streaming.+--+-- @+-- forEach conn listAllUsers () $ \\user ->+-- putStrLn (userName user)+-- @+forEach :: Connection -> Statement p r -> p -> (r -> IO ()) -> IO ()+forEach conn stmt params action = do+ stmtName <- ensurePrepared conn stmt+ let encodedParams = stmtEncode stmt params+ submitExclusive (connAsync conn) $ \wc txRef -> do+ sendFrontendMsgs wc+ [ Bind "" stmtName binaryFmtVec encodedParams binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ let go = do+ msg <- recvBackendMsg wc+ case msg of+ BindComplete -> go+ ParseComplete -> go+ DataRow vals -> case stmtDecode stmt vals of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> do+ action val+ go+ CommandComplete _ -> go+ EmptyQueryResponse -> go+ ReadyForQuery status -> writeIORef txRef status+ ErrorResponse err -> throwPgWire (QueryError err)+ NoticeResponse _ -> go+ other -> throwPgWire (ProtocolError ("Unexpected in forEach: " <> BS8.pack (show other)))+ go++------------------------------------------------------------------------+-- Fast decode variants+------------------------------------------------------------------------++-- | Like 'fetchAll' but uses 'FromRowFast' for faster decoding.+-- Throws exceptions on decode errors instead of returning 'Either'.+-- ~10-20% faster on wide rows (10+ columns) due to eliminating+-- per-column 'Either' wrapper allocations.+--+-- @+-- users <- fetchAllFast conn listAllUsers ()+-- @+fetchAllFast :: (FromRowFast r) => Connection -> Statement p r -> p -> IO [r]+fetchAllFast conn stmt params = do+ rows <- fetchRowsRaw conn stmt params+ pure (map fromRowFast rows)+{-# INLINABLE fetchAllFast #-}++-- | Like 'fetchOne' but uses 'FromRowFast' for faster decoding.+fetchOneFast :: (FromRowFast r) => Connection -> Statement p r -> p -> IO (Maybe r)+fetchOneFast conn stmt params = do+ mRow <- fetchFirstRowRaw conn stmt params+ pure (fmap fromRowFast mRow)+{-# INLINABLE fetchOneFast #-}++------------------------------------------------------------------------+-- Commands+------------------------------------------------------------------------++-- | Execute a command (INSERT\/UPDATE\/DELETE) and return the number of+-- rows affected. The statement result type is fixed to @()@ since no+-- rows are decoded.+--+-- @+-- n <- execute conn insertUser ("Alice", Just "alice\@example.com")+-- @+execute :: Connection -> Statement p () -> p -> IO Int64+execute conn stmt params = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encodedParams = stmtEncode stmt params+ bindExec =+ [ Bind "" name binaryFmtVec encodedParams V.empty+ , Execute "" 0+ , Sync+ ]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExec+ else bindExec+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectCommand+ case resp of+ RespCommand tag -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ pure (tagRows tag)+ _ -> throwPgWire (ProtocolError "execute: unexpected response type")++-- | Execute a command with a RETURNING clause. Returns the number of rows+-- affected and the decoded result rows.+--+-- Unlike 'execute' (which discards result rows) or 'fetchAll' (which+-- discards the command tag), this function captures both.+--+-- @+-- (n, ids) <- executeReturning conn insertUsers ("Alice", Just "alice\@example.com")+-- @+executeReturning :: Connection -> Statement p r -> p -> IO (Int64, [r])+executeReturning conn stmt params = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encodedParams = stmtEncode stmt params+ bindExec =+ [ Bind "" name binaryFmtVec encodedParams binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExec+ else bindExec+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectRowsAndCommand+ case resp of+ RespRowsAndCommand rawRows tag -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ decoded <- decodeRows (stmtDecode stmt) rawRows+ pure (tagRows tag, decoded)+ _ -> throwPgWire (ProtocolError "executeReturning: unexpected response type")++-- | Execute a statement once for each parameter set. Returns the total+-- number of rows affected across all executions. This is an alias for+-- 'executeBatch' with a more common name (matching postgresql-simple's API).+--+-- @+-- total <- executeMany conn insertUser [(\"Alice\", Nothing), (\"Bob\", Just \"b\@x.com\")]+-- @+executeMany :: Connection -> Statement p () -> [p] -> IO Int64+executeMany = executeBatch++-- | Execute a statement with RETURNING for each parameter set, pipelined+-- into a single round-trip. Returns the total rows affected and all+-- decoded RETURNING rows concatenated.+--+-- @+-- (total, ids) <- executeReturningMany conn insertUserReturningId+-- [(\"Alice\", Nothing), (\"Bob\", Just \"b\@x.com\")]+-- @+executeReturningMany :: Connection -> Statement p r -> [p] -> IO (Int64, [r])+executeReturningMany _ _ [] = pure (0, [])+executeReturningMany conn stmt paramsList = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encode = stmtEncode stmt+ bindExecs = concatMap (\p ->+ [ Bind "" name binaryFmtVec (encode p) binaryFmtVec+ , Execute "" 0+ ]) paramsList ++ [Sync]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExecs+ else bindExecs+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs (CollectBatch (length paramsList))+ case resp of+ RespBatchRows results -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ allRows <- concat <$> mapM (decodeRows (stmtDecode stmt)) results+ let total = fromIntegral (length allRows)+ pure (total, allRows)+ _ -> throwPgWire (ProtocolError "executeReturningMany: unexpected response type")++------------------------------------------------------------------------+-- Raw (unchecked) queries+------------------------------------------------------------------------++-- | Execute a raw SQL query and return all rows as untyped column vectors.+-- This is the escape hatch for dynamic SQL or queries that don't fit the+-- 'Statement' model. No compile-time type checking is performed.+--+-- Parameters are passed as pre-encoded binary 'ByteString' values (or+-- 'Nothing' for NULL). Parameter OIDs tell Postgres how to interpret them;+-- use OID 0 to let Postgres infer the type.+--+-- @+-- rows <- rawFetchAll conn+-- "SELECT id, name FROM users WHERE role = $1"+-- [25] -- OID 25 = text+-- [Just "admin"] -- $1 value+-- @+rawFetchAll+ :: Connection+ -> ByteString+ -- ^ SQL query text+ -> [Word32]+ -- ^ Parameter OIDs (use 0 for Postgres to infer)+ -> [Maybe ByteString]+ -- ^ Parameter values (binary-encoded, or Nothing for NULL)+ -> IO [Vector (Maybe ByteString)]+rawFetchAll conn sql oids params = do+ (name, needsParse) <- lookupOrAllocRaw conn sql+ let paramVec = V.fromList params+ msgs =+ [Parse name sql (V.fromList oids) | needsParse]+ ++ [ Bind "" name binaryFmtVec paramVec binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectRows+ case resp of+ RespRows rows -> do+ when needsParse $ cacheStmt conn sql name+ pure rows+ _ -> throwPgWire (ProtocolError "rawFetchAll: unexpected response type")++-- | Execute a raw SQL query and return the first row, or 'Nothing' if the+-- query produces no results. Like 'rawFetchAll' but discards all rows after+-- the first.+rawFetchOne+ :: Connection+ -> ByteString+ -- ^ SQL query text+ -> [Word32]+ -- ^ Parameter OIDs (use 0 for Postgres to infer)+ -> [Maybe ByteString]+ -- ^ Parameter values (binary-encoded, or Nothing for NULL)+ -> IO (Maybe (Vector (Maybe ByteString)))+rawFetchOne conn sql oids params = do+ rows <- rawFetchAll conn sql oids params+ case rows of+ [] -> pure Nothing+ (row : _) -> pure (Just row)++-- | Execute a raw SQL command (INSERT\/UPDATE\/DELETE) and return the number+-- of rows affected. The raw counterpart of 'execute', with no compile-time+-- type checking.+rawExecute+ :: Connection+ -> ByteString+ -- ^ SQL command text+ -> [Word32]+ -- ^ Parameter OIDs (use 0 for Postgres to infer)+ -> [Maybe ByteString]+ -- ^ Parameter values (binary-encoded, or Nothing for NULL)+ -> IO Int64+rawExecute conn sql oids params = do+ (name, needsParse) <- lookupOrAllocRaw conn sql+ let paramVec = V.fromList params+ msgs =+ [Parse name sql (V.fromList oids) | needsParse]+ ++ [ Bind "" name binaryFmtVec paramVec V.empty+ , Execute "" 0+ , Sync+ ]+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectCommand+ case resp of+ RespCommand tag -> do+ when needsParse $ cacheStmt conn sql name+ pure (tagRows tag)+ _ -> throwPgWire (ProtocolError "rawExecute: unexpected response type")++-- | Look up or allocate a statement name for raw SQL (same cache as typed).+lookupOrAllocRaw :: Connection -> ByteString -> IO (ByteString, Bool)+lookupOrAllocRaw conn sql = do+ cache <- readIORef (connStmtCache conn)+ case PSQ.lookup sql cache of+ Just (_prio, name) -> do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)+ pure (name, False)+ Nothing -> do+ evictIfNeeded conn cache+ counter <- atomicModifyIORef' (connStmtCounter conn) (\n -> (n + 1, n))+ let name = "s" <> BS8.pack (show counter)+ pure (name, True)++------------------------------------------------------------------------+-- Pipelined batch execution+------------------------------------------------------------------------++-- | Execute a batch of commands using pipeline mode, returning the total+-- number of rows affected. All Bind+Execute pairs share a single Sync,+-- eliminating per-item round-trip overhead.+--+-- For batches larger than 256 items, automatically switches to streaming+-- mode (exclusive wire access, chunked sends) to bound memory usage.+--+-- @+-- total <- executeBatch conn insertUser [("Alice", Nothing), ("Bob", Just "b\@x.com")]+-- @+executeBatch :: Connection -> Statement p () -> [p] -> IO Int64+executeBatch _ _ [] = pure 0+executeBatch conn stmt paramsList = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ case splitAtEnd batchStreamThreshold paramsList of+ -- Small batch: coalesce through async channel+ (small, Nothing) -> do+ let encode = stmtEncode stmt+ bindExecs = concatMap (\p ->+ [ Bind "" name binaryFmtVec (encode p) V.empty+ , Execute "" 0+ ]) small ++ [Sync]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExecs+ else bindExecs+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs (CollectBatchCommand (length small))+ case resp of+ RespBatchCommand total -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ pure total+ _ -> throwPgWire (ProtocolError "executeBatch: unexpected response type")++ -- Large batch: stream in chunks via exclusive mode+ (_, Just _) -> submitExclusive (connAsync conn) $ \wc txRef -> do+ when needsParse $ do+ sendFrontendMsgs wc [Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)), Flush]+ waitParseWire wc+ streamBatchChunks wc name stmt paramsList+ sendFrontendMsg wc Sync+ total <- collectBatchCmdWire wc txRef (length paramsList)+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ pure total++-- | Fetch zero or one row for each parameter set, pipelined into a single+-- round trip. Returns a list of results parallel to the input list, where+-- each element is 'Nothing' if that particular query returned no rows.+--+-- @+-- results <- fetchBatchOne conn findById [1, 2, 3]+-- -- results :: [Maybe User]+-- @+fetchBatchOne :: Connection -> Statement p r -> [p] -> IO [Maybe r]+fetchBatchOne _ _ [] = pure []+fetchBatchOne conn stmt paramsList = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encode = stmtEncode stmt+ bindExecs = concatMap (\p ->+ [ Bind "" name binaryFmtVec (encode p) binaryFmtVec+ , Execute "" 0+ ]) paramsList ++ [Sync]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExecs+ else bindExecs+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs (CollectBatch (length paramsList))+ case resp of+ RespBatchRows results -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ mapM (\case+ [] -> pure Nothing+ (row : _) -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure (Just val)) results+ _ -> throwPgWire (ProtocolError "fetchBatchOne: unexpected response type")++-- | Fetch all rows for each parameter set, pipelined into a single round+-- trip. Returns a list of result lists parallel to the input list.+--+-- @+-- results <- fetchBatchAll conn listPostsByUser [1, 2, 3]+-- -- results :: [[Post]]+-- @+fetchBatchAll :: Connection -> Statement p r -> [p] -> IO [[r]]+fetchBatchAll _ _ [] = pure []+fetchBatchAll conn stmt paramsList = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encode = stmtEncode stmt+ bindExecs = concatMap (\p ->+ [ Bind "" name binaryFmtVec (encode p) binaryFmtVec+ , Execute "" 0+ ]) paramsList ++ [Sync]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExecs+ else bindExecs+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs (CollectBatch (length paramsList))+ case resp of+ RespBatchRows results -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ mapM (decodeRows (stmtDecode stmt)) results+ _ -> throwPgWire (ProtocolError "fetchBatchAll: unexpected response type")++------------------------------------------------------------------------+-- Internal+------------------------------------------------------------------------++-- | Fetch at most the first raw row, discarding the rest at the wire level.+fetchFirstRowRaw :: Connection -> Statement p r -> p -> IO (Maybe (Vector (Maybe ByteString)))+fetchFirstRowRaw conn stmt params = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encodedParams = stmtEncode stmt params+ bindExec =+ [ Bind "" name binaryFmtVec encodedParams binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExec+ else bindExec+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectFirstRow+ case resp of+ RespFirstRow mRow -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ pure mRow+ _ -> throwPgWire (ProtocolError "fetchFirstRowRaw: unexpected response type")++-- | Fetch raw rows, coalescing Parse+Bind+Execute for cache misses.+fetchRowsRaw :: Connection -> Statement p r -> p -> IO [Vector (Maybe ByteString)]+fetchRowsRaw conn stmt params = do+ (name, needsParse) <- lookupOrAllocStmt conn stmt+ let encodedParams = stmtEncode stmt params+ bindExec =+ [ Bind "" name binaryFmtVec encodedParams binaryFmtVec+ , Execute "" 0+ , Sync+ ]+ msgs = if needsParse+ then Parse name (stmtSQL stmt) (V.map Oid.unOid (stmtParamOids stmt)) : bindExec+ else bindExec+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs CollectRows+ case resp of+ RespRows rows -> do+ when needsParse $ cacheStmt conn (stmtSQL stmt) name+ pure rows+ _ -> throwPgWire (ProtocolError "fetchRows: unexpected response type")++-- | Check the statement cache. On miss, allocate a name and evict if needed.+-- When prepared statements are disabled (PgBouncer compatibility), always+-- returns the unnamed statement ("", True) so queries are parsed each time.+lookupOrAllocStmt :: Connection -> Statement p r -> IO (ByteString, Bool)+lookupOrAllocStmt conn stmt+ | not (connPreparedStatements conn) =+ -- Unprepared mode: always use unnamed statement, always parse+ pure ("", True)+ | otherwise = do+ cache <- readIORef (connStmtCache conn)+ let sql = stmtSQL stmt+ case PSQ.lookup sql cache of+ Just (_prio, name) -> do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)+ pure (name, False)+ Nothing -> do+ evictIfNeeded conn cache+ counter <- atomicModifyIORef' (connStmtCounter conn) (\n -> (n + 1, n))+ let name = "s" <> BS8.pack (show counter)+ pure (name, True)++-- | Cache a statement name after successful Parse.+-- No-op when prepared statements are disabled (unnamed statements are never cached).+cacheStmt :: Connection -> ByteString -> ByteString -> IO ()+cacheStmt conn sql name+ | not (connPreparedStatements conn) = pure ()+ | otherwise = do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)++-- | Decode a list of raw row vectors into typed values.+decodeRows :: (Vector (Maybe ByteString) -> Either String r) -> [Vector (Maybe ByteString)] -> IO [r]+decodeRows decode = mapM $ \row -> case decode row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> pure val++-- | Batch size threshold for switching to streaming mode.+batchStreamThreshold :: Int+batchStreamThreshold = 256++-- | Ensure a statement is prepared on this connection. Returns the statement name.+-- Uses Parse+Flush (no ReadyForQuery overhead). Prefer the coalesced path in+-- 'fetchRowsRaw' / 'execute' for better latency; this function exists for+-- callers that need the statement name before building messages (Pipeline, Fold).+ensurePrepared :: Connection -> Statement p r -> IO ByteString+ensurePrepared conn stmt = do+ cache <- readIORef (connStmtCache conn)+ let sql = stmtSQL stmt+ case PSQ.lookup sql cache of+ Just (_prio, name) -> do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)+ pure name+ Nothing -> do+ evictIfNeeded conn cache+ counter <- atomicModifyIORef' (connStmtCounter conn) (\n -> (n + 1, n))+ let name = "s" <> BS8.pack (show counter)+ oids = V.map Oid.unOid (stmtParamOids stmt)+ resp <- submitRequest (connAsync conn) $ ReqPrepare (Parse name sql oids)+ case resp of+ RespParsed -> do+ tick <- nextTick conn+ modifyIORef' (connStmtCache conn) (PSQ.insert sql tick name)+ pure name+ _ -> throwPgWire (ProtocolError "ensurePrepared: unexpected response type")++-- | Bump the monotonic tick counter and return the new value.+nextTick :: Connection -> IO Word64+nextTick conn = atomicModifyIORef' (connStmtTick conn) (\n -> (n + 1, n + 1))++-- | If the cache has reached its limit, evict the least recently used+-- prepared statement (the entry with the lowest tick priority).+evictIfNeeded :: Connection -> HashPSQ ByteString Word64 ByteString -> IO ()+evictIfNeeded conn cache+ | PSQ.size cache < connMaxPreparedStatements conn = pure ()+ | otherwise = case PSQ.findMin cache of+ Nothing -> pure ()+ Just (oldSql, _prio, oldName) -> do+ resp <- submitRequest (connAsync conn) $ ReqClose (Close DescribeStatement oldName)+ case resp of+ RespClosed -> pure ()+ _ -> pure () -- best effort+ modifyIORef' (connStmtCache conn) (PSQ.delete oldSql)++------------------------------------------------------------------------+-- Streaming batch helpers (exclusive mode, direct wire access)+------------------------------------------------------------------------++-- | Stream Bind+Execute pairs in chunks of 'batchStreamThreshold'.+streamBatchChunks :: WireConn -> ByteString -> Statement p () -> [p] -> IO ()+streamBatchChunks _ _ _ [] = pure ()+streamBatchChunks wc name stmt paramsList = do+ let (chunk, rest) = splitAt batchStreamThreshold paramsList+ msgs = concatMap (\params ->+ let ep = stmtEncode stmt params+ in [ Bind "" name binaryFmtVec ep V.empty+ , Execute "" 0+ ]) chunk+ sendFrontendMsgs wc msgs+ streamBatchChunks wc name stmt rest++-- | Wait for ParseComplete on the wire (used after Flush in exclusive mode).+waitParseWire :: WireConn -> IO ()+waitParseWire wc = do+ msg <- recvBackendMsg wc+ case msg of+ ParseComplete -> pure ()+ NoticeResponse _ -> waitParseWire wc+ ErrorResponse err -> throwPgWire (QueryError err)+ other -> throwPgWire (ProtocolError ("Expected ParseComplete, got: " <> BS8.pack (show other)))++-- | Collect batch command results on the wire (used in streaming mode).+collectBatchCmdWire :: WireConn -> IORef TxStatus -> Int -> IO Int64+collectBatchCmdWire wc txRef = go 0+ where+ go !total 0 = do+ waitReadyWire wc txRef+ pure total+ go !total !n = do+ msg <- recvBackendMsg wc+ case msg of+ BindComplete -> go total n+ CommandComplete tag -> go (total + tagRows tag) (n - 1)+ ErrorResponse err -> do+ waitReadyWire wc txRef+ throwPgWire (QueryError err)+ NoticeResponse _ -> go total n+ other -> throwPgWire (ProtocolError ("Unexpected in batch: " <> BS8.pack (show other)))++-- | Wait for ReadyForQuery on the wire.+waitReadyWire :: WireConn -> IORef TxStatus -> IO ()+waitReadyWire wc txRef = do+ msg <- recvBackendMsg wc+ case msg of+ ReadyForQuery status -> writeIORef txRef status+ _ -> waitReadyWire wc txRef++-- | Split a list, returning (prefix, Nothing) if length <= n,+-- or (full list, Just ()) if length > n. Avoids forcing the full spine+-- for the common small-batch case.+splitAtEnd :: Int -> [a] -> ([a], Maybe ())+splitAtEnd _ [] = ([], Nothing)+splitAtEnd 0 xs = (xs, Just ())+splitAtEnd n (x : xs) = case splitAtEnd (n - 1) xs of+ (rest, flag) -> (x : rest, flag)++tagRows :: CommandTag -> Int64+tagRows (InsertTag n) = n+tagRows (UpdateTag n) = n+tagRows (DeleteTag n) = n+tagRows (SelectTag n) = n+tagRows (OtherTag _) = 0
+ src/Valiant/Fold.hs view
@@ -0,0 +1,99 @@+{-# OPTIONS_GHC -fno-full-laziness #-}++-- | Constant-memory row processing via strict folds.+--+-- Process large result sets without buffering all rows in memory+-- and without requiring a transaction or cursor.+--+-- Folds require exclusive access to the connection's wire so that rows+-- can be processed one-at-a-time directly from the socket.+--+-- @+-- -- Count and sum in one pass, constant memory:+-- (count, total) <- executeWithFold conn stmt params $+-- RowFold (0, 0) (\\(!c, !t) (_, _, score) -> (c + 1, t + score))+-- @+module Valiant.Fold+ ( RowFold (..)+ , executeWithFold+ ) where++import Control.Exception (SomeException, catch, throwIO)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Vector (Vector)+import Data.Vector qualified as V+import PgWire.Async (submitExclusive)+import PgWire.Connection (Connection (..))+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend+import PgWire.Protocol.Frontend+import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsgs)+import Valiant.Execute (ensurePrepared)+import Valiant.Statement (Statement (..))++-- | A strict left fold over result rows. Processes rows in constant+-- memory as they arrive from the wire — no list, no buffering.+data RowFold a b = RowFold+ { foldInit :: !b+ -- ^ Initial accumulator value.+ , foldStep :: !(b -> a -> b)+ -- ^ Strict step function. Apply bang patterns to accumulator components.+ }++-- | Execute a statement and fold over the results in constant memory.+--+-- Each row is decoded and fed to the fold as it arrives from the wire.+-- The entire result set is never held in memory.+--+-- @+-- total <- executeWithFold conn countStmt () $+-- RowFold 0 (\\acc (Only n) -> acc + n)+-- @+executeWithFold :: Connection -> Statement p r -> p -> RowFold r b -> IO b+executeWithFold conn stmt params (RowFold z0 step) = do+ stmtName <- ensurePrepared conn stmt+ let encodedParams = stmtEncode stmt params+ submitExclusive (connAsync conn) $ \wc txRef -> do+ sendFrontendMsgs wc+ [ Bind "" stmtName (V.singleton BinaryFormat) encodedParams (V.singleton BinaryFormat)+ , Execute "" 0+ , Sync+ ]+ collectFold wc txRef (stmtDecode stmt) z0 step+ `catch` \(e :: SomeException) -> do+ -- Drain remaining messages up to ReadyForQuery so the connection+ -- is left in a clean state after a fold exception (e.g., decode+ -- error or user step function throwing).+ drainUntilReady wc txRef+ `catch` \(_ :: SomeException) -> pure ()+ throwIO e++collectFold :: WireConn -> IORef TxStatus -> (Vector (Maybe ByteString) -> Either String r) -> b -> (b -> r -> b) -> IO b+collectFold wc txRef decode = go+ where+ go !acc step' = do+ msg <- recvBackendMsg wc+ case msg of+ BindComplete -> go acc step'+ DataRow vals -> case decode vals of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> go (step' acc val) step'+ CommandComplete _ -> go acc step'+ EmptyQueryResponse -> go acc step'+ ReadyForQuery status -> do+ writeIORef txRef status+ pure acc+ ErrorResponse err -> throwPgWire (QueryError err)+ NoticeResponse _ -> go acc step'+ other -> throwPgWire (ProtocolError ("Unexpected in fold: " <> BS8.pack (show other)))++-- | Drain messages until ReadyForQuery, updating transaction status.+-- Used for cleanup after fold exceptions.+drainUntilReady :: WireConn -> IORef TxStatus -> IO ()+drainUntilReady wc txRef = do+ msg <- recvBackendMsg wc+ case msg of+ ReadyForQuery status -> writeIORef txRef status+ _ -> drainUntilReady wc txRef
+ src/Valiant/FromRow.hs view
@@ -0,0 +1,210 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Decode result rows into Haskell types.+module Valiant.FromRow+ ( FromRow (..)+ , DecodeColumn (..)+ , GFromRow (..)+ ) where++import Data.ByteString (ByteString)+import Data.Kind (Type)+import Data.Vector (Vector)+import Data.Vector qualified as V+import GHC.Generics+import Valiant.Binary.Decode ()+import PgWire.Binary.Types (PgDecode (..))++-- | Closed type family: is this type @Maybe a@?+type family Nullable (a :: Type) :: Bool where+ Nullable (Maybe _) = 'True+ Nullable _ = 'False++-- | Decode a single column value, handling NULL appropriately.+--+-- For non-Maybe types, NULL produces an error.+-- For @Maybe a@, NULL produces @Nothing@.+--+-- This uses a closed type family ('Nullable') to dispatch without+-- overlapping instances.+class DecodeColumn a where+ decodeColumn :: Vector (Maybe ByteString) -> Int -> Either String a++-- | Implementation dispatcher — selects nullable vs non-nullable decoding+-- based on the closed type family 'Nullable'.+instance (Nullable a ~ flag, DecodeColumnImpl flag a) => DecodeColumn a where+ decodeColumn = decodeColumnImpl @flag+ {-# INLINE decodeColumn #-}++-- | Internal class parameterized by the 'Nullable' flag.+-- Not exported — users only see 'DecodeColumn'.+class DecodeColumnImpl (flag :: Bool) a where+ decodeColumnImpl :: Vector (Maybe ByteString) -> Int -> Either String a++-- | Non-nullable: NULL is an error.+instance (PgDecode a) => DecodeColumnImpl 'False a where+ decodeColumnImpl row idx+ | idx >= V.length row = Left $ colOutOfRange idx (V.length row)+ | otherwise = case V.unsafeIndex row idx of+ Nothing -> Left $ "Column " <> show idx <> " is NULL but expected a non-nullable value"+ Just bs -> pgDecode bs+ {-# INLINE decodeColumnImpl #-}++-- | Nullable: NULL becomes @Nothing@, non-NULL becomes @Just a@.+instance (PgDecode a) => DecodeColumnImpl 'True (Maybe a) where+ decodeColumnImpl row idx+ | idx >= V.length row = Left $ colOutOfRange idx (V.length row)+ | otherwise = case V.unsafeIndex row idx of+ Nothing -> Right Nothing+ Just bs -> Just <$> pgDecode bs+ {-# INLINE decodeColumnImpl #-}++colOutOfRange :: Int -> Int -> String+colOutOfRange idx len =+ "Column index " <> show idx <> " out of range (row has " <> show len <> " columns)"+{-# INLINE colOutOfRange #-}++-- | Decode a result row into a Haskell value.+--+-- For custom record types, derive via @Generic@:+--+-- @+-- data User = User+-- { userId :: Int32+-- , userName :: Text+-- , userEmail :: Maybe Text+-- } deriving stock (Generic)+-- deriving anyclass (FromRow)+-- @+--+-- Fields are decoded positionally (column 0 → first field, etc.).+class FromRow a where+ fromRow :: Vector (Maybe ByteString) -> Either String a+ default fromRow :: (Generic a, GFromRow (Rep a)) => Vector (Maybe ByteString) -> Either String a+ fromRow row = to <$> gFromRow row 0++-- | Generic helper for positional row decoding. Internal — not exported.+class GFromRow f where+ gFromRow :: Vector (Maybe ByteString) -> Int -> Either String (f p)+ gFieldCount :: proxy f -> Int++-- Datatype metadata — delegate+instance (GFromRow f) => GFromRow (M1 D c f) where+ gFromRow row idx = M1 <$> gFromRow row idx+ gFieldCount _ = gFieldCount (undefined :: proxy f)+ {-# INLINE gFromRow #-}+ {-# INLINE gFieldCount #-}++-- Constructor metadata — delegate+instance (GFromRow f) => GFromRow (M1 C c f) where+ gFromRow row idx = M1 <$> gFromRow row idx+ gFieldCount _ = gFieldCount (undefined :: proxy f)+ {-# INLINE gFromRow #-}+ {-# INLINE gFieldCount #-}++-- Selector metadata — delegate+instance (GFromRow f) => GFromRow (M1 S c f) where+ gFromRow row idx = M1 <$> gFromRow row idx+ gFieldCount _ = gFieldCount (undefined :: proxy f)+ {-# INLINE gFromRow #-}+ {-# INLINE gFieldCount #-}++-- Leaf field — decode one column at the current index+instance (DecodeColumn a) => GFromRow (K1 R a) where+ gFromRow row idx = K1 <$> decodeColumn row idx+ gFieldCount _ = 1+ {-# INLINE gFromRow #-}+ {-# INLINE gFieldCount #-}++-- Product — decode left fields, then right fields+instance (GFromRow f, GFromRow g) => GFromRow (f :*: g) where+ gFromRow row idx = do+ l <- gFromRow row idx+ r <- gFromRow row (idx + gFieldCount (undefined :: proxy f))+ Right (l :*: r)+ gFieldCount _ = gFieldCount (undefined :: proxy f) + gFieldCount (undefined :: proxy g)+ {-# INLINE gFromRow #-}+ {-# INLINE gFieldCount #-}++-- Unit — no fields+instance GFromRow U1 where+ gFromRow _ _ = Right U1+ gFieldCount _ = 0+ {-# INLINE gFromRow #-}+ {-# INLINE gFieldCount #-}++-- Unit instance for commands (INSERT/UPDATE/DELETE) that return no rows+instance {-# OVERLAPPING #-} FromRow () where+ fromRow _ = Right ()++-- Single column+instance (DecodeColumn a) => FromRow a where+ fromRow row = decodeColumn row 0++-- Single-value wrapper for scalar queries+instance {-# OVERLAPPING #-} (DecodeColumn a) => FromRow (Only a) where+ fromRow row = Only <$> decodeColumn row 0++-- | Wrapper for single-column results.+newtype Only a = Only {fromOnly :: a}+ deriving stock (Show, Eq)++-- Tuple instances ---------------------------------------------------------+-- These need OVERLAPPING because they also match the single-column+-- FromRow a instance above. This is the only place we use overlap,+-- and it's benign: tuples are always more specific than a bare type variable.++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b) => FromRow (a, b) where+ fromRow row = (,) <$> decodeColumn row 0 <*> decodeColumn row 1++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c) => FromRow (a, b, c) where+ fromRow row = (,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d) => FromRow (a, b, c, d) where+ fromRow row = (,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e) => FromRow (a, b, c, d, e) where+ fromRow row = (,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f) => FromRow (a, b, c, d, e, f) where+ fromRow row = (,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g) => FromRow (a, b, c, d, e, f, g) where+ fromRow row = (,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h) => FromRow (a, b, c, d, e, f, g, h) where+ fromRow row = (,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i) => FromRow (a, b, c, d, e, f, g, h, i) where+ fromRow row = (,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j) => FromRow (a, b, c, d, e, f, g, h, i, j) where+ fromRow row = (,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j, DecodeColumn k) => FromRow (a, b, c, d, e, f, g, h, i, j, k) where+ fromRow row = (,,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9 <*> decodeColumn row 10++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j, DecodeColumn k, DecodeColumn l) => FromRow (a, b, c, d, e, f, g, h, i, j, k, l) where+ fromRow row = (,,,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9 <*> decodeColumn row 10 <*> decodeColumn row 11++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j, DecodeColumn k, DecodeColumn l, DecodeColumn m) => FromRow (a, b, c, d, e, f, g, h, i, j, k, l, m) where+ fromRow row = (,,,,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9 <*> decodeColumn row 10 <*> decodeColumn row 11 <*> decodeColumn row 12++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j, DecodeColumn k, DecodeColumn l, DecodeColumn m, DecodeColumn n) => FromRow (a, b, c, d, e, f, g, h, i, j, k, l, m, n) where+ fromRow row = (,,,,,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9 <*> decodeColumn row 10 <*> decodeColumn row 11 <*> decodeColumn row 12 <*> decodeColumn row 13++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j, DecodeColumn k, DecodeColumn l, DecodeColumn m, DecodeColumn n, DecodeColumn o) => FromRow (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) where+ fromRow row = (,,,,,,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9 <*> decodeColumn row 10 <*> decodeColumn row 11 <*> decodeColumn row 12 <*> decodeColumn row 13 <*> decodeColumn row 14++instance {-# OVERLAPPING #-} (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e, DecodeColumn f, DecodeColumn g, DecodeColumn h, DecodeColumn i, DecodeColumn j, DecodeColumn k, DecodeColumn l, DecodeColumn m, DecodeColumn n, DecodeColumn o, DecodeColumn p) => FromRow (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) where+ fromRow row = (,,,,,,,,,,,,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4 <*> decodeColumn row 5 <*> decodeColumn row 6 <*> decodeColumn row 7 <*> decodeColumn row 8 <*> decodeColumn row 9 <*> decodeColumn row 10 <*> decodeColumn row 11 <*> decodeColumn row 12 <*> decodeColumn row 13 <*> decodeColumn row 14 <*> decodeColumn row 15++-- PgDecode (Maybe a) for the single-column FromRow a path+instance (PgDecode a) => PgDecode (Maybe a) where+ pgDecode bs = Just <$> pgDecode bs+ {-# INLINE pgDecode #-}
+ src/Valiant/FromRowFast.hs view
@@ -0,0 +1,174 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Fast row decoding that throws exceptions instead of returning Either.+--+-- 'FromRowFast' eliminates the per-column @Right@ wrapper allocation+-- that 'FromRow' requires. On wide rows (10+ columns), this can be+-- 10-20% faster because:+--+-- 1. No @Either String a@ allocation per column decode+-- 2. No @<*>@ chain threading the @Either@ through Applicative+-- 3. Direct constructor application without wrapping+--+-- The trade-off: decode errors become imprecise exceptions instead of+-- structured @Left@ values. Since decode errors indicate a bug (schema+-- mismatch), not a runtime condition, this is acceptable for most code.+--+-- @+-- -- Use fetchAllFast instead of fetchAll for hot paths:+-- users <- fetchAllFast conn listAllUsers ()+-- @+module Valiant.FromRowFast+ ( FromRowFast (..)+ , DecodeColumnFast (..)+ ) where++import Data.ByteString (ByteString)+import Data.Kind (Type)+import Data.Vector (Vector)+import Data.Vector qualified as V+import GHC.Generics+import Valiant.Binary.Decode ()+import PgWire.Binary.Types (PgDecode (..))++-- | Closed type family: is this type @Maybe a@?+type family Nullable (a :: Type) :: Bool where+ Nullable (Maybe _) = 'True+ Nullable _ = 'False++-- | Decode a single column, throwing on errors instead of returning Either.+-- No bounds checking — the DataRow parser already validated column count.+class DecodeColumnFast a where+ decodeColumnFast :: Vector (Maybe ByteString) -> Int -> a++instance (Nullable a ~ flag, DecodeColumnFastImpl flag a) => DecodeColumnFast a where+ decodeColumnFast = decodeColumnFastImpl @flag+ {-# INLINE decodeColumnFast #-}++class DecodeColumnFastImpl (flag :: Bool) a where+ decodeColumnFastImpl :: Vector (Maybe ByteString) -> Int -> a++-- Non-nullable: NULL throws, decode errors throw.+instance (PgDecode a) => DecodeColumnFastImpl 'False a where+ decodeColumnFastImpl row idx =+ case V.unsafeIndex row idx of+ Nothing -> error $ "Column " <> show idx <> " is NULL"+ Just bs -> case pgDecode bs of+ Left err -> error err+ Right !val -> val+ {-# INLINE decodeColumnFastImpl #-}++-- Nullable: NULL → Nothing, decode errors throw.+instance (PgDecode a) => DecodeColumnFastImpl 'True (Maybe a) where+ decodeColumnFastImpl row idx =+ case V.unsafeIndex row idx of+ Nothing -> Nothing+ Just bs -> case pgDecode bs of+ Left err -> error err+ Right !val -> Just val+ {-# INLINE decodeColumnFastImpl #-}++-- | Fast row decoder. Throws on errors instead of returning Either.+--+-- Derive via @Generic@:+--+-- @+-- data User = User+-- { userId :: Int32, userName :: Text, userEmail :: Maybe Text }+-- deriving stock (Generic)+-- deriving anyclass (FromRowFast)+-- @+class FromRowFast a where+ fromRowFast :: Vector (Maybe ByteString) -> a+ default fromRowFast :: (Generic a, GFromRowFast (Rep a)) => Vector (Maybe ByteString) -> a+ fromRowFast row = to (gFromRowFast row 0)++-- Generic machinery (same structure as FromRow but without Either)+class GFromRowFast f where+ gFromRowFast :: Vector (Maybe ByteString) -> Int -> f p+ gFieldCountFast :: proxy f -> Int++instance (GFromRowFast f) => GFromRowFast (M1 D c f) where+ gFromRowFast row idx = M1 (gFromRowFast row idx)+ gFieldCountFast _ = gFieldCountFast (undefined :: proxy f)+ {-# INLINE gFromRowFast #-}+ {-# INLINE gFieldCountFast #-}++instance (GFromRowFast f) => GFromRowFast (M1 C c f) where+ gFromRowFast row idx = M1 (gFromRowFast row idx)+ gFieldCountFast _ = gFieldCountFast (undefined :: proxy f)+ {-# INLINE gFromRowFast #-}+ {-# INLINE gFieldCountFast #-}++instance (GFromRowFast f) => GFromRowFast (M1 S c f) where+ gFromRowFast row idx = M1 (gFromRowFast row idx)+ gFieldCountFast _ = gFieldCountFast (undefined :: proxy f)+ {-# INLINE gFromRowFast #-}+ {-# INLINE gFieldCountFast #-}++instance (DecodeColumnFast a) => GFromRowFast (K1 R a) where+ gFromRowFast row idx = K1 (decodeColumnFast row idx)+ gFieldCountFast _ = 1+ {-# INLINE gFromRowFast #-}+ {-# INLINE gFieldCountFast #-}++instance (GFromRowFast f, GFromRowFast g) => GFromRowFast (f :*: g) where+ gFromRowFast row idx =+ let !l = gFromRowFast row idx+ !r = gFromRowFast row (idx + gFieldCountFast (undefined :: proxy f))+ in l :*: r+ gFieldCountFast _ = gFieldCountFast (undefined :: proxy f) + gFieldCountFast (undefined :: proxy g)+ {-# INLINE gFromRowFast #-}+ {-# INLINE gFieldCountFast #-}++instance GFromRowFast U1 where+ gFromRowFast _ _ = U1+ gFieldCountFast _ = 0+ {-# INLINE gFromRowFast #-}+ {-# INLINE gFieldCountFast #-}++-- ── Instances ──────────────────────────────────────────────────────++instance {-# OVERLAPPING #-} FromRowFast () where+ fromRowFast _ = ()++instance (DecodeColumnFast a) => FromRowFast a where+ fromRowFast row = decodeColumnFast row 0++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b) => FromRowFast (a, b) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c) => FromRowFast (a, b, c) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d) => FromRowFast (a, b, c, d) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e) => FromRowFast (a, b, c, d, e) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f) => FromRowFast (a, b, c, d, e, f) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f, DecodeColumnFast g) => FromRowFast (a, b, c, d, e, f, g) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5, decodeColumnFast row 6)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f, DecodeColumnFast g, DecodeColumnFast h) => FromRowFast (a, b, c, d, e, f, g, h) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5, decodeColumnFast row 6, decodeColumnFast row 7)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f, DecodeColumnFast g, DecodeColumnFast h, DecodeColumnFast i) => FromRowFast (a, b, c, d, e, f, g, h, i) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5, decodeColumnFast row 6, decodeColumnFast row 7, decodeColumnFast row 8)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f, DecodeColumnFast g, DecodeColumnFast h, DecodeColumnFast i, DecodeColumnFast j) => FromRowFast (a, b, c, d, e, f, g, h, i, j) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5, decodeColumnFast row 6, decodeColumnFast row 7, decodeColumnFast row 8, decodeColumnFast row 9)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f, DecodeColumnFast g, DecodeColumnFast h, DecodeColumnFast i, DecodeColumnFast j, DecodeColumnFast k) => FromRowFast (a, b, c, d, e, f, g, h, i, j, k) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5, decodeColumnFast row 6, decodeColumnFast row 7, decodeColumnFast row 8, decodeColumnFast row 9, decodeColumnFast row 10)++instance {-# OVERLAPPING #-} (DecodeColumnFast a, DecodeColumnFast b, DecodeColumnFast c, DecodeColumnFast d, DecodeColumnFast e, DecodeColumnFast f, DecodeColumnFast g, DecodeColumnFast h, DecodeColumnFast i, DecodeColumnFast j, DecodeColumnFast k, DecodeColumnFast l) => FromRowFast (a, b, c, d, e, f, g, h, i, j, k, l) where+ fromRowFast row = (decodeColumnFast row 0, decodeColumnFast row 1, decodeColumnFast row 2, decodeColumnFast row 3, decodeColumnFast row 4, decodeColumnFast row 5, decodeColumnFast row 6, decodeColumnFast row 7, decodeColumnFast row 8, decodeColumnFast row 9, decodeColumnFast row 10, decodeColumnFast row 11)
+ src/Valiant/FromRowStrict.hs view
@@ -0,0 +1,65 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Strict variant of 'FromRow' that validates column count.+--+-- 'FromRow' silently ignores extra columns in the result row. This is+-- convenient but can mask schema evolution bugs. 'FromRowStrict' checks+-- that the number of columns matches exactly, failing with an error if+-- there are extra or missing columns.+--+-- @+-- data User = User { userId :: Int32, userName :: Text }+-- deriving stock (Generic)+-- deriving anyclass (FromRowStrict)+--+-- -- SELECT id, name FROM users → OK+-- -- SELECT id, name, email FROM users → Error: expected 2 columns, got 3+-- @+module Valiant.FromRowStrict+ ( FromRowStrict (..)+ ) where++import Data.ByteString (ByteString)+import Data.Vector (Vector)+import Data.Vector qualified as V+import GHC.Generics+import Valiant.FromRow (DecodeColumn (..), GFromRow (..))++-- | Like 'FromRow' but validates that the row has exactly the expected+-- number of columns. Derive via @Generic@ the same way as 'FromRow'.+class FromRowStrict a where+ fromRowStrict :: Vector (Maybe ByteString) -> Either String a+ default fromRowStrict :: (Generic a, GFromRow (Rep a)) => Vector (Maybe ByteString) -> Either String a+ fromRowStrict row+ | V.length row /= expected =+ Left $ "Column count mismatch: expected " <> show expected+ <> " but got " <> show (V.length row)+ | otherwise = to <$> gFromRow row 0+ where+ expected = gFieldCount (undefined :: proxy (Rep a))++-- Tuple instances with exact column count validation.++instance (DecodeColumn a, DecodeColumn b) => FromRowStrict (a, b) where+ fromRowStrict row+ | V.length row /= 2 = Left $ "Column count mismatch: expected 2 but got " <> show (V.length row)+ | otherwise = (,) <$> decodeColumn row 0 <*> decodeColumn row 1++instance (DecodeColumn a, DecodeColumn b, DecodeColumn c) => FromRowStrict (a, b, c) where+ fromRowStrict row+ | V.length row /= 3 = Left $ "Column count mismatch: expected 3 but got " <> show (V.length row)+ | otherwise = (,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2++instance (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d) => FromRowStrict (a, b, c, d) where+ fromRowStrict row+ | V.length row /= 4 = Left $ "Column count mismatch: expected 4 but got " <> show (V.length row)+ | otherwise = (,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3++instance (DecodeColumn a, DecodeColumn b, DecodeColumn c, DecodeColumn d, DecodeColumn e) => FromRowStrict (a, b, c, d, e) where+ fromRowStrict row+ | V.length row /= 5 = Left $ "Column count mismatch: expected 5 but got " <> show (V.length row)+ | otherwise = (,,,,) <$> decodeColumn row 0 <*> decodeColumn row 1 <*> decodeColumn row 2 <*> decodeColumn row 3 <*> decodeColumn row 4
+ src/Valiant/LargeObject.hs view
@@ -0,0 +1,196 @@+-- | PostgreSQL large object API.+--+-- Large objects provide a streaming interface for binary data too large+-- to fit in a single column value. Operations must occur within a+-- transaction.+--+-- @+-- withTransaction pool $ \\tx -> do+-- oid <- loCreate (txConn tx)+-- fd <- loOpen (txConn tx) oid WriteMode+-- loWrite (txConn tx) fd someBytes+-- loClose (txConn tx) fd+-- @+module Valiant.LargeObject+ ( -- * Types+ LoFd (..)+ , LoMode (..)+ -- * Create / delete+ , loCreate+ , loUnlink+ -- * Open / close+ , loOpen+ , loClose+ , withLargeObject+ -- * Read / write+ , loRead+ , loWrite+ -- * Seek / tell+ , loSeek+ , loTell+ , loTruncate+ -- * Import / export+ , loImport+ , loExport+ ) where++import Control.Exception (bracket)+import Data.ByteString (ByteString)+import Data.ByteString qualified as BS+import Data.ByteString.Char8 qualified as BS8+import Data.Int (Int32, Int64)+import Data.Word (Word32)+import PgWire.Connection (Connection, escapeLiteral, simpleQuery)+import PgWire.Error (PgWireError (..), throwPgWire)++-- | A large object file descriptor, returned by 'loOpen'.+newtype LoFd = LoFd { unLoFd :: Int32 }+ deriving stock (Show, Eq)++-- | Access mode for 'loOpen'.+data LoMode+ = ReadMode+ | WriteMode+ | ReadWriteMode+ deriving stock (Show, Eq)++loModeToInt :: LoMode -> Int32+loModeToInt ReadMode = 0x00040000 -- INV_READ+loModeToInt WriteMode = 0x00020000 -- INV_WRITE+loModeToInt ReadWriteMode = 0x00060000 -- INV_READ | INV_WRITE++-- | Create a new large object with a system-assigned OID.+-- Returns the OID of the new object.+loCreate :: Connection -> IO Word32+loCreate conn = do+ (rows, _) <- simpleQuery conn "SELECT lo_create(0)"+ case rows of+ [[Just oidBs]] -> case BS8.readInt oidBs of+ Just (n, _) -> pure (fromIntegral n)+ Nothing -> throwPgWire (DecodeError "loCreate: could not parse OID")+ _ -> throwPgWire (DecodeError "loCreate: unexpected result")++-- | Delete a large object.+loUnlink :: Connection -> Word32 -> IO ()+loUnlink conn oid = do+ _ <- simpleQuery conn ("SELECT lo_unlink(" <> BS8.pack (show oid) <> ")")+ pure ()++-- | Open a large object for reading and/or writing.+-- Must be called within a transaction.+loOpen :: Connection -> Word32 -> LoMode -> IO LoFd+loOpen conn oid mode = do+ (rows, _) <- simpleQuery conn+ ("SELECT lo_open(" <> BS8.pack (show oid) <> ", " <> BS8.pack (show (loModeToInt mode)) <> ")")+ case rows of+ [[Just fdBs]] -> case BS8.readInt fdBs of+ Just (n, _) -> pure (LoFd (fromIntegral n))+ Nothing -> throwPgWire (DecodeError "loOpen: could not parse fd")+ _ -> throwPgWire (DecodeError "loOpen: unexpected result")++-- | Close a large object file descriptor.+loClose :: Connection -> LoFd -> IO ()+loClose conn (LoFd fd) = do+ _ <- simpleQuery conn ("SELECT lo_close(" <> BS8.pack (show fd) <> ")")+ pure ()++-- | Open a large object, run an action, then close. Guarantees the file+-- descriptor is closed even if the action throws an exception.+-- Must be called within a transaction.+withLargeObject :: Connection -> Word32 -> LoMode -> (LoFd -> IO a) -> IO a+withLargeObject conn oid mode =+ bracket (loOpen conn oid mode) (loClose conn)++-- | Read up to @n@ bytes from a large object.+loRead :: Connection -> LoFd -> Int32 -> IO ByteString+loRead conn (LoFd fd) n = do+ (rows, _) <- simpleQuery conn+ ("SELECT loread(" <> BS8.pack (show fd) <> ", " <> BS8.pack (show n) <> ")")+ case rows of+ [[Just bs]] -> pure (unescapeBytea bs)+ _ -> throwPgWire (DecodeError "loRead: unexpected result")++-- | Write bytes to a large object. Returns number of bytes written.+loWrite :: Connection -> LoFd -> ByteString -> IO Int32+loWrite conn (LoFd fd) bs = do+ let escaped = escapeBytea bs+ (rows, _) <- simpleQuery conn+ ("SELECT lowrite(" <> BS8.pack (show fd) <> ", " <> escaped <> ")")+ case rows of+ [[Just nBs]] -> case BS8.readInt nBs of+ Just (n, _) -> pure (fromIntegral n)+ Nothing -> throwPgWire (DecodeError "loWrite: could not parse result")+ _ -> throwPgWire (DecodeError "loWrite: unexpected result")++-- | Seek to a position in a large object. Returns the new position.+loSeek :: Connection -> LoFd -> Int64 -> Int32 -> IO Int64+loSeek conn (LoFd fd) offset whence = do+ (rows, _) <- simpleQuery conn+ ("SELECT lo_lseek64(" <> BS8.pack (show fd) <> ", " <> BS8.pack (show offset) <> ", " <> BS8.pack (show whence) <> ")")+ case rows of+ [[Just posBs]] -> case BS8.readInteger posBs of+ Just (n, _) -> pure (fromIntegral n)+ Nothing -> throwPgWire (DecodeError "loSeek: could not parse position")+ _ -> throwPgWire (DecodeError "loSeek: unexpected result")++-- | Get the current position in a large object.+loTell :: Connection -> LoFd -> IO Int64+loTell conn (LoFd fd) = do+ (rows, _) <- simpleQuery conn ("SELECT lo_tell64(" <> BS8.pack (show fd) <> ")")+ case rows of+ [[Just posBs]] -> case BS8.readInteger posBs of+ Just (n, _) -> pure (fromIntegral n)+ Nothing -> throwPgWire (DecodeError "loTell: could not parse position")+ _ -> throwPgWire (DecodeError "loTell: unexpected result")++-- | Truncate a large object to the given length.+loTruncate :: Connection -> LoFd -> Int64 -> IO ()+loTruncate conn (LoFd fd) len = do+ _ <- simpleQuery conn+ ("SELECT lo_truncate64(" <> BS8.pack (show fd) <> ", " <> BS8.pack (show len) <> ")")+ pure ()++-- | Import a file from the server filesystem as a large object.+-- Returns the OID of the new object.+loImport :: Connection -> ByteString -> IO Word32+loImport conn path = do+ (rows, _) <- simpleQuery conn ("SELECT lo_import(" <> escapeLiteral conn path <> ")")+ case rows of+ [[Just oidBs]] -> case BS8.readInt oidBs of+ Just (n, _) -> pure (fromIntegral n)+ Nothing -> throwPgWire (DecodeError "loImport: could not parse OID")+ _ -> throwPgWire (DecodeError "loImport: unexpected result")++-- | Export a large object to a file on the server filesystem.+loExport :: Connection -> Word32 -> ByteString -> IO ()+loExport conn oid path = do+ _ <- simpleQuery conn+ ("SELECT lo_export(" <> BS8.pack (show oid) <> ", " <> escapeLiteral conn path <> ")")+ pure ()++-- Simple bytea escape/unescape for the text protocol used by simpleQuery.+-- In the simple query protocol, bytea is returned as hex or escape format.+escapeBytea :: ByteString -> ByteString+escapeBytea bs = "'\\x" <> BS8.pack (concatMap toHex (BS.unpack bs)) <> "'"+ where+ toHex w = [hexDigit (w `div` 16), hexDigit (w `mod` 16)]+ hexDigit n+ | n < 10 = toEnum (fromEnum '0' + fromIntegral n)+ | otherwise = toEnum (fromEnum 'a' + fromIntegral n - 10)++-- Unescape hex-encoded bytea from simple query response.+unescapeBytea :: ByteString -> ByteString+unescapeBytea bs+ | "\\x" `BS.isPrefixOf` bs = decodeHex (BS.drop 2 bs)+ | otherwise = bs -- already raw or escape format+ where+ decodeHex hex = BS.pack (go (BS.unpack hex))+ go [] = []+ go [_] = [] -- odd number of chars, ignore trailing+ go (h : l : rest) = fromHex h l : go rest+ fromHex h l = fromIntegral (hexVal h * 16 + hexVal l)+ hexVal c+ | c >= 48 && c <= 57 = c - 48 -- '0'-'9'+ | c >= 97 && c <= 102 = c - 87 -- 'a'-'f'+ | c >= 65 && c <= 70 = c - 55 -- 'A'-'F'+ | otherwise = 0
+ src/Valiant/Logging.hs view
@@ -0,0 +1,156 @@+-- | Logging hooks for query timing and connection events.+module Valiant.Logging+ ( LogEvent (..)+ , LogLevel (..)+ , Logger+ , nullLogger+ , stderrLogger+ , withQueryLogging+ , poolLoggerFromLogger+ ) where++import Control.Exception (SomeException, displayException, throwIO, try)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.Time (NominalDiffTime, diffUTCTime, getCurrentTime)+import PgWire.Pool.Config (PoolLogger)+import System.IO (hPutStrLn, stderr)++-- | Log severity levels.+data LogLevel = Debug | Info | Warn | Error+ deriving stock (Show, Eq, Ord)++-- | Events that can be logged.+data LogEvent+ = QueryStart ByteString+ | QueryComplete ByteString NominalDiffTime Int+ -- ^ SQL, duration, rows returned+ | QueryError ByteString NominalDiffTime ByteString+ -- ^ SQL, duration, error message+ | ConnectionOpened ByteString+ -- ^ Connection string (masked)+ | ConnectionClosed+ | PoolAcquire NominalDiffTime+ -- ^ Time waited for a connection+ | PoolRelease+ | PoolTimeout+ | PoolCreatedConn+ -- ^ A new connection was created in the pool+ | PoolDestroyedConn ByteString+ -- ^ A connection was destroyed. Includes reason.+ | PoolRecycled+ -- ^ A connection was recycled (health check passed)+ | PoolReaperSwept Int+ -- ^ Reaper swept N idle/expired connections+ | PoolWarmerCreated Int+ -- ^ Warmer created N connections to maintain min-idle+ | PoolResized Int Int+ -- ^ Pool resized from old size to new size+ deriving stock (Show)++-- | A logging callback.+type Logger = LogLevel -> LogEvent -> IO ()++-- | Logger that discards all events.+nullLogger :: Logger+nullLogger _ _ = pure ()++-- | Simple logger that writes to stderr.+stderrLogger :: LogLevel -> Logger+stderrLogger minLevel level event+ | level >= minLevel = hPutStrLn stderr (formatEvent level event)+ | otherwise = pure ()++-- | Bracket a query with timing and logging. Emits 'QueryStart',+-- then either 'QueryComplete' on success or 'QueryError' on+-- exception. Exceptions are re-raised after logging.+withQueryLogging :: Logger -> ByteString -> IO a -> IO a+withQueryLogging logger sql action = do+ logger Debug (QueryStart sql)+ start <- getCurrentTime+ outcome <- try action+ end <- getCurrentTime+ let elapsed = diffUTCTime end start+ case outcome of+ Right result -> do+ logger Debug (QueryComplete sql elapsed 0)+ pure result+ Left (e :: SomeException) -> do+ logger Warn (QueryError sql elapsed (BS8.pack (displayException e)))+ throwIO e++-- | Create a 'PoolLogger' that forwards pool events to a 'Logger' as+-- structured 'LogEvent's. Pass this as the @poolLogger@ field when+-- constructing a 'PgWire.Pool.Config.PoolConfig'.+--+-- @+-- let cfg = defaultPoolConfig+-- { poolLogger = poolLoggerFromLogger (stderrLogger Info)+-- }+-- @+poolLoggerFromLogger :: Logger -> PoolLogger+poolLoggerFromLogger logger level msg = logger valiantLevel event+ where+ valiantLevel = case level of+ "debug" -> Debug+ "info" -> Info+ "warn" -> Warn+ _ -> Info+ event = parsePoolEvent msg++-- | Parse the pool's log message into a structured 'LogEvent'.+-- Falls back to a generic 'PoolDestroyedConn' with the raw message+-- if the pattern doesn't match a known event.+parsePoolEvent :: ByteString -> LogEvent+parsePoolEvent msg+ | msg == "created connection" = PoolCreatedConn+ | msg == "recycled connection" = PoolRecycled+ | msg == "acquire timeout" = PoolTimeout+ | "destroyed connection: " `BS8.isPrefixOf` msg =+ PoolDestroyedConn (BS8.drop 22 msg)+ | "reaper swept " `BS8.isPrefixOf` msg =+ case BS8.readInt (BS8.drop 14 msg) of+ Just (n, _) -> PoolReaperSwept n+ Nothing -> PoolDestroyedConn msg+ | "warmer created " `BS8.isPrefixOf` msg =+ case BS8.readInt (BS8.drop 15 msg) of+ Just (n, _) -> PoolWarmerCreated n+ Nothing -> PoolDestroyedConn msg+ | "resized from " `BS8.isPrefixOf` msg =+ case parseResized (BS8.drop 13 msg) of+ Just (old, new') -> PoolResized old new'+ Nothing -> PoolDestroyedConn msg+ | otherwise = PoolDestroyedConn msg+ where+ parseResized bs = do+ (old, rest) <- BS8.readInt bs+ -- rest should be " to N"+ let rest' = BS8.drop 4 rest -- drop " to "+ (new', _) <- BS8.readInt rest'+ pure (old, new')++formatEvent :: LogLevel -> LogEvent -> String+formatEvent level event =+ "[valiant:" <> show level <> "] " <> case event of+ QueryStart sql -> "query start: " <> trunc 80 (BS8.unpack sql)+ QueryComplete sql dur rows ->+ "query complete: " <> trunc 60 (BS8.unpack sql)+ <> " (" <> show dur <> ", " <> show rows <> " rows)"+ QueryError sql dur msg ->+ "query error: " <> trunc 60 (BS8.unpack sql)+ <> " (" <> show dur <> "): " <> BS8.unpack msg+ ConnectionOpened cs -> "connection opened: " <> BS8.unpack cs+ ConnectionClosed -> "connection closed"+ PoolAcquire dur -> "pool acquire (" <> show dur <> ")"+ PoolRelease -> "pool release"+ PoolTimeout -> "pool timeout"+ PoolCreatedConn -> "pool created connection"+ PoolDestroyedConn reason -> "pool destroyed connection: " <> BS8.unpack reason+ PoolRecycled -> "pool recycled connection"+ PoolReaperSwept n -> "pool reaper swept " <> show n <> " connections"+ PoolWarmerCreated n -> "pool warmer created " <> show n <> " connections"+ PoolResized old new' -> "pool resized from " <> show old <> " to " <> show new'+ where+ trunc n s+ | length s <= n = s+ | otherwise = take n s <> "..."
+ src/Valiant/NamedParams.hs view
@@ -0,0 +1,174 @@+{-# LANGUAGE AllowAmbiguousTypes #-}+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE TypeFamilies #-}+{-# LANGUAGE UndecidableInstances #-}++-- | Named parameter support for query execution.+--+-- This module provides 'ToNamedParams', a type class for encoding+-- Haskell records as named query parameters. Field names in the record+-- correspond to @:name@ parameters in the SQL file.+--+-- @+-- data FindParams = FindParams+-- { orgId :: Int32+-- , role :: Text+-- , active :: Bool+-- } deriving (Generic, ToNamedParams)+--+-- findByRole :: NamedStatement FindParams [(Int32, Text)]+-- findByRole = mkStatementNamed+-- "SELECT id, name FROM users WHERE org_id = $1 AND role = $2 AND active = $3"+-- [23, 25, 16]+-- ["id", "name"]+-- ["orgId", "role", "active"] -- maps $1=orgId, $2=role, $3=active+-- "users/find_by_role.sql"+-- @+module Valiant.NamedParams+ ( ToNamedParams (..)+ , NamedStatement+ , mkStatementNamed+ ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.List (intercalate)+import Data.Map.Strict qualified as Map+import Data.Set qualified as Set+import Data.Text (Text)+import Data.Text qualified as T+import Data.Vector (Vector)+import Data.Vector qualified as V+import GHC.Generics+import Valiant.FromRow (FromRow (..))+import Valiant.Statement (Statement (..))+import Valiant.ToParams (EncodeField (..))+import PgWire.Protocol.Oid (Oid (..))++-- | A statement with named parameters. This is just a type alias —+-- at the wire level it's the same as a positional 'Statement'.+type NamedStatement p r = Statement p r++-- | Encode a record's fields as named parameters.+--+-- The default 'Generic' implementation extracts field names and values,+-- producing an association list of @(fieldName, encodedValue)@.+--+-- Derive with:+--+-- @+-- data MyParams = MyParams { userId :: Int32, role :: Text }+-- deriving stock (Generic)+-- deriving anyclass (ToNamedParams)+-- @+class ToNamedParams a where+ toNamedParamList :: a -> [(Text, Maybe ByteString)]++ default toNamedParamList :: (Generic a, GToNamedParams (Rep a)) => a -> [(Text, Maybe ByteString)]+ toNamedParamList = gToNamedParams . from+ {-# INLINE toNamedParamList #-}++-- | Generic implementation of 'ToNamedParams'.+class GToNamedParams f where+ gToNamedParams :: f p -> [(Text, Maybe ByteString)]++-- Datatype metadata: pass through+instance (GToNamedParams f) => GToNamedParams (D1 c f) where+ gToNamedParams (M1 x) = gToNamedParams x+ {-# INLINE gToNamedParams #-}++-- Constructor metadata: pass through+instance (GToNamedParams f) => GToNamedParams (C1 c f) where+ gToNamedParams (M1 x) = gToNamedParams x+ {-# INLINE gToNamedParams #-}++-- Product: concatenate fields+instance (GToNamedParams f, GToNamedParams g) => GToNamedParams (f :*: g) where+ gToNamedParams (f :*: g) = gToNamedParams f ++ gToNamedParams g+ {-# INLINE gToNamedParams #-}++-- Selector (record field): extract name and encode value+instance (Selector s, EncodeField a) => GToNamedParams (S1 s (Rec0 a)) where+ gToNamedParams m@(M1 (K1 x)) =+ let name = T.pack (selName m)+ in [(name, encodeField x)]+ {-# INLINE gToNamedParams #-}++-- Unit: no fields+instance GToNamedParams U1 where+ gToNamedParams U1 = []+ {-# INLINE gToNamedParams #-}++-- | Construct a 'Statement' that accepts named parameters via a record type.+--+-- The @paramNames@ list defines the mapping: @paramNames !! 0@ is the+-- name for @$1@, @paramNames !! 1@ for @$2@, etc. At runtime,+-- 'ToNamedParams' extracts the record's fields by name and reorders+-- them to match the positional indices.+--+-- If a record field name doesn't match any expected parameter name,+-- or an expected parameter has no matching field, a runtime error+-- is thrown with a clear message.+mkStatementNamed+ :: (ToNamedParams p, FromRow r)+ => String+ -- ^ SQL text (with positional $1, $2, ... — already preprocessed)+ -> [Int]+ -- ^ Parameter OIDs+ -> [String]+ -- ^ Result column names+ -> [String]+ -- ^ Parameter names in positional order ($1=first, $2=second, ...)+ -> String+ -- ^ Source .sql file path+ -> Statement p r+mkStatementNamed sqlStr oids colNames paramNames path =+ Statement+ { stmtSQL = BS8.pack sqlStr+ , stmtFile = path+ , stmtParamOids = V.fromList (map (Oid . fromIntegral) oids)+ , stmtEncode = encodeNamed (map T.pack paramNames)+ , stmtDecode = fromRow+ , stmtColumns = V.fromList (map BS8.pack colNames)+ }++-- | Build the encoder that reorders named params to positional order.+-- Validates at first use that all SQL param names have matching record+-- fields and vice versa, with clear error messages on mismatch.+encodeNamed :: (ToNamedParams p) => [Text] -> p -> Vector (Maybe ByteString)+encodeNamed paramNames p =+ let pairs = toNamedParamList p+ pairMap = Map.fromList pairs+ expectedNames = Set.fromList paramNames+ actualNames = Set.fromList (map fst pairs)+ missing = Set.toList (Set.difference expectedNames actualNames)+ extra = Set.toList (Set.difference actualNames expectedNames)+ n = length paramNames+ in case (missing, extra) of+ ([], []) ->+ V.fromListN n (map (\name -> Map.findWithDefault Nothing name pairMap) paramNames)+ _ ->+ error $ unlines $ concatMap (filter (not . null))+ [ [ "valiant: named parameter mismatch"+ , ""+ , " SQL parameters: " <> showNames paramNames+ , " Record fields: " <> showNames (map fst pairs)+ ]+ , if null missing then []+ else [ ""+ , " Missing fields (SQL expects these but the record doesn't have them):"+ , " " <> intercalate ", " (map (\n' -> ":" <> T.unpack n') missing)+ ]+ , if null extra then []+ else [ ""+ , " Extra fields (record has these but the SQL doesn't use them):"+ , " " <> intercalate ", " (map T.unpack extra)+ , ""+ , " Hint: remove unused fields from the record, or add"+ , " corresponding :name parameters to the SQL query."+ ]+ ]+ where+ showNames ns = intercalate ", " (map T.unpack ns)+{-# INLINE encodeNamed #-}
+ src/Valiant/Notify.hs view
@@ -0,0 +1,112 @@+-- | LISTEN/NOTIFY support for PostgreSQL asynchronous notifications.+--+-- == Delivery model+--+-- valiant's reader thread parks between queries, so NOTIFY messages that+-- arrive while the connection is idle are not dispatched until the next+-- query wakes the reader. 'waitForNotification' and its timeout variant+-- issue an empty simple query to flush pending NOTIFYs, then block on a+-- one-shot MVar filled by the handler.+--+-- In practice this means:+--+-- * Call 'waitForNotificationTimeout' (not 'waitForNotification') in+-- production code. A hung delivery returns 'Nothing' after the+-- deadline rather than blocking the caller forever.+--+-- * For fire-and-forget \"wake me whenever a NOTIFY arrives\" semantics,+-- poll: call 'waitForNotificationTimeout' in a loop with a short+-- timeout. Every iteration re-flushes the socket.+module Valiant.Notify+ ( Notification (..)+ , listen+ , unlisten+ , waitForNotification+ , waitForNotificationTimeout+ ) where++import Control.Concurrent (threadDelay)+import Control.Concurrent.Async (race)+import Control.Concurrent.MVar+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Int (Int32)+import PgWire.Async (AsyncWireConn (..))+import PgWire.Connection (Connection (..), simpleQuery)++-- | A notification received from PostgreSQL via @NOTIFY@ or @pg_notify()@.+data Notification = Notification+ { notifPid :: Int32+ -- ^ Process ID of the notifying backend.+ , notifChannel :: ByteString+ -- ^ Channel name the notification was sent on.+ , notifPayload :: ByteString+ -- ^ Optional payload string (empty if none was sent).+ }+ deriving stock (Show, Eq)++-- | Subscribe to a PostgreSQL notification channel by issuing @LISTEN@.+-- The channel name is automatically quoted as an identifier.+-- Use 'waitForNotification' or 'waitForNotificationTimeout' to receive+-- notifications after subscribing.+listen :: Connection -> ByteString -> IO ()+listen conn channel = do+ _ <- simpleQuery conn ("LISTEN " <> quoteIdent channel)+ pure ()++-- | Unsubscribe from a PostgreSQL notification channel by issuing @UNLISTEN@.+-- The channel name is automatically quoted as an identifier.+unlisten :: Connection -> ByteString -> IO ()+unlisten conn channel = do+ _ <- simpleQuery conn ("UNLISTEN " <> quoteIdent channel)+ pure ()++-- | Block until a notification arrives on any subscribed channel.+--+-- Registers a one-shot handler, sends an empty simple query to flush+-- any NOTIFYs already queued on the socket, then blocks on the handler.+-- If no NOTIFY was already pending at the time of the flush, this call+-- blocks until the next query on this connection wakes the reader.+-- Prefer 'waitForNotificationTimeout' for production use.+waitForNotification :: Connection -> IO Notification+waitForNotification conn = do+ notifVar <- newEmptyMVar+ installNotifyHandler conn notifVar+ -- Send an empty query to flush any pending notifications from the server+ _ <- simpleQuery conn ""+ takeMVar notifVar++-- | Wait for a notification with a timeout (in seconds).+-- Returns 'Nothing' if the timeout expires.+waitForNotificationTimeout :: Connection -> Double -> IO (Maybe Notification)+waitForNotificationTimeout conn seconds = do+ notifVar <- newEmptyMVar+ installNotifyHandler conn notifVar+ _ <- simpleQuery conn ""+ let micros = round (seconds * 1000000) :: Int+ result <- race (threadDelay micros) (takeMVar notifVar)+ -- Restore default handler regardless of outcome+ writeIORef (awcNotifyHandler (connAsync conn)) (\_ _ _ -> pure ())+ pure $ case result of+ Left () -> Nothing+ Right n -> Just n++-- Internal ------------------------------------------------------------------++-- | Install a one-shot notification handler that fills the given MVar.+installNotifyHandler :: Connection -> MVar Notification -> IO ()+installNotifyHandler conn notifVar = do+ let handler pid channel payload = do+ let notif = Notification pid channel payload+ _ <- tryPutMVar notifVar notif+ -- Restore the default no-op handler after delivery+ writeIORef (awcNotifyHandler (connAsync conn)) (\_ _ _ -> pure ())+ writeIORef (awcNotifyHandler (connAsync conn)) handler++-- | Simple identifier quoting (double-quote).+quoteIdent :: ByteString -> ByteString+quoteIdent ident = "\"" <> BS8.concatMap escapeQuote ident <> "\""+ where+ escapeQuote '"' = "\"\""+ escapeQuote c = BS8.singleton c
+ src/Valiant/Pipeline.hs view
@@ -0,0 +1,160 @@+{-# LANGUAGE ExistentialQuantification #-}+{-# LANGUAGE LambdaCase #-}+{-# OPTIONS_GHC -fno-full-laziness #-}++-- | Pipelined query execution for batching multiple independent queries+-- into a single network round-trip.+--+-- The PostgreSQL extended query protocol allows sending multiple+-- Bind+Execute sequences before a single Sync. This module provides+-- an 'Applicative' interface for composing independent queries that+-- are sent together and collected in order.+--+-- @+-- (user, posts) <- 'runPipeline' conn $ (,)+-- '<$>' 'pipeFetchAll' listPostsByUser 42+-- '<*>' 'pipeFetchOne' findUserById 42+-- -- Both queries sent in one round-trip, results collected in order+-- @+--+-- This eliminates the N+1 query problem at the driver level.+module Valiant.Pipeline+ ( Pipeline+ , pipeFetchOne+ , pipeFetchAll+ , pipeFetchScalar+ , pipeExecute+ , runPipeline+ ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Vector (Vector)+import Data.Vector qualified as V+import PgWire.Async (Request (..), Response (..), ResponseCollector (..), submitRequest)+import PgWire.Connection (Connection (..))+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Frontend+import Valiant.Execute (ensurePrepared)+import Valiant.Statement (Statement (..))++-- | A pipeline of queries to execute in a single round-trip.+-- Use the 'Applicative' interface to compose independent queries.+data Pipeline a+ = PureP a+ | forall p r. QueryP+ (Statement p r)+ p+ ([Vector (Maybe ByteString)] -> IO a)+ -- ^ Decoder: takes collected rows and produces the result+ | forall b. ApP (Pipeline (b -> a)) (Pipeline b)++instance Functor Pipeline where+ fmap f (PureP a) = PureP (f a)+ fmap f (QueryP s p dec) = QueryP s p (fmap f . dec)+ fmap f (ApP pf px) = ApP (fmap (f .) pf) px++instance Applicative Pipeline where+ pure = PureP+ (<*>) = ApP++-- | Pipeline a query that returns zero or one row.+pipeFetchOne :: Statement p r -> p -> Pipeline (Maybe r)+pipeFetchOne stmt params = QueryP stmt params $ \case+ [] -> pure Nothing+ (row : _) -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure (Just val)+{-# INLINE pipeFetchOne #-}++-- | Pipeline a query that returns all rows.+pipeFetchAll :: Statement p r -> p -> Pipeline [r]+pipeFetchAll stmt params = QueryP stmt params $ \rows ->+ mapM (\row -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure val) rows+{-# INLINE pipeFetchAll #-}++-- | Pipeline a query that returns a single scalar value.+pipeFetchScalar :: Statement p r -> p -> Pipeline r+pipeFetchScalar stmt params = QueryP stmt params $ \case+ [row] -> case stmtDecode stmt row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right val -> pure val+ [] -> throwPgWire (DecodeError "pipeFetchScalar: no rows")+ _ -> throwPgWire (DecodeError "pipeFetchScalar: more than one row")+{-# INLINE pipeFetchScalar #-}++-- | Pipeline a command (INSERT\/UPDATE\/DELETE).+--+-- Row counts are not available in pipeline mode because the batch collector+-- does not track individual command tags. Use 'Valiant.Execute.execute' or+-- 'Valiant.Execute.executeBatch' outside a pipeline if you need the+-- rows-affected count.+pipeExecute :: Statement p () -> p -> Pipeline ()+pipeExecute stmt params = QueryP stmt params $ \_ -> pure ()+{-# INLINE pipeExecute #-}++-- | Execute all queries in the pipeline in a single network round-trip.+--+-- @+-- (user, posts, count) <- runPipeline conn $ (,,)+-- '<$>' pipeFetchOne findUserById 42+-- '<*>' pipeFetchAll listRecentPosts 10+-- '<*>' pipeFetchScalar countUsers ()+-- @+runPipeline :: Connection -> Pipeline a -> IO a+runPipeline conn pipeline = do+ let queries = collectQueries pipeline+ case queries of+ [] -> evalPipeline pipeline []+ _ -> do+ -- Ensure all statements are prepared+ names <- mapM (ensurePreparedPipe conn) queries++ -- Build all Bind+Execute messages with one Sync+ let msgs = concatMap (\(PipeQuery stmt p, name) ->+ let encodedParams = stmtEncode stmt p+ in [ Bind "" name (V.singleton BinaryFormat) encodedParams (V.singleton BinaryFormat)+ , Execute "" 0+ ]) (zip queries names)+ ++ [Sync]++ -- Submit as a single batch request+ resp <- submitRequest (connAsync conn) $ ReqExtendedQuery msgs (CollectBatch (length queries))++ case resp of+ RespBatchRows results -> evalPipeline pipeline results+ _ -> throwPgWire (ProtocolError "runPipeline: unexpected response type")++-- Internal ----------------------------------------------------------------++data PipeQuery = forall p r. PipeQuery (Statement p r) p++collectQueries :: Pipeline a -> [PipeQuery]+collectQueries (PureP _) = []+collectQueries (QueryP stmt params _) = [PipeQuery stmt params]+collectQueries (ApP pf px) = collectQueries pf ++ collectQueries px++evalPipeline :: Pipeline a -> [[Vector (Maybe ByteString)]] -> IO a+evalPipeline pipeline allResults = do+ ref <- newIORef allResults+ eval ref pipeline++eval :: IORef [[Vector (Maybe ByteString)]] -> Pipeline a -> IO a+eval _ (PureP a) = pure a+eval ref (QueryP _ _ decode) = do+ rs <- readIORef ref+ case rs of+ (rows : rest) -> do+ writeIORef ref rest+ decode rows+ [] -> throwPgWire (DecodeError "pipeline: result underflow")+eval ref (ApP pf px) = do+ f <- eval ref pf+ x <- eval ref px+ pure (f x)++ensurePreparedPipe :: Connection -> PipeQuery -> IO ByteString+ensurePreparedPipe conn (PipeQuery stmt _) = ensurePrepared conn stmt
+ src/Valiant/Statement.hs view
@@ -0,0 +1,116 @@+{-# LANGUAGE DataKinds #-}+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-simplifiable-class-constraints #-}++-- | The 'Statement' type: a compile-time validated SQL query.+--+-- @p@ is the parameter type, @r@ is the result type. The GHC source plugin+-- verifies that these types match the SQL query's parameters and columns.+--+-- Create statements with 'queryFile' (validated by the plugin) or+-- 'mkStatement' (for manual\/test construction).+module Valiant.Statement+ ( Statement (..)+ , query+ , queryFile+ , queryFileAs+ , mkStatement+ ) where++import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.Vector (Vector)+import Data.Vector qualified as V+import GHC.TypeLits (ErrorMessage (..), TypeError)+import Valiant.FromRow (FromRow (..))+import PgWire.Protocol.Oid (Oid (..))+import Valiant.ToParams (ToParams (..))++-- | A compile-time validated SQL statement.+--+-- @p@ is the parameter type (e.g., @Int32@ or @(Text, Maybe Text)@).+-- @r@ is the result type (e.g., @(Int32, Text)@ or @()@ for commands).+data Statement p r = Statement+ { stmtSQL :: ByteString+ -- ^ The raw SQL text.+ , stmtFile :: FilePath+ -- ^ The source @.sql@ file path (for error messages).+ , stmtParamOids :: Vector Oid+ -- ^ PostgreSQL type OIDs for each parameter.+ , stmtEncode :: p -> Vector (Maybe ByteString)+ -- ^ Encode parameters to binary format.+ , stmtDecode :: Vector (Maybe ByteString) -> Either String r+ -- ^ Decode a result row from binary format.+ , stmtColumns :: Vector ByteString+ -- ^ Column names from the query result.+ }++-- | Constraint that fires a helpful 'TypeError' if the valiant plugin is+-- not loaded. The plugin removes this constraint during the parse-phase+-- rewrite (by replacing @queryFile@ with @mkStatement@). If the plugin+-- isn't loaded, GHC tries to solve this constraint and hits the+-- 'TypeError' instance.+class ValiantPluginRequired++instance+ TypeError+ ( 'Text "queryFile/queryFileAs requires the Valiant.Plugin GHC source plugin."+ ':$$: 'Text ""+ ':$$: 'Text "Add this to your module:"+ ':$$: 'Text " {-# OPTIONS_GHC -fplugin=Valiant.Plugin #-}"+ ':$$: 'Text ""+ ':$$: 'Text "Or add to your .cabal file:"+ ':$$: 'Text " ghc-options: -fplugin=Valiant.Plugin"+ )+ => ValiantPluginRequired++-- | Reference a @.sql@ file for compile-time validation.+--+-- The GHC source plugin intercepts this call and rewrites it to+-- 'mkStatement' with embedded SQL, OIDs, column names, and file path.+--+-- Without the plugin, this produces a compile-time type error with+-- instructions on how to enable the plugin.+queryFile :: ValiantPluginRequired => FilePath -> Statement p r+queryFile _ = error "unreachable: valiant plugin not loaded"++-- | Like 'queryFile' but for named result types with 'FromRow'.+queryFileAs :: ValiantPluginRequired => FilePath -> Statement p r+queryFileAs _ = error "unreachable: valiant plugin not loaded"++-- | Inline SQL with compile-time validation.+--+-- Write SQL directly in Haskell source. The GHC plugin intercepts this+-- call, hashes the SQL text, looks up the type metadata in @.valiant\/@,+-- and rewrites to 'mkStatement' with the correct types.+--+-- Requires running @valiant prepare@ with the same SQL text present in a+-- @.sql@ file or registered via @valiant prepare --inline@.+--+-- @+-- findById :: Statement Int32 (Int32, Text, Maybe Text)+-- findById = query "SELECT id, name, email FROM users WHERE id = $1"+-- @+query :: ValiantPluginRequired => String -> Statement p r+query _ = error "unreachable: valiant plugin not loaded"++-- | Construct a 'Statement' from literal data embedded by the plugin.+-- The plugin rewrites @queryFile "path.sql"@ into a call to this function+-- with the SQL text, parameter OIDs, column names, and file path baked in.+-- Type class constraints resolve the encoder and decoder at compile time.+mkStatement+ :: (ToParams p, FromRow r)+ => String+ -> [Int]+ -> [String]+ -> String+ -> Statement p r+mkStatement sqlStr oids colNames path =+ Statement+ { stmtSQL = BS8.pack sqlStr+ , stmtFile = path+ , stmtParamOids = V.fromList (map (Oid . fromIntegral) oids)+ , stmtEncode = toParams+ , stmtDecode = fromRow+ , stmtColumns = V.fromList (map BS8.pack colNames)+ }
+ src/Valiant/Streaming.hs view
@@ -0,0 +1,220 @@+{-# OPTIONS_GHC -fno-full-laziness #-}++-- | Streaming query results using server-side cursors.+--+-- Instead of fetching all rows at once, a cursor fetches rows in batches,+-- allowing processing of large result sets without loading everything into+-- memory. Must be called inside a transaction.+--+-- Cursors require exclusive access to the connection's wire, so the entire+-- cursor lifecycle runs inside 'submitExclusive'.+--+-- @+-- 'Valiant.Transaction.withTransaction' pool $ \\tx -> do+-- 'withCursor' (txConn tx) myQuery params 100 $ \\cursor -> do+-- let loop = do+-- batch <- 'fetchBatch' cursor 100+-- unless (null batch) $ do+-- mapM_ processRow batch+-- loop+-- loop+-- @+module Valiant.Streaming+ ( withCursor+ , fetchBatch+ , fetchAllCursor+ , CursorState+ ) where++import Control.Exception (SomeException, catch, onException)+import Control.Monad (when)+import Data.ByteString (ByteString)+import Data.ByteString.Char8 qualified as BS8+import Data.IORef+import Data.Vector (Vector)+import Data.Vector qualified as V+import Data.Word (Word64)+import PgWire.Async (submitExclusive)+import PgWire.Connection (Connection (..))+import PgWire.Error (PgWireError (..), throwPgWire)+import PgWire.Protocol.Backend+import PgWire.Protocol.Frontend+import PgWire.Protocol.Oid qualified as Oid+import Valiant.Statement (Statement (..))+import PgWire.Wire (WireConn, recvBackendMsg, sendFrontendMsg, sendFrontendMsgs)+import System.IO.Unsafe (unsafePerformIO)++-- | State of an open cursor, valid only within the 'withCursor' callback.+-- Do not retain references to this value outside the callback scope.+data CursorState = CursorState+ { csName :: ByteString+ -- ^ The cursor name on the server.+ , csWire :: WireConn+ -- ^ Direct wire access (valid only inside exclusive mode).+ , csTxRef :: IORef TxStatus+ -- ^ Transaction status ref.+ , csExhausted :: IORef Bool+ -- ^ Whether the cursor has returned all rows.+ , csPreparedBatch :: IORef Int+ -- ^ Batch size currently baked into the unnamed prepared FETCH+ -- statement on the server. 'fetchBatch' re-parses only when the+ -- requested size differs from this value.+ }++-- | Open a parameterized cursor for a statement, run an action, then close.+--+-- The statement's parameters are bound via the extended query protocol+-- (Parse\/Bind), so parameterized queries work correctly.+--+-- Must be called inside a transaction. The entire cursor lifecycle runs+-- in exclusive mode (direct socket access).+withCursor+ :: Connection+ -> Statement p r+ -> p+ -> Int+ -- ^ Batch size hint (used as default for 'fetchBatch')+ -> (CursorState -> IO a)+ -> IO a+withCursor conn stmt params batchSize action =+ submitExclusive (connAsync conn) $ \wc txRef -> do+ cursorName <- freshCursorName++ -- Declare the cursor and prepare a FETCH statement for the hinted+ -- batch size in a single pipelined round-trip. Subsequent+ -- 'fetchBatch' calls reuse the unnamed prepared FETCH as long as+ -- the requested batch size matches.+ let declareSql = "DECLARE " <> cursorName <> " NO SCROLL CURSOR FOR " <> stmtSQL stmt+ fetchSql = "FETCH FORWARD " <> BS8.pack (show batchSize) <> " FROM " <> cursorName+ encodedParams = stmtEncode stmt params+ paramOids = V.map Oid.unOid (stmtParamOids stmt)++ sendFrontendMsgs wc+ [ Parse "" declareSql paramOids+ , Bind "" "" (V.singleton BinaryFormat) encodedParams V.empty+ , Execute "" 0+ , Parse "" fetchSql V.empty+ , Sync+ ]++ waitDeclareComplete wc txRef++ exhausted <- newIORef False+ preparedBatch <- newIORef batchSize+ let cs = CursorState cursorName wc txRef exhausted preparedBatch+ closeCursor =+ (do sendFrontendMsg wc (Query ("CLOSE " <> cursorName))+ collectSimpleDiscard wc txRef+ ) `catch` \(_ :: SomeException) -> pure ()+ result <- action cs `onException` closeCursor+ closeCursor+ pure result++-- | Fetch the next batch of rows from a cursor.+--+-- Returns an empty list when the cursor is exhausted. Subsequent calls+-- after exhaustion return empty immediately without a round-trip.+fetchBatch :: CursorState -> Int -> IO [Vector (Maybe ByteString)]+fetchBatch cs n = do+ done <- readIORef (csExhausted cs)+ if done+ then pure []+ else do+ prepared <- readIORef (csPreparedBatch cs)+ let bindExecSync =+ [ Bind "" "" V.empty V.empty (V.singleton BinaryFormat)+ , Execute "" 0+ , Sync+ ]+ msgs+ | prepared == n = bindExecSync+ | otherwise =+ let fetchSql = "FETCH FORWARD " <> BS8.pack (show n) <> " FROM " <> csName cs+ in Parse "" fetchSql V.empty : bindExecSync+ when (prepared /= n) $ writeIORef (csPreparedBatch cs) n+ sendFrontendMsgs (csWire cs) msgs+ rows <- collectFetchResults (csWire cs) (csTxRef cs)+ if null rows+ then do+ writeIORef (csExhausted cs) True+ pure []+ else pure rows++-- | Open a cursor, drain every batch, decode and return all rows.+--+-- Convenience wrapper used by streaming adapters that need to materialise+-- a cursor result as a list. Must be called inside a transaction.+fetchAllCursor :: Connection -> Statement p r -> p -> Int -> IO [r]+fetchAllCursor conn stmt params batchSize =+ withCursor conn stmt params batchSize $ \cs ->+ -- Accumulate batches as a list-of-lists, concat once at the end.+ -- This is O(N) total work, vs ~3-5x more with `reverse b ++ acc`,+ -- because we avoid reversing each batch and the final list.+ let go !acc = do+ batch <- fetchBatch cs batchSize+ if null batch+ then pure (concat (reverse acc))+ else do+ decoded <- mapM (decodeRow (stmtDecode stmt)) batch+ go (decoded : acc)+ in go []+ where+ decodeRow decode row = case decode row of+ Left err -> throwPgWire (DecodeError (BS8.pack err))+ Right !val -> pure val++-- Internal ------------------------------------------------------------------++waitDeclareComplete :: WireConn -> IORef TxStatus -> IO ()+waitDeclareComplete wc txRef = go+ where+ go = do+ msg <- recvBackendMsg wc+ case msg of+ ParseComplete -> go+ BindComplete -> go+ CommandComplete _ -> go+ ReadyForQuery status -> writeIORef txRef status+ ErrorResponse err -> throwPgWire (QueryError err)+ NoticeResponse _ -> go+ other -> throwPgWire (ProtocolError ("Unexpected in DECLARE cursor: " <> BS8.pack (show other)))++collectFetchResults :: WireConn -> IORef TxStatus -> IO [Vector (Maybe ByteString)]+collectFetchResults wc txRef = go []+ where+ go !acc = do+ msg <- recvBackendMsg wc+ case msg of+ ParseComplete -> go acc+ BindComplete -> go acc+ RowDescription _ -> go acc+ DataRow vals -> go (vals : acc)+ CommandComplete _ -> go acc+ EmptyQueryResponse -> go acc+ ReadyForQuery status -> do+ writeIORef txRef status+ pure (reverse acc)+ ErrorResponse err -> throwPgWire (QueryError err)+ NoticeResponse _ -> go acc+ other -> throwPgWire (ProtocolError ("Unexpected in cursor fetch: " <> BS8.pack (show other)))++-- | Discard all results from a simple query.+collectSimpleDiscard :: WireConn -> IORef TxStatus -> IO ()+collectSimpleDiscard wc txRef = go+ where+ go = do+ msg <- recvBackendMsg wc+ case msg of+ ReadyForQuery status -> writeIORef txRef status+ ErrorResponse err -> throwPgWire (QueryError err)+ _ -> go++-- | Global counter for unique cursor names.+{-# NOINLINE cursorCounter #-}+cursorCounter :: IORef Word64+cursorCounter = unsafePerformIO (newIORef 0)++freshCursorName :: IO ByteString+freshCursorName = do+ n <- atomicModifyIORef' cursorCounter (\n -> (n + 1, n))+ pure ("valiant_cur_" <> BS8.pack (show n))
+ src/Valiant/ToParams.hs view
@@ -0,0 +1,154 @@+{-# LANGUAGE DefaultSignatures #-}+{-# LANGUAGE UndecidableInstances #-}+{-# OPTIONS_GHC -Wno-orphans #-}++-- | Encode Haskell values into PostgreSQL binary parameter format.+--+-- 'ToParams' encodes query parameter tuples, while 'EncodeField' handles+-- individual fields. @Maybe a@ encodes as SQL NULL when @Nothing@.+-- Instances are provided for @()@, single values, tuples up to 16,+-- and any @Generic@ type via @DefaultSignatures@.+--+-- For custom record types, derive via @Generic@:+--+-- @+-- data InsertParams = InsertParams+-- { ipName :: Text+-- , ipEmail :: Maybe Text+-- } deriving stock (Generic)+-- deriving anyclass (ToParams)+-- @+--+-- Fields are encoded positionally (first field → $1, etc.).+module Valiant.ToParams+ ( ToParams (..)+ , EncodeField (..)+ ) where++import Data.ByteString (ByteString)+import Data.Vector (Vector)+import Data.Vector qualified as V+import GHC.Generics+import Valiant.Binary.Encode ()+import PgWire.Binary.Types (PgEncode (..))++-- | Encode a single field, handling @Maybe@ for nullable parameters.+class EncodeField a where+ encodeField :: a -> Maybe ByteString++instance {-# OVERLAPPABLE #-} (PgEncode a) => EncodeField a where+ encodeField = Just . pgEncode+ {-# INLINE encodeField #-}++instance (PgEncode a) => EncodeField (Maybe a) where+ encodeField Nothing = Nothing+ encodeField (Just a) = Just (pgEncode a)+ {-# INLINE encodeField #-}++-- | Encode query parameters into a vector of nullable byte strings.+class ToParams a where+ toParams :: a -> Vector (Maybe ByteString)+ default toParams :: (Generic a, GToParams (Rep a)) => a -> Vector (Maybe ByteString)+ toParams = gToParams . from++-- | Generic helper for positional parameter encoding. Internal — not exported.+class GToParams f where+ gToParams :: f p -> Vector (Maybe ByteString)++-- Datatype metadata — delegate+instance (GToParams f) => GToParams (M1 D c f) where+ gToParams (M1 x) = gToParams x+ {-# INLINE gToParams #-}++-- Constructor metadata — delegate+instance (GToParams f) => GToParams (M1 C c f) where+ gToParams (M1 x) = gToParams x+ {-# INLINE gToParams #-}++-- Selector metadata — delegate+instance (GToParams f) => GToParams (M1 S c f) where+ gToParams (M1 x) = gToParams x+ {-# INLINE gToParams #-}++-- Leaf field — encode one field+instance (EncodeField a) => GToParams (K1 R a) where+ gToParams (K1 a) = V.singleton (encodeField a)+ {-# INLINE gToParams #-}++-- Product — concatenate left and right.+-- V.++ is O(n) per call, but with INLINE GHC fuses the chain for small records.+-- For records with > ~20 fields, a DList-based approach would be better.+instance (GToParams f, GToParams g) => GToParams (f :*: g) where+ gToParams (l :*: r) = gToParams l V.++ gToParams r+ {-# INLINE gToParams #-}++-- Unit — no fields+instance GToParams U1 where+ gToParams U1 = V.empty+ {-# INLINE gToParams #-}++instance ToParams () where+ toParams () = V.empty++instance {-# OVERLAPPABLE #-} (EncodeField a) => ToParams a where+ toParams a = V.singleton (encodeField a)++instance (EncodeField a, EncodeField b) => ToParams (a, b) where+ toParams (a, b) = V.fromListN 2 [encodeField a, encodeField b]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c) => ToParams (a, b, c) where+ toParams (a, b, c) = V.fromListN 3 [encodeField a, encodeField b, encodeField c]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d) => ToParams (a, b, c, d) where+ toParams (a, b, c, d) = V.fromListN 4 [encodeField a, encodeField b, encodeField c, encodeField d]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e) => ToParams (a, b, c, d, e) where+ toParams (a, b, c, d, e) = V.fromListN 5 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f) => ToParams (a, b, c, d, e, f) where+ toParams (a, b, c, d, e, f) = V.fromListN 6 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g) => ToParams (a, b, c, d, e, f, g) where+ toParams (a, b, c, d, e, f, g) = V.fromListN 7 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h) => ToParams (a, b, c, d, e, f, g, h) where+ toParams (a, b, c, d, e, f, g, h) = V.fromListN 8 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i) => ToParams (a, b, c, d, e, f, g, h, i) where+ toParams (a, b, c, d, e, f, g, h, i) = V.fromListN 9 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j) => ToParams (a, b, c, d, e, f, g, h, i, j) where+ toParams (a, b, c, d, e, f, g, h, i, j) = V.fromListN 10 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j, EncodeField k) => ToParams (a, b, c, d, e, f, g, h, i, j, k) where+ toParams (a, b, c, d, e, f, g, h, i, j, k) = V.fromListN 11 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j, encodeField k]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j, EncodeField k, EncodeField l) => ToParams (a, b, c, d, e, f, g, h, i, j, k, l) where+ toParams (a, b, c, d, e, f, g, h, i, j, k, l) = V.fromListN 12 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j, encodeField k, encodeField l]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j, EncodeField k, EncodeField l, EncodeField m) => ToParams (a, b, c, d, e, f, g, h, i, j, k, l, m) where+ toParams (a, b, c, d, e, f, g, h, i, j, k, l, m) = V.fromListN 13 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j, encodeField k, encodeField l, encodeField m]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j, EncodeField k, EncodeField l, EncodeField m, EncodeField n) => ToParams (a, b, c, d, e, f, g, h, i, j, k, l, m, n) where+ toParams (a, b, c, d, e, f, g, h, i, j, k, l, m, n) = V.fromListN 14 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j, encodeField k, encodeField l, encodeField m, encodeField n]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j, EncodeField k, EncodeField l, EncodeField m, EncodeField n, EncodeField o) => ToParams (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) where+ toParams (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o) = V.fromListN 15 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j, encodeField k, encodeField l, encodeField m, encodeField n, encodeField o]+ {-# INLINE toParams #-}++instance (EncodeField a, EncodeField b, EncodeField c, EncodeField d, EncodeField e, EncodeField f, EncodeField g, EncodeField h, EncodeField i, EncodeField j, EncodeField k, EncodeField l, EncodeField m, EncodeField n, EncodeField o, EncodeField p) => ToParams (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) where+ toParams (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) = V.fromListN 16 [encodeField a, encodeField b, encodeField c, encodeField d, encodeField e, encodeField f, encodeField g, encodeField h, encodeField i, encodeField j, encodeField k, encodeField l, encodeField m, encodeField n, encodeField o, encodeField p]+ {-# INLINE toParams #-}
+ src/Valiant/Transaction.hs view
@@ -0,0 +1,312 @@+-- | Transaction management with savepoint support.+--+-- @+-- 'withTransaction' pool $ \\tx -> do+-- execute (txConn tx) insertUser (\"Alice\", email)+-- 'withSavepoint' tx $ \\_ -> do+-- execute (txConn tx) riskyOperation params+-- -- If this throws, only the savepoint is rolled back,+-- -- not the entire transaction.+-- @+module Valiant.Transaction+ ( Transaction (..)+ , IsolationLevel (..)+ , TransactionMode (..)+ , defaultTransactionMode+ , withTransaction+ , withTransaction_+ , withTransactionLevel+ , withTransactionMode+ , withTransactionConn+ , withTransactionLevelConn+ , withTransactionModeConn+ , withReadOnlyTransaction+ , withDeferrableTransaction+ , withTransactionRetry+ , withTransactionRetryIf+ , withSavepoint+ ) where++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.Word (Word64)+import PgWire.Connection (Connection, simpleQuery)+import PgWire.Error (PgWireError (..))+import PgWire.Protocol.Backend (PgError (..))+import PgWire.Pool (Pool, withResource)+import System.IO.Unsafe (unsafePerformIO)++-- | A database transaction handle. Wraps a 'Connection' that has an active+-- @BEGIN@ block. Use 'txConn' to obtain the underlying connection for+-- executing queries within the transaction.+newtype Transaction = Transaction+ { txConn :: Connection+ -- ^ The underlying connection with an active transaction.+ }++-- | PostgreSQL transaction isolation level.+--+-- See the <https://www.postgresql.org/docs/current/transaction-iso.html PostgreSQL documentation>+-- for the semantics of each level.+data IsolationLevel+ = ReadCommitted+ -- ^ The default. Each statement sees a snapshot as of the start of that statement.+ | RepeatableRead+ -- ^ All statements in the transaction see a snapshot as of the first non-transaction-control+ -- statement. Throws a serialization error on write conflicts.+ | Serializable+ -- ^ The strictest level. Transactions behave as if executed sequentially.+ -- Throws a serialization error on detected conflicts.+ deriving stock (Show, Eq)++-- | Run an action inside a transaction with default isolation ('ReadCommitted').+--+-- Acquires a connection from the pool, issues @BEGIN@, runs the action, and+-- commits. If the action throws an exception, the transaction is rolled back+-- and the exception is re-raised. The connection is always returned to the pool.+--+-- @+-- withTransaction pool $ \\tx -> do+-- execute (txConn tx) insertStmt params+-- execute (txConn tx) updateStmt params+-- -- Auto-committed here, or rolled back on exception.+-- @+withTransaction :: Pool -> (Transaction -> IO a) -> IO a+withTransaction = withTransactionLevel ReadCommitted++-- | Like 'withTransaction' but discards the return value.+--+-- @+-- withTransaction_ pool $ \\tx ->+-- execute (txConn tx) insertStmt params+-- @+withTransaction_ :: Pool -> (Transaction -> IO ()) -> IO ()+withTransaction_ = withTransaction+{-# INLINE withTransaction_ #-}++-- | Run an action inside a transaction with the given isolation level.+--+-- Behaves like 'withTransaction' but issues @BEGIN ISOLATION LEVEL ...@ with+-- the specified level. Uses 'mask' to ensure the @BEGIN@\/@COMMIT@\/@ROLLBACK@+-- sequence cannot be interrupted by async exceptions.+--+-- If COMMIT fails (e.g., deferred constraint violation), the transaction is+-- rolled back so the connection is returned to the pool in a clean state.+-- Note: if the network drops after the server commits but before the client+-- receives the response, the client will attempt a ROLLBACK on an already-+-- committed transaction. This is an inherent TCP limitation — without two-+-- phase commit, the client cannot distinguish "committed but response lost"+-- from "failed to commit."+withTransactionLevel :: IsolationLevel -> Pool -> (Transaction -> IO a) -> IO a+withTransactionLevel level pool action =+ withResource pool $ \conn -> mask $ \restore -> do+ _ <- simpleQuery conn (beginStatement level)+ result <- restore (action (Transaction conn)) `onException` rollback conn+ _ <- simpleQuery conn "COMMIT" `onException` rollback conn+ pure result++-- | Run an action inside a transaction on an existing connection.+-- Unlike 'withTransaction' which acquires a connection from a pool, this+-- function operates on a connection you already hold (e.g., from 'withResource').+--+-- @+-- withResource pool $ \\conn -> do+-- -- some setup work on conn ...+-- withTransactionConn conn $ \\tx ->+-- execute (txConn tx) insertStmt params+-- @+withTransactionConn :: Connection -> (Transaction -> IO a) -> IO a+withTransactionConn = withTransactionLevelConn ReadCommitted++-- | Like 'withTransactionConn' but with a specific isolation level.+withTransactionLevelConn :: IsolationLevel -> Connection -> (Transaction -> IO a) -> IO a+withTransactionLevelConn level conn action = mask $ \restore -> do+ _ <- simpleQuery conn (beginStatement level)+ result <- restore (action (Transaction conn)) `onException` rollback conn+ _ <- simpleQuery conn "COMMIT" `onException` rollback conn+ pure result++-- | Full transaction mode configuration, bundling isolation level,+-- read/write mode, and deferrable flag.+--+-- @+-- let mode = defaultTransactionMode+-- { tmIsolation = Serializable+-- , tmReadOnly = True+-- , tmDeferrable = True+-- }+-- withTransactionMode mode pool $ \\tx -> ...+-- @+data TransactionMode = TransactionMode+ { tmIsolation :: !IsolationLevel+ -- ^ Isolation level. Default: 'ReadCommitted'.+ , tmReadOnly :: !Bool+ -- ^ If 'True', issues @READ ONLY@. Default: 'False'.+ , tmDeferrable :: !Bool+ -- ^ If 'True', issues @DEFERRABLE@. Only meaningful with+ -- 'Serializable' + 'tmReadOnly'. Default: 'False'.+ }+ deriving stock (Show, Eq)++-- | Default transaction mode: 'ReadCommitted', read-write, not deferrable.+defaultTransactionMode :: TransactionMode+defaultTransactionMode = TransactionMode+ { tmIsolation = ReadCommitted+ , tmReadOnly = False+ , tmDeferrable = False+ }++-- | Run an action inside a transaction with the given 'TransactionMode'.+--+-- This is the most general transaction function, supporting all combinations+-- of isolation level, read/write mode, and deferrable flag.+withTransactionMode :: TransactionMode -> Pool -> (Transaction -> IO a) -> IO a+withTransactionMode mode pool action =+ withResource pool $ \conn -> withTransactionModeConn mode conn action++-- | Like 'withTransactionMode' but on an existing connection.+withTransactionModeConn :: TransactionMode -> Connection -> (Transaction -> IO a) -> IO a+withTransactionModeConn mode conn action = mask $ \restore -> do+ _ <- simpleQuery conn (beginModeStatement mode)+ result <- restore (action (Transaction conn)) `onException` rollback conn+ _ <- simpleQuery conn "COMMIT" `onException` rollback conn+ pure result++-- | Run an action inside a @READ ONLY@ transaction. Postgres guarantees+-- no writes can occur, which enables use of standby replicas.+--+-- Uses the default 'ReadCommitted' isolation level. Combine with+-- 'withTransactionMode' if you need a different level.+--+-- @+-- users <- withReadOnlyTransaction pool $ \\tx ->+-- fetchAll (txConn tx) listAllUsers ()+-- @+withReadOnlyTransaction :: Pool -> (Transaction -> IO a) -> IO a+withReadOnlyTransaction =+ withTransactionMode defaultTransactionMode { tmReadOnly = True }++-- | Run an action inside a @SERIALIZABLE READ ONLY DEFERRABLE@ transaction.+--+-- PostgreSQL guarantees a consistent snapshot without risk of serialization+-- failure, making this ideal for long-running analytics or reporting queries.+-- May block briefly at the start while PostgreSQL finds a safe snapshot.+--+-- @+-- report <- withDeferrableTransaction pool $ \\tx ->+-- fetchAll (txConn tx) bigAnalyticsQuery ()+-- @+withDeferrableTransaction :: Pool -> (Transaction -> IO a) -> IO a+withDeferrableTransaction =+ withTransactionMode+ TransactionMode+ { tmIsolation = Serializable+ , tmReadOnly = True+ , tmDeferrable = True+ }++-- | Run a 'Serializable' transaction, automatically retrying on+-- serialization failures (SQLSTATE 40001) up to @maxRetries@ times.+--+-- The action is re-run from scratch on each retry (a new BEGIN is issued).+-- If all retries are exhausted, the last serialization error is thrown.+--+-- @+-- withTransactionRetry 3 pool $ \\tx -> do+-- balance <- fetchScalar (txConn tx) getBalance userId+-- execute (txConn tx) setBalance (balance + amount, userId)+-- @+withTransactionRetry :: Int -> Pool -> (Transaction -> IO a) -> IO a+withTransactionRetry maxRetries =+ withTransactionRetryIf isSerializationErr maxRetries+ where+ isSerializationErr (QueryError err) = pgCode err == "40001"+ isSerializationErr _ = False++-- | Like 'withTransactionRetry' but with a custom predicate to decide+-- which errors should trigger a retry.+--+-- @+-- -- Retry on both serialization failures and deadlocks+-- let shouldRetry err = isSerializationError err || isDeadlockError err+-- withTransactionRetryIf shouldRetry 3 pool $ \\tx -> ...+-- @+withTransactionRetryIf+ :: (PgWireError -> Bool)+ -> Int+ -> Pool+ -> (Transaction -> IO a)+ -> IO a+withTransactionRetryIf shouldRetry maxRetries pool action = go 0+ where+ go !attempt = do+ result <- try (withTransactionLevel Serializable pool action)+ case result of+ Right val -> pure val+ Left err+ | shouldRetry err && attempt < maxRetries -> go (attempt + 1)+ | otherwise -> throwIO err++-- | Run an action inside a savepoint within an existing transaction.+--+-- If the action throws an exception, the savepoint is rolled back but+-- the outer transaction remains active. If the action succeeds, the+-- savepoint is released.+--+-- Savepoints can be nested.+--+-- @+-- 'withTransaction' pool $ \\tx -> do+-- execute (txConn tx) stmt1 params1+-- result <- 'withSavepoint' tx $ \\_ -> do+-- execute (txConn tx) riskyStmt params2+-- -- If riskyStmt threw, we're still in the transaction+-- execute (txConn tx) stmt3 params3+-- @+withSavepoint :: Transaction -> (Transaction -> IO a) -> IO a+withSavepoint tx action = do+ name <- freshSavepointName+ let conn = txConn tx+ _ <- simpleQuery conn ("SAVEPOINT " <> name)+ mask $ \restore -> do+ result <- restore (action tx) `onException` do+ _ <- simpleQuery conn ("ROLLBACK TO SAVEPOINT " <> name)+ pure ()+ _ <- simpleQuery conn ("RELEASE SAVEPOINT " <> name)+ pure result++-- Internal ------------------------------------------------------------------++rollback :: Connection -> IO ()+rollback conn =+ void (simpleQuery conn "ROLLBACK") `catch` \(_ :: SomeException) -> pure ()++beginStatement :: IsolationLevel -> ByteString+beginStatement = \case+ ReadCommitted -> "BEGIN"+ RepeatableRead -> "BEGIN ISOLATION LEVEL REPEATABLE READ"+ Serializable -> "BEGIN ISOLATION LEVEL SERIALIZABLE"++beginModeStatement :: TransactionMode -> ByteString+beginModeStatement (TransactionMode iso ro def) =+ "BEGIN" <> isoClause <> roClause <> defClause+ where+ isoClause = case iso of+ ReadCommitted -> ""+ RepeatableRead -> " ISOLATION LEVEL REPEATABLE READ"+ Serializable -> " ISOLATION LEVEL SERIALIZABLE"+ roClause = if ro then " READ ONLY" else ""+ defClause = if def then " DEFERRABLE" else ""++-- | Global counter for unique savepoint names.+{-# NOINLINE savepointCounter #-}+savepointCounter :: IORef Word64+savepointCounter = unsafePerformIO (newIORef 0)++freshSavepointName :: IO ByteString+freshSavepointName = do+ n <- atomicModifyIORef' savepointCounter (\n -> (n + 1, n))+ pure ("valiant_sp_" <> BS8.pack (show n))
+ test/Main.hs view
@@ -0,0 +1,44 @@+module Main where++import Valiant.ErrorSpec qualified as ErrorSpec+import Valiant.Binary.ArraySpec qualified as ArraySpec+import Valiant.Binary.CompositeSpec qualified as CompositeSpec+import Valiant.Binary.DecodeSpec qualified as DecodeSpec+import Valiant.Binary.EncodeSpec qualified as EncodeSpec+import Valiant.Binary.HStoreSpec qualified as HStoreSpec+import Valiant.Binary.InetSpec qualified as InetSpec+import Valiant.Binary.MacAddrSpec qualified as MacAddrSpec+import Valiant.Binary.IntervalSpec qualified as IntervalSpec+import Valiant.Binary.JSONSpec qualified as JSONSpec+import Valiant.Binary.NewtypeSpec qualified as NewtypeSpec+import Valiant.Binary.PropertyHedgehogSpec qualified as PropertyHedgehogSpec+import Valiant.Binary.RangeSpec qualified as RangeSpec+import Valiant.Binary.RefineSpec qualified as RefineSpec+import Valiant.Binary.ScientificSpec qualified as ScientificSpec+import Valiant.Binary.UUIDSpec qualified as UUIDSpec+import Valiant.FromRowSpec qualified as FromRowSpec+import Valiant.NamedParamsSpec qualified as NamedParamsSpec+import Valiant.TupleSpec qualified as TupleSpec+import Test.Hspec++main :: IO ()+main = hspec $ do+ describe "Valiant.Binary.Encode" EncodeSpec.spec+ describe "Valiant.Binary.Decode" DecodeSpec.spec+ describe "Valiant.Binary.Array" ArraySpec.spec+ describe "Valiant.Binary.Composite" CompositeSpec.spec+ describe "Valiant.Binary.Range" RangeSpec.spec+ describe "Valiant.Binary.Refine" RefineSpec.spec+ describe "Valiant.Binary.Scientific" ScientificSpec.spec+ describe "Valiant.Binary.HStore" HStoreSpec.spec+ describe "Valiant.Binary.Inet" InetSpec.spec+ describe "Valiant.Binary.MacAddr" MacAddrSpec.spec+ describe "Valiant.Binary.Interval" IntervalSpec.spec+ describe "Valiant.Binary.UUID" UUIDSpec.spec+ describe "Valiant.Binary.JSON" JSONSpec.spec+ describe "Valiant.Binary.Newtype" NewtypeSpec.spec+ describe "Valiant.Binary.Property" PropertyHedgehogSpec.spec+ describe "Valiant.FromRow" FromRowSpec.spec+ describe "Valiant.Tuple" TupleSpec.spec+ describe "Valiant.NamedParams" NamedParamsSpec.spec+ describe "Valiant.Error" ErrorSpec.spec
+ test/Valiant/Binary/ArraySpec.hs view
@@ -0,0 +1,84 @@+module Valiant.Binary.ArraySpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int32, Int64)+import Data.Text (Text)+import Data.Vector (Vector)+import Data.Vector qualified as V+import Valiant.Binary.Array (pgDecodeArray, pgEncodeArray)+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import PgWire.Protocol.Oid+import Test.Hspec++spec :: Spec+spec = do+ describe "Int32 array round-trip" $ do+ it "empty array" $ do+ let arr = V.empty :: Vector Int32+ encoded = pgEncodeArray oidInt4 arr+ pgDecodeArray encoded `shouldBe` Right arr++ it "single element" $ do+ let arr = V.singleton (42 :: Int32)+ encoded = pgEncodeArray oidInt4 arr+ pgDecodeArray encoded `shouldBe` Right arr++ it "multiple elements" $ do+ let arr = V.fromList [1, 2, 3, 4, 5] :: Vector Int32+ encoded = pgEncodeArray oidInt4 arr+ pgDecodeArray encoded `shouldBe` Right arr++ it "negative values" $ do+ let arr = V.fromList [-100, 0, 100] :: Vector Int32+ encoded = pgEncodeArray oidInt4 arr+ pgDecodeArray encoded `shouldBe` Right arr++ describe "Int64 array round-trip" $ do+ it "multiple elements" $ do+ let arr = V.fromList [1000000, 2000000, 3000000] :: Vector Int64+ encoded = pgEncodeArray oidInt8 arr+ pgDecodeArray encoded `shouldBe` Right arr++ describe "Text array round-trip" $ do+ it "multiple strings" $ do+ let arr = V.fromList ["hello", "world", "foo"] :: Vector Text+ encoded = pgEncodeArray oidText arr+ pgDecodeArray encoded `shouldBe` Right arr++ it "empty strings" $ do+ let arr = V.fromList ["", "", "x"] :: Vector Text+ encoded = pgEncodeArray oidText arr+ pgDecodeArray encoded `shouldBe` Right arr++ describe "Bool array round-trip" $ do+ it "mixed booleans" $ do+ let arr = V.fromList [True, False, True, False] :: Vector Bool+ encoded = pgEncodeArray oidBool arr+ pgDecodeArray encoded `shouldBe` Right arr++ describe "decode errors" $ do+ it "fails on truncated input" $ do+ let result = pgDecodeArray (BS.pack [0, 0, 0, 1]) :: Either String (Vector Int32)+ result `shouldSatisfy` isLeft++ it "fails on empty input" $ do+ let result = pgDecodeArray BS.empty :: Either String (Vector Int32)+ result `shouldSatisfy` isLeft++ describe "binary format" $ do+ it "produces correct header for empty 1-D array" $ do+ let encoded = pgEncodeArray oidInt4 (V.empty :: Vector Int32)+ -- ndim=1, no null, oid=23, size=0, lbound=1+ BS.length encoded `shouldBe` 20+ -- First 4 bytes: ndim = 1+ BS.take 4 encoded `shouldBe` BS.pack [0, 0, 0, 1]++ it "element data follows header" $ do+ let encoded = pgEncodeArray oidInt4 (V.singleton (42 :: Int32))+ -- 20 bytes header + 4 bytes element length + 4 bytes element data = 28+ BS.length encoded `shouldBe` 28++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/CompositeSpec.hs view
@@ -0,0 +1,55 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+module Valiant.Binary.CompositeSpec (spec) where++import Data.ByteString qualified as BS+import Valiant.Binary.Composite+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import PgWire.Binary.Types (PgEncode (..))+import Data.Int (Int32)+import Test.Hspec++spec :: Spec+spec = do+ describe "composite round-trip" $ do+ it "empty composite (no fields)" $ do+ let fields = [] :: [CompositeField]+ pgDecodeComposite (pgEncodeComposite fields) `shouldBe` Right fields++ it "single non-null field" $ do+ let val = pgEncode (42 :: Int32)+ fields = [CompositeField 23 (Just val)]+ pgDecodeComposite (pgEncodeComposite fields) `shouldBe` Right fields++ it "single null field" $ do+ let fields = [CompositeField 23 Nothing]+ pgDecodeComposite (pgEncodeComposite fields) `shouldBe` Right fields++ it "multiple fields with mixed nulls" $ do+ let fields =+ [ CompositeField 23 (Just (pgEncode (1 :: Int32)))+ , CompositeField 25 Nothing+ , CompositeField 23 (Just (pgEncode (99 :: Int32)))+ ]+ pgDecodeComposite (pgEncodeComposite fields) `shouldBe` Right fields++ it "preserves field OIDs" $ do+ let fields =+ [ CompositeField 23 (Just (pgEncode (1 :: Int32)))+ , CompositeField 25 (Just (pgEncode ("hello" :: BS.ByteString)))+ ]+ Right decoded = pgDecodeComposite (pgEncodeComposite fields)+ map cfOid decoded `shouldBe` [23, 25]++ describe "composite decode errors" $ do+ it "fails on empty input" $ do+ let result = pgDecodeComposite BS.empty+ result `shouldSatisfy` isLeft++ it "fails on truncated input" $ do+ let result = pgDecodeComposite (BS.pack [0, 0, 0, 2, 0])+ result `shouldSatisfy` isLeft++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/DecodeSpec.hs view
@@ -0,0 +1,157 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+module Valiant.Binary.DecodeSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int16, Int32, Int64)+import Data.Text (Text)+import Data.Time+ ( Day+ , LocalTime (..)+ , TimeOfDay (..)+ , UTCTime (..)+ , fromGregorian+ )+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "Bool round-trip" $ do+ it "True" $ pgDecode (pgEncode True) `shouldBe` Right True+ it "False" $ pgDecode (pgEncode False) `shouldBe` Right False++ describe "Int16 round-trip" $ do+ it "0" $ pgDecode (pgEncode (0 :: Int16)) `shouldBe` Right (0 :: Int16)+ it "32767" $ pgDecode (pgEncode (32767 :: Int16)) `shouldBe` Right (32767 :: Int16)+ it "-1" $ pgDecode (pgEncode (-1 :: Int16)) `shouldBe` Right (-1 :: Int16)++ describe "Int32 round-trip" $ do+ it "0" $ pgDecode (pgEncode (0 :: Int32)) `shouldBe` Right (0 :: Int32)+ it "42" $ pgDecode (pgEncode (42 :: Int32)) `shouldBe` Right (42 :: Int32)+ it "-100" $ pgDecode (pgEncode (-100 :: Int32)) `shouldBe` Right (-100 :: Int32)++ describe "Int64 round-trip" $ do+ it "0" $ pgDecode (pgEncode (0 :: Int64)) `shouldBe` Right (0 :: Int64)+ it "large value" $ pgDecode (pgEncode (9999999999 :: Int64)) `shouldBe` Right (9999999999 :: Int64)++ describe "Float round-trip" $ do+ it "0.0" $ pgDecode (pgEncode (0.0 :: Float)) `shouldBe` Right (0.0 :: Float)+ it "3.14" $ do+ let Right v = pgDecode (pgEncode (3.14 :: Float)) :: Either String Float+ abs (v - 3.14) `shouldSatisfy` (< 0.001)++ describe "Double round-trip" $ do+ it "0.0" $ pgDecode (pgEncode (0.0 :: Double)) `shouldBe` Right (0.0 :: Double)+ it "3.14159" $ pgDecode (pgEncode (3.14159 :: Double)) `shouldBe` Right (3.14159 :: Double)++ describe "Text round-trip" $ do+ it "empty" $ pgDecode (pgEncode ("" :: Text)) `shouldBe` Right ("" :: Text)+ it "hello" $ pgDecode (pgEncode ("hello" :: Text)) `shouldBe` Right ("hello" :: Text)++ describe "ByteString round-trip" $ do+ it "identity" $ pgDecode (pgEncode ("bytes" :: BS.ByteString)) `shouldBe` Right ("bytes" :: BS.ByteString)++ describe "error on wrong size" $ do+ it "Int32 with 3 bytes fails" $ do+ let result = pgDecode (BS.pack [0, 0, 0]) :: Either String Int32+ result `shouldSatisfy` isLeft+ it "Bool with 2 bytes fails" $ do+ let result = pgDecode (BS.pack [0, 0]) :: Either String Bool+ result `shouldSatisfy` isLeft++ describe "UTCTime round-trip" $ do+ it "PG epoch (2000-01-01 00:00:00 UTC)" $ do+ let t = UTCTime (fromGregorian 2000 1 1) 0+ pgDecode (pgEncode t) `shouldBe` Right t++ it "one second after PG epoch" $ do+ let t = UTCTime (fromGregorian 2000 1 1) 1+ pgDecode (pgEncode t) `shouldBe` Right t++ it "a date in 2024" $ do+ -- 2024-06-15 12:00:00 UTC+ let t = UTCTime (fromGregorian 2024 6 15) (12 * 3600)+ pgDecode (pgEncode t) `shouldBe` Right t++ it "midnight on a non-epoch day" $ do+ let t = UTCTime (fromGregorian 2020 3 15) 0+ pgDecode (pgEncode t) `shouldBe` Right t++ it "time with sub-second precision (microseconds)" $ do+ -- 0.5 seconds = 500000 microseconds+ let t = UTCTime (fromGregorian 2000 1 1) 0.5+ -- PG stores microseconds, so round-trip should preserve microsecond precision+ pgDecode (pgEncode t) `shouldBe` Right t++ it "time before PG epoch (1999-12-31)" $ do+ let t = UTCTime (fromGregorian 1999 12 31) 0+ pgDecode (pgEncode t) `shouldBe` Right t++ describe "Day round-trip" $ do+ it "PG epoch day (2000-01-01)" $ do+ let d = fromGregorian 2000 1 1+ pgDecode (pgEncode d) `shouldBe` Right d++ it "2000-01-02" $ do+ let d = fromGregorian 2000 1 2+ pgDecode (pgEncode d) `shouldBe` Right d++ it "1999-12-31 (before PG epoch)" $ do+ let d = fromGregorian 1999 12 31+ pgDecode (pgEncode d) `shouldBe` Right d++ it "2024-06-15" $ do+ let d = fromGregorian 2024 6 15+ pgDecode (pgEncode d) `shouldBe` Right d++ it "fails on wrong byte count" $ do+ let result = pgDecode (BS.pack [0, 0]) :: Either String Day+ result `shouldSatisfy` isLeft++ describe "TimeOfDay round-trip" $ do+ it "midnight (00:00:00)" $ do+ let tod = TimeOfDay 0 0 0+ pgDecode (pgEncode tod) `shouldBe` Right tod++ it "00:00:01" $ do+ let tod = TimeOfDay 0 0 1+ pgDecode (pgEncode tod) `shouldBe` Right tod++ it "12:30:45" $ do+ let tod = TimeOfDay 12 30 45+ pgDecode (pgEncode tod) `shouldBe` Right tod++ it "23:59:59" $ do+ let tod = TimeOfDay 23 59 59+ pgDecode (pgEncode tod) `shouldBe` Right tod++ it "fails on wrong byte count" $ do+ let result = pgDecode (BS.pack [0, 0, 0, 0]) :: Either String TimeOfDay+ result `shouldSatisfy` isLeft++ describe "LocalTime round-trip" $ do+ it "PG epoch midnight" $ do+ let lt = LocalTime (fromGregorian 2000 1 1) (TimeOfDay 0 0 0)+ pgDecode (pgEncode lt) `shouldBe` Right lt++ it "PG epoch + 1 second" $ do+ let lt = LocalTime (fromGregorian 2000 1 1) (TimeOfDay 0 0 1)+ pgDecode (pgEncode lt) `shouldBe` Right lt++ it "2024-03-15 14:30:00" $ do+ let lt = LocalTime (fromGregorian 2024 3 15) (TimeOfDay 14 30 0)+ pgDecode (pgEncode lt) `shouldBe` Right lt++ it "date before PG epoch" $ do+ let lt = LocalTime (fromGregorian 1999 6 15) (TimeOfDay 8 0 0)+ pgDecode (pgEncode lt) `shouldBe` Right lt++ it "fails on wrong byte count" $ do+ let result = pgDecode (BS.pack [0, 0, 0, 0]) :: Either String LocalTime+ result `shouldSatisfy` isLeft++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/EncodeSpec.hs view
@@ -0,0 +1,120 @@+module Valiant.Binary.EncodeSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int16, Int32, Int64)+import Data.Text (Text)+import Data.Time+ ( LocalTime (..)+ , TimeOfDay (..)+ , UTCTime (..)+ , fromGregorian+ , secondsToDiffTime+ )+import Valiant.Binary.Encode ()+import PgWire.Binary.Types (PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "Bool" $ do+ it "encodes True as 0x01" $ do+ pgEncode True `shouldBe` BS.singleton 1+ it "encodes False as 0x00" $ do+ pgEncode False `shouldBe` BS.singleton 0++ describe "Int16" $ do+ it "encodes 0" $ do+ pgEncode (0 :: Int16) `shouldBe` BS.pack [0, 0]+ it "encodes 256 as big-endian" $ do+ pgEncode (256 :: Int16) `shouldBe` BS.pack [1, 0]+ it "encodes -1" $ do+ pgEncode (-1 :: Int16) `shouldBe` BS.pack [0xFF, 0xFF]++ describe "Int32" $ do+ it "encodes 1 as big-endian" $ do+ pgEncode (1 :: Int32) `shouldBe` BS.pack [0, 0, 0, 1]+ it "encodes 16777216 (0x01000000)" $ do+ pgEncode (16777216 :: Int32) `shouldBe` BS.pack [1, 0, 0, 0]++ describe "Int64" $ do+ it "encodes 1 as 8 bytes big-endian" $ do+ pgEncode (1 :: Int64) `shouldBe` BS.pack [0, 0, 0, 0, 0, 0, 0, 1]++ describe "Text" $ do+ it "encodes as UTF-8" $ do+ pgEncode ("hello" :: Text) `shouldBe` "hello"++ describe "ByteString" $ do+ it "encodes as identity" $ do+ pgEncode ("raw bytes" :: BS.ByteString) `shouldBe` "raw bytes"++ describe "UTCTime" $ do+ it "encodes PG epoch (2000-01-01 00:00:00 UTC) as zero" $ do+ let pgEpoch = UTCTime (fromGregorian 2000 1 1) 0+ bs = pgEncode pgEpoch+ -- PG epoch should encode to 0 microseconds+ BS.length bs `shouldBe` 8+ bs `shouldBe` BS.pack [0, 0, 0, 0, 0, 0, 0, 0]++ it "encodes a time after PG epoch as positive" $ do+ -- 2000-01-01 00:00:01 UTC = 1 second = 1000000 microseconds after PG epoch+ let t = UTCTime (fromGregorian 2000 1 1) 1+ bs = pgEncode t+ BS.length bs `shouldBe` 8+ -- 1000000 = 0x000000000F4240+ bs `shouldBe` BS.pack [0, 0, 0, 0, 0, 0x0F, 0x42, 0x40]++ it "encodes a known date correctly (2024-01-15 12:30:00 UTC)" $ do+ let t = UTCTime (fromGregorian 2024 1 15) (secondsToDiffTime (12 * 3600 + 30 * 60))+ bs = pgEncode t+ BS.length bs `shouldBe` 8++ describe "Day" $ do+ it "encodes PG epoch day (2000-01-01) as zero" $ do+ let d = fromGregorian 2000 1 1+ pgEncode d `shouldBe` BS.pack [0, 0, 0, 0]++ it "encodes 2000-01-02 as 1" $ do+ let d = fromGregorian 2000 1 2+ pgEncode d `shouldBe` BS.pack [0, 0, 0, 1]++ it "encodes a date before PG epoch as negative" $ do+ -- 1999-12-31 is -1 days from PG epoch+ let d = fromGregorian 1999 12 31+ pgEncode d `shouldBe` BS.pack [0xFF, 0xFF, 0xFF, 0xFF]++ it "produces 4 bytes" $ do+ let d = fromGregorian 2024 6 15+ BS.length (pgEncode d) `shouldBe` 4++ describe "TimeOfDay" $ do+ it "encodes midnight as zero" $ do+ let tod = TimeOfDay 0 0 0+ pgEncode tod `shouldBe` BS.pack [0, 0, 0, 0, 0, 0, 0, 0]++ it "encodes 00:00:01 as 1000000 microseconds" $ do+ let tod = TimeOfDay 0 0 1+ pgEncode tod `shouldBe` BS.pack [0, 0, 0, 0, 0, 0x0F, 0x42, 0x40]++ it "encodes 01:00:00 as 3600000000 microseconds" $ do+ let tod = TimeOfDay 1 0 0+ bs = pgEncode tod+ BS.length bs `shouldBe` 8++ it "produces 8 bytes" $ do+ let tod = TimeOfDay 12 30 45+ BS.length (pgEncode tod) `shouldBe` 8++ describe "LocalTime" $ do+ it "encodes PG epoch midnight as zero" $ do+ let lt = LocalTime (fromGregorian 2000 1 1) (TimeOfDay 0 0 0)+ pgEncode lt `shouldBe` BS.pack [0, 0, 0, 0, 0, 0, 0, 0]++ it "encodes PG epoch + 1 second" $ do+ -- 1 second = 1000000 microseconds+ let lt = LocalTime (fromGregorian 2000 1 1) (TimeOfDay 0 0 1)+ pgEncode lt `shouldBe` BS.pack [0, 0, 0, 0, 0, 0x0F, 0x42, 0x40]++ it "produces 8 bytes" $ do+ let lt = LocalTime (fromGregorian 2024 3 15) (TimeOfDay 14 30 0)+ BS.length (pgEncode lt) `shouldBe` 8
+ test/Valiant/Binary/HStoreSpec.hs view
@@ -0,0 +1,46 @@+module Valiant.Binary.HStoreSpec (spec) where++import Data.Map.Strict qualified as Map+import Valiant.Binary.HStore (PgHStore (..), hstoreFromList, hstoreToList)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "hstore round-trip" $ do+ it "empty map" $ do+ let v = PgHStore Map.empty+ pgDecode (pgEncode v) `shouldBe` Right v++ it "single pair" $ do+ let v = hstoreFromList [("key", "value")]+ pgDecode (pgEncode v) `shouldBe` Right v++ it "multiple pairs" $ do+ let v = hstoreFromList [("a", "1"), ("b", "2"), ("c", "3")]+ pgDecode (pgEncode v) `shouldBe` Right v++ it "preserves NULL values" $ do+ let v = PgHStore (Map.fromList [("present", Just "yes"), ("absent", Nothing)])+ pgDecode (pgEncode v) `shouldBe` Right v++ it "empty key" $ do+ let v = hstoreFromList [("", "empty-key")]+ pgDecode (pgEncode v) `shouldBe` Right v++ it "empty value" $ do+ let v = hstoreFromList [("key", "")]+ pgDecode (pgEncode v) `shouldBe` Right v++ it "unicode" $ do+ let v = hstoreFromList [("名前", "太郎"), ("stadt", "München")]+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "hstoreToList" $ do+ it "filters out NULL values" $ do+ let v = PgHStore (Map.fromList [("a", Just "1"), ("b", Nothing), ("c", Just "3")])+ hstoreToList v `shouldBe` [("a", "1"), ("c", "3")]++ it "empty on all NULLs" $ do+ let v = PgHStore (Map.fromList [("x", Nothing)])+ hstoreToList v `shouldBe` []
+ test/Valiant/Binary/InetSpec.hs view
@@ -0,0 +1,87 @@+module Valiant.Binary.InetSpec (spec) where++import Data.ByteString qualified as BS+import Valiant.Binary.Inet (PgInet (..), ipv4, ipv4Host, ipv6, ipv6Host, inetToText)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "inet IPv4 round-trip" $ do+ it "loopback /8" $ do+ let v = ipv4 127 0 0 1 8+ pgDecode (pgEncode v) `shouldBe` Right v++ it "private network /24" $ do+ let v = ipv4 192 168 1 0 24+ pgDecode (pgEncode v) `shouldBe` Right v++ it "host /32" $ do+ let v = ipv4Host 10 0 0 1+ pgDecode (pgEncode v) `shouldBe` Right v++ it "default route /0" $ do+ let v = ipv4 0 0 0 0 0+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "inet IPv6 round-trip" $ do+ it "loopback ::1/128" $ do+ let addr = BS.pack [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1]+ v = ipv6Host addr+ pgDecode (pgEncode v) `shouldBe` Right v++ it "network /64" $ do+ let addr = BS.pack [0x20,0x01,0x0d,0xb8,0,0,0,0,0,0,0,0,0,0,0,0]+ v = ipv6 addr 64+ pgDecode (pgEncode v) `shouldBe` Right v++ it "all zeros ::/0" $ do+ let addr = BS.replicate 16 0+ v = ipv6 addr 0+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "inet encode size" $ do+ it "IPv4 produces 8 bytes (4 header + 4 address)" $ do+ BS.length (pgEncode (ipv4 192 168 1 0 24)) `shouldBe` 8++ it "IPv6 produces 20 bytes (4 header + 16 address)" $ do+ let addr = BS.replicate 16 0+ BS.length (pgEncode (ipv6 addr 64)) `shouldBe` 20++ describe "inet decode errors" $ do+ it "fails on too-short input" $ do+ let result = pgDecode (BS.pack [2, 24, 0]) :: Either String PgInet+ result `shouldSatisfy` isLeft++ it "fails on empty input" $ do+ let result = pgDecode BS.empty :: Either String PgInet+ result `shouldSatisfy` isLeft++ it "fails on address family mismatch" $ do+ -- Claim IPv6 family (3) but only 4 address bytes+ let result = pgDecode (BS.pack [3, 32, 0, 4, 10, 0, 0, 1]) :: Either String PgInet+ result `shouldSatisfy` isLeft++ describe "inet cidr flag" $ do+ it "preserves cidr flag on round-trip" $ do+ let v = (ipv4 192 168 1 0 24) { inetIsCidr = True }+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "inetToText" $ do+ it "renders IPv4 in CIDR notation" $+ inetToText (ipv4 192 168 1 0 24) `shouldBe` "192.168.1.0/24"++ it "renders IPv4 host" $+ inetToText (ipv4Host 10 0 0 1) `shouldBe` "10.0.0.1/32"++ it "renders IPv6 loopback" $ do+ let addr = BS.pack [0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1]+ inetToText (ipv6Host addr) `shouldBe` "::1/128"++ it "renders IPv6 all zeros" $ do+ let addr = BS.replicate 16 0+ inetToText (ipv6 addr 0) `shouldBe` "::/0"++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/IntervalSpec.hs view
@@ -0,0 +1,54 @@+module Valiant.Binary.IntervalSpec (spec) where++import Data.ByteString qualified as BS+import Valiant.Binary.Interval (PgInterval (..))+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "interval round-trip" $ do+ it "zero interval" $ do+ let iv = PgInterval 0 0 0+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "1 second" $ do+ let iv = PgInterval 1000000 0 0+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "1 day" $ do+ let iv = PgInterval 0 1 0+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "1 month" $ do+ let iv = PgInterval 0 0 1+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "complex interval (1 year, 2 months, 3 days, 4 hours)" $ do+ let iv = PgInterval (4 * 3600 * 1000000) 3 14 -- 14 months = 1y2m+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "negative microseconds" $ do+ let iv = PgInterval (-500000) 0 0+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "negative days" $ do+ let iv = PgInterval 0 (-7) 0+ pgDecode (pgEncode iv) `shouldBe` Right iv++ it "negative months" $ do+ let iv = PgInterval 0 0 (-3)+ pgDecode (pgEncode iv) `shouldBe` Right iv++ describe "interval binary format" $ do+ it "produces exactly 16 bytes" $ do+ let iv = PgInterval 12345 67 8+ BS.length (pgEncode iv) `shouldBe` 16++ it "fails on wrong byte count" $ do+ let result = pgDecode (BS.pack [0, 0, 0, 0]) :: Either String PgInterval+ result `shouldSatisfy` isLeft++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/JSONSpec.hs view
@@ -0,0 +1,61 @@+module Valiant.Binary.JSONSpec (spec) where++import Data.Aeson (Value (..), object, (.=))+import Data.ByteString qualified as BS+import Valiant.Binary.JSON ()+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "JSON round-trip (json format)" $ do+ it "null" $ do+ pgDecode (pgEncode Null) `shouldBe` Right Null++ it "boolean true" $ do+ pgDecode (pgEncode (Bool True)) `shouldBe` Right (Bool True)++ it "number" $ do+ pgDecode (pgEncode (Number 42)) `shouldBe` Right (Number 42)++ it "string" $ do+ pgDecode (pgEncode (String "hello")) `shouldBe` Right (String "hello")++ it "empty object" $ do+ pgDecode (pgEncode (object [])) `shouldBe` Right (object [])++ it "object with fields" $ do+ let v = object ["name" .= String "Alice", "age" .= Number 30]+ pgDecode (pgEncode v) `shouldBe` Right v++ it "array" $ do+ let v = Array mempty+ pgDecode (pgEncode v) `shouldBe` Right v++ it "nested structure" $ do+ let v = object ["users" .= [object ["id" .= Number 1, "name" .= String "Bob"]]]+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "JSONB decode (version byte prefix)" $ do+ it "decodes jsonb with version byte 0x01" $ do+ let jsonBytes = pgEncode (String "test")+ jsonbBytes = BS.singleton 1 <> jsonBytes+ pgDecode jsonbBytes `shouldBe` Right (String "test")++ it "decodes jsonb object with version byte" $ do+ let jsonBytes = pgEncode (object ["k" .= String "v"])+ jsonbBytes = BS.singleton 1 <> jsonBytes+ pgDecode jsonbBytes `shouldBe` Right (object ["k" .= String "v"])++ describe "JSON decode errors" $ do+ it "fails on invalid JSON" $ do+ let result = pgDecode "not json {{{" :: Either String Value+ result `shouldSatisfy` isLeft++ it "fails on empty input" $ do+ let result = pgDecode BS.empty :: Either String Value+ result `shouldSatisfy` isLeft++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/MacAddrSpec.hs view
@@ -0,0 +1,53 @@+module Valiant.Binary.MacAddrSpec (spec) where++import Data.ByteString qualified as BS+import Valiant.Binary.MacAddr (PgMacAddr (..), macAddr, macAddrToText)+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "macaddr round-trip" $ do+ it "standard 6-byte address" $ do+ let v = macAddr 0x08 0x00 0x2b 0x01 0x02 0x03+ pgDecode (pgEncode v) `shouldBe` Right v++ it "all zeros" $ do+ let v = macAddr 0 0 0 0 0 0+ pgDecode (pgEncode v) `shouldBe` Right v++ it "all ones" $ do+ let v = macAddr 0xff 0xff 0xff 0xff 0xff 0xff+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "macaddr8 round-trip" $ do+ it "8-byte address" $ do+ let v = PgMacAddr (BS.pack [0x08, 0x00, 0x2b, 0xfe, 0xff, 0x01, 0x02, 0x03])+ pgDecode (pgEncode v) `shouldBe` Right v++ describe "macaddr encode size" $ do+ it "6-byte produces 6 bytes" $+ BS.length (pgEncode (macAddr 1 2 3 4 5 6)) `shouldBe` 6++ it "8-byte produces 8 bytes" $+ BS.length (pgEncode (PgMacAddr (BS.replicate 8 0))) `shouldBe` 8++ describe "macaddr decode errors" $ do+ it "fails on 5 bytes" $ do+ let result = pgDecode (BS.pack [1, 2, 3, 4, 5]) :: Either String PgMacAddr+ result `shouldSatisfy` isLeft++ it "fails on empty" $ do+ let result = pgDecode BS.empty :: Either String PgMacAddr+ result `shouldSatisfy` isLeft++ describe "macAddrToText" $ do+ it "renders with colons" $+ macAddrToText (macAddr 0x08 0x00 0x2b 0x01 0x02 0x03) `shouldBe` "08:00:2b:01:02:03"++ it "pads single hex digits" $+ macAddrToText (macAddr 0x01 0x02 0x03 0x04 0x05 0x06) `shouldBe` "01:02:03:04:05:06"++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/NewtypeSpec.hs view
@@ -0,0 +1,131 @@+{-# LANGUAGE DerivingVia #-}++module Valiant.Binary.NewtypeSpec (spec) where++import Data.Int (Int32, Int64)+import Data.Proxy (Proxy (..))+import Data.Text (Text)+import Data.Vector qualified as V+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import Valiant.FromRow (DecodeColumn (..), FromRow (..))+import Valiant.ToParams (EncodeField (..), ToParams (..))+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid (oidInt4, oidInt8, oidText)+import Test.Hspec++-- | A newtype that should inherit PgEncode/PgDecode via GeneralizedNewtypeDeriving.+newtype UserId = UserId Int32+ deriving stock (Show, Eq)+ deriving newtype (PgEncode, PgDecode)++-- | Another newtype wrapping Text.+newtype Email = Email Text+ deriving stock (Show, Eq)+ deriving newtype (PgEncode, PgDecode)++-- | A newtype wrapping Int64.+newtype PostId = PostId Int64+ deriving stock (Show, Eq)+ deriving newtype (PgEncode, PgDecode)++-- | A newtype using DerivingVia instead of GeneralizedNewtypeDeriving.+newtype OrgId = OrgId Int32+ deriving stock (Show, Eq)+ deriving (PgEncode, PgDecode) via Int32++spec :: Spec+spec = do+ describe "GeneralizedNewtypeDeriving for PgEncode" $ do+ it "UserId encodes identically to Int32" $ do+ pgEncode (UserId 42) `shouldBe` pgEncode (42 :: Int32)++ it "UserId has the same OID as Int32" $ do+ pgOid (Proxy @UserId) `shouldBe` oidInt4++ it "Email encodes identically to Text" $ do+ pgEncode (Email "alice@example.com") `shouldBe` pgEncode ("alice@example.com" :: Text)++ it "Email has the same OID as Text" $ do+ pgOid (Proxy @Email) `shouldBe` oidText++ it "PostId encodes identically to Int64" $ do+ pgEncode (PostId 999) `shouldBe` pgEncode (999 :: Int64)++ it "PostId has the same OID as Int64" $ do+ pgOid (Proxy @PostId) `shouldBe` oidInt8++ describe "GeneralizedNewtypeDeriving for PgDecode" $ do+ it "decodes UserId from Int32 bytes" $ do+ let bs = pgEncode (42 :: Int32)+ pgDecode bs `shouldBe` Right (UserId 42)++ it "decodes Email from Text bytes" $ do+ let bs = pgEncode ("test@example.com" :: Text)+ pgDecode bs `shouldBe` Right (Email "test@example.com")++ it "decodes PostId from Int64 bytes" $ do+ let bs = pgEncode (123 :: Int64)+ pgDecode bs `shouldBe` Right (PostId 123)++ describe "DerivingVia for PgEncode/PgDecode" $ do+ it "OrgId encodes identically to Int32" $ do+ pgEncode (OrgId 7) `shouldBe` pgEncode (7 :: Int32)++ it "OrgId has the same OID as Int32" $ do+ pgOid (Proxy @OrgId) `shouldBe` oidInt4++ it "OrgId round-trips" $ do+ let bs = pgEncode (OrgId 7)+ pgDecode bs `shouldBe` Right (OrgId 7)++ describe "Newtype round-trip" $ do+ it "UserId encode/decode is identity" $ do+ let uid = UserId 12345+ pgDecode (pgEncode uid) `shouldBe` Right uid++ it "Email encode/decode is identity" $ do+ let email = Email "hello@world.com"+ pgDecode (pgEncode email) `shouldBe` Right email++ describe "Newtypes work with EncodeField" $ do+ it "EncodeField encodes UserId" $ do+ encodeField (UserId 42) `shouldBe` Just (pgEncode (42 :: Int32))++ it "EncodeField encodes Maybe UserId (Just)" $ do+ encodeField (Just (UserId 42)) `shouldBe` Just (pgEncode (42 :: Int32))++ it "EncodeField encodes Maybe UserId (Nothing)" $ do+ encodeField (Nothing @UserId) `shouldBe` Nothing++ describe "Newtypes work with ToParams" $ do+ it "single UserId as ToParams" $ do+ let params = toParams (UserId 42)+ V.length params `shouldBe` 1+ V.head params `shouldBe` Just (pgEncode (42 :: Int32))++ it "tuple of newtypes as ToParams" $ do+ let params = toParams (UserId 1, Email "a@b.com")+ V.length params `shouldBe` 2++ describe "Newtypes work with DecodeColumn" $ do+ it "decodes UserId from a row" $ do+ let row = V.singleton (Just (pgEncode (42 :: Int32)))+ decodeColumn row 0 `shouldBe` Right (UserId 42)++ it "decodes Maybe UserId from NULL" $ do+ let row = V.singleton Nothing+ decodeColumn row 0 `shouldBe` Right (Nothing @UserId)++ it "decodes Maybe UserId from non-NULL" $ do+ let row = V.singleton (Just (pgEncode (42 :: Int32)))+ decodeColumn row 0 `shouldBe` Right (Just (UserId 42))++ describe "Newtypes work with FromRow" $ do+ it "decodes single UserId via FromRow" $ do+ let row = V.singleton (Just (pgEncode (42 :: Int32)))+ fromRow row `shouldBe` Right (UserId 42)++ it "decodes tuple with newtypes via FromRow" $ do+ let row = V.fromList [Just (pgEncode (1 :: Int32)), Just (pgEncode ("alice" :: Text))]+ fromRow row `shouldBe` Right (UserId 1, Email "alice")
+ test/Valiant/Binary/PropertyHedgehogSpec.hs view
@@ -0,0 +1,278 @@+module Valiant.Binary.PropertyHedgehogSpec (spec) where++import Data.Aeson (Value (..), object, (.=))+import Data.Aeson.Key qualified as Key+import Data.ByteString qualified as BS+import Data.Int (Int32)+import Data.Scientific (Scientific, scientific)+import Data.Text qualified as T+import Data.Time+ ( Day+ , LocalTime (..)+ , TimeOfDay (..)+ , UTCTime (..)+ , fromGregorian+ , picosecondsToDiffTime+ )+import Data.UUID.Types (UUID)+import Data.UUID.Types qualified as UUID+import Data.Vector qualified as V+import Hedgehog+import Hedgehog.Gen qualified as Gen+import Hedgehog.Range qualified as Range+import Valiant.Binary.Array (pgDecodeArray, pgEncodeArray)+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import Valiant.Binary.Interval (PgInterval (..))+import Valiant.Binary.JSON ()+import Valiant.Binary.Range (PgRange (..), RangeBound (..), pgDecodeRange, pgEncodeRange)+import Valiant.Binary.Scientific ()+import Valiant.Binary.UUID ()+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import PgWire.Protocol.Oid+import Test.Hspec+import Test.Hspec.Hedgehog (hedgehog)++spec :: Spec+spec = do+ -- ── Scalar round-trips ──────────────────────────────────────────+ describe "Bool round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ b <- forAll Gen.bool+ pgDecode (pgEncode b) === Right b++ describe "Int16 round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ n <- forAll (Gen.int16 Range.linearBounded)+ pgDecode (pgEncode n) === Right n++ describe "Int32 round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ n <- forAll (Gen.int32 Range.linearBounded)+ pgDecode (pgEncode n) === Right n++ describe "Int64 round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ n <- forAll (Gen.int64 Range.linearBounded)+ pgDecode (pgEncode n) === Right n++ describe "Float round-trip (hedgehog)" $+ it "encode/decode is identity (non-NaN)" $ hedgehog $ do+ n <- forAll (Gen.filter (not . isNaN) (Gen.float (Range.linearFrac (-1e6) 1e6)))+ pgDecode (pgEncode n) === Right n++ describe "Double round-trip (hedgehog)" $+ it "encode/decode is identity (non-NaN)" $ hedgehog $ do+ n <- forAll (Gen.filter (not . isNaN) (Gen.double (Range.linearFrac (-1e15) 1e15)))+ pgDecode (pgEncode n) === Right n++ describe "Text round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ t <- forAll (Gen.text (Range.linear 0 200) Gen.unicode)+ pgDecode (pgEncode t) === Right t++ describe "ByteString round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ bs <- forAll (Gen.bytes (Range.linear 0 200))+ pgDecode (pgEncode bs) === Right bs++ -- ── Time types ──────────────────────────────────────────────────+ describe "Day round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ d <- forAll genDay+ pgDecode (pgEncode d) === Right d++ describe "TimeOfDay round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ tod <- forAll genTimeOfDay+ pgDecode (pgEncode tod) === Right tod++ describe "UTCTime round-trip (hedgehog)" $+ it "encode/decode is identity (microsecond precision)" $ hedgehog $ do+ t <- forAll genUTCTime+ pgDecode (pgEncode t) === Right t++ describe "LocalTime round-trip (hedgehog)" $+ it "encode/decode is identity (microsecond precision)" $ hedgehog $ do+ lt <- forAll genLocalTime+ pgDecode (pgEncode lt) === Right lt++ describe "PgInterval round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ iv <- forAll genInterval+ pgDecode (pgEncode iv) === Right iv++ -- ── Numeric types ───────────────────────────────────────────────+ describe "Scientific round-trip (hedgehog)" $+ it "encode/decode is identity (reasonable range)" $ hedgehog $ do+ s <- forAll genScientific+ pgDecode (pgEncode s) === Right s++ -- ── Array types ─────────────────────────────────────────────────+ describe "Int32 array round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ xs <- forAll (Gen.list (Range.linear 0 50) (Gen.int32 Range.linearBounded))+ let v = V.fromList xs+ pgDecodeArray (pgEncodeArray oidInt4 v) === Right v++ describe "Text array round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ xs <- forAll (Gen.list (Range.linear 0 50) (Gen.text (Range.linear 0 20) Gen.alphaNum))+ let v = V.fromList xs+ pgDecodeArray (pgEncodeArray oidText v) === Right v++ describe "Bool array round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ xs <- forAll (Gen.list (Range.linear 0 50) Gen.bool)+ let v = V.fromList xs+ pgDecodeArray (pgEncodeArray oidBool v) === Right v++ describe "Int64 array round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ xs <- forAll (Gen.list (Range.linear 0 50) (Gen.int64 Range.linearBounded))+ let v = V.fromList xs+ pgDecodeArray (pgEncodeArray oidInt8 v) === Right v++ describe "Double array round-trip (hedgehog)" $+ it "encode/decode is identity (non-NaN)" $ hedgehog $ do+ xs <- forAll (Gen.list (Range.linear 0 50)+ (Gen.filter (not . isNaN) (Gen.double (Range.linearFrac (-1e15) 1e15))))+ let v = V.fromList xs+ pgDecodeArray (pgEncodeArray oidFloat8 v) === Right v++ -- ── UUID ────────────────────────────────────────────────────────+ describe "UUID round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ u <- forAll genUUID+ pgDecode (pgEncode u) === Right u++ -- ── Range ───────────────────────────────────────────────────────+ describe "PgRange Int32 round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ r <- forAll genInt32Range+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) === Right r++ -- ── JSON ────────────────────────────────────────────────────────+ describe "JSON Value round-trip (hedgehog)" $+ it "encode/decode is identity" $ hedgehog $ do+ v <- forAll genJsonValue+ pgDecode (pgEncode v) === Right v++ -- ── Encode sizes ────────────────────────────────────────────────+ describe "Encode output size (hedgehog)" $ do+ it "Bool encodes to 1 byte" $ hedgehog $ do+ b <- forAll Gen.bool+ BS.length (pgEncode b) === 1++ it "Int16 encodes to 2 bytes" $ hedgehog $ do+ n <- forAll (Gen.int16 Range.linearBounded)+ BS.length (pgEncode n) === 2++ it "Int32 encodes to 4 bytes" $ hedgehog $ do+ n <- forAll (Gen.int32 Range.linearBounded)+ BS.length (pgEncode n) === 4++ it "Int64 encodes to 8 bytes" $ hedgehog $ do+ n <- forAll (Gen.int64 Range.linearBounded)+ BS.length (pgEncode n) === 8++ it "Day encodes to 4 bytes" $ hedgehog $ do+ d <- forAll genDay+ BS.length (pgEncode d) === 4++ it "TimeOfDay encodes to 8 bytes" $ hedgehog $ do+ tod <- forAll genTimeOfDay+ BS.length (pgEncode tod) === 8++ it "PgInterval encodes to 16 bytes" $ hedgehog $ do+ iv <- forAll genInterval+ BS.length (pgEncode iv) === 16++ it "UUID encodes to 16 bytes" $ hedgehog $ do+ u <- forAll genUUID+ BS.length (pgEncode u) === 16++ it "Empty range encodes to 1 byte" $+ BS.length (pgEncodeRange (pgEncode @Int32) EmptyRange) `shouldBe` 1++-- ── Generators ──────────────────────────────────────────────────────++genDay :: Gen Day+genDay = do+ y <- Gen.integral (Range.linear 1900 2100 :: Range Integer)+ m <- Gen.int (Range.linear 1 12)+ d <- Gen.int (Range.linear 1 28)+ pure (fromGregorian (fromIntegral y) m d)++genTimeOfDay :: Gen TimeOfDay+genTimeOfDay = do+ h <- Gen.int (Range.linear 0 23)+ m <- Gen.int (Range.linear 0 59)+ micros <- Gen.int (Range.linear 0 59999999)+ pure (TimeOfDay h m (fromIntegral micros / 1000000))++genUTCTime :: Gen UTCTime+genUTCTime = do+ d <- genDay+ micros <- Gen.int64 (Range.linear 0 86399999999)+ let picos = fromIntegral micros * 1000000+ pure (UTCTime d (picosecondsToDiffTime picos))++genLocalTime :: Gen LocalTime+genLocalTime = LocalTime <$> genDay <*> genTimeOfDay++genInterval :: Gen PgInterval+genInterval =+ PgInterval+ <$> Gen.int64 Range.linearBounded+ <*> Gen.int32 Range.linearBounded+ <*> Gen.int32 Range.linearBounded++genScientific :: Gen Scientific+genScientific = do+ coeff <- Gen.integral (Range.linear (-999999999) 999999999)+ expo <- Gen.int (Range.linear (-8) 8)+ pure (scientific coeff expo)++genUUID :: Gen UUID+genUUID = do+ w1 <- Gen.word32 Range.linearBounded+ w2 <- Gen.word32 Range.linearBounded+ w3 <- Gen.word32 Range.linearBounded+ w4 <- Gen.word32 Range.linearBounded+ pure (UUID.fromWords w1 w2 w3 w4)++genRangeBound :: Gen a -> Gen (Maybe (RangeBound a))+genRangeBound gen =+ Gen.choice+ [ pure Nothing+ , Just . Inclusive <$> gen+ , Just . Exclusive <$> gen+ ]++genInt32Range :: Gen (PgRange Int32)+genInt32Range =+ Gen.frequency+ [ (1, pure EmptyRange)+ , (9, PgRange <$> genRangeBound (Gen.int32 Range.linearBounded) <*> genRangeBound (Gen.int32 Range.linearBounded))+ ]++genJsonValue :: Gen Value+genJsonValue = Gen.recursive Gen.choice+ -- base cases+ [ pure Null+ , Bool <$> Gen.bool+ , Number . fromIntegral <$> Gen.int (Range.linear (-1000) 1000)+ , String . T.pack <$> Gen.list (Range.linear 0 20) (Gen.element (['a'..'z'] ++ ['0'..'9'] ++ " "))+ ]+ -- recursive cases+ [ Array . V.fromList <$> Gen.list (Range.linear 0 4) genJsonValue+ , do+ pairs <- Gen.list (Range.linear 0 4) genJsonPair+ pure (object pairs)+ ]+ where+ genJsonPair = do+ k <- Key.fromText . T.pack <$> Gen.list (Range.linear 1 10) (Gen.element ['a'..'z'])+ v <- genJsonValue+ pure (k .= v)
+ test/Valiant/Binary/RangeSpec.hs view
@@ -0,0 +1,59 @@+module Valiant.Binary.RangeSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int32)+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import Valiant.Binary.Range+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "range round-trip" $ do+ it "empty range" $ do+ let r = EmptyRange :: PgRange Int32+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "inclusive-inclusive [1, 10]" $ do+ let r = PgRange (Just (Inclusive (1 :: Int32))) (Just (Inclusive 10))+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "exclusive-exclusive (1, 10)" $ do+ let r = PgRange (Just (Exclusive (1 :: Int32))) (Just (Exclusive 10))+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "inclusive-exclusive [1, 10)" $ do+ let r = PgRange (Just (Inclusive (1 :: Int32))) (Just (Exclusive 10))+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "unbounded lower (,10]" $ do+ let r = PgRange Nothing (Just (Inclusive (10 :: Int32)))+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "unbounded upper [1,)" $ do+ let r = PgRange (Just (Inclusive (1 :: Int32))) Nothing+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "fully unbounded (,)" $ do+ let r = PgRange Nothing Nothing :: PgRange Int32+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ it "negative values [-100, -1]" $ do+ let r = PgRange (Just (Inclusive (-100 :: Int32))) (Just (Inclusive (-1)))+ pgDecodeRange pgDecode (pgEncodeRange pgEncode r) `shouldBe` Right r++ describe "range decode errors" $ do+ it "fails on empty input" $ do+ let result = pgDecodeRange (pgDecode @Int32) BS.empty+ result `shouldSatisfy` isLeft++ describe "range binary format" $ do+ it "empty range is a single byte" $ do+ let encoded = pgEncodeRange (pgEncode @Int32) EmptyRange+ BS.length encoded `shouldBe` 1+ BS.index encoded 0 `shouldBe` 0x01++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/Binary/RefineSpec.hs view
@@ -0,0 +1,58 @@+module Valiant.Binary.RefineSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int32)+import Data.Text (Text)+import Data.Text qualified as T+import Valiant.Binary.Decode (refine, refineWith)+import Valiant.Binary.Encode ()+import PgWire.Binary.Types (PgEncode (..))+import Test.Hspec++newtype PositiveInt = PositiveInt Int32+ deriving (Eq, Show)++mkPositive :: Int32 -> Either String PositiveInt+mkPositive n+ | n > 0 = Right (PositiveInt n)+ | otherwise = Left ("expected positive, got " <> show n)++newtype Email = Email Text+ deriving (Eq, Show)++mkEmail :: Text -> Either String Email+mkEmail t+ | T.any (== '@') t = Right (Email t)+ | otherwise = Left "missing '@'"++spec :: Spec+spec = do+ describe "refine" $ do+ it "passes through valid values" $+ refine mkPositive (pgEncode (42 :: Int32))+ `shouldBe` Right (PositiveInt 42)++ it "fails when the predicate rejects" $+ refine mkPositive (pgEncode (0 :: Int32))+ `shouldBe` (Left "expected positive, got 0" :: Either String PositiveInt)++ it "fails when the underlying decode fails" $+ refine mkPositive (BS.pack [0, 0])+ `shouldBe` (Left "int32: expected 4 bytes, got 2" :: Either String PositiveInt)++ it "works for Text refinements" $+ refine mkEmail (pgEncode ("a@b" :: Text))+ `shouldBe` Right (Email "a@b")++ describe "refineWith" $ do+ it "prepends context to predicate failures" $+ refineWith "PositiveInt" mkPositive (pgEncode (-1 :: Int32))+ `shouldBe` (Left "PositiveInt: expected positive, got -1" :: Either String PositiveInt)++ it "prepends context to decode failures" $+ refineWith "PositiveInt" mkPositive (BS.pack [0])+ `shouldBe` (Left "PositiveInt: int32: expected 4 bytes, got 1" :: Either String PositiveInt)++ it "passes valid values unchanged" $+ refineWith "Email" mkEmail (pgEncode ("x@y" :: Text))+ `shouldBe` Right (Email "x@y")
+ test/Valiant/Binary/ScientificSpec.hs view
@@ -0,0 +1,48 @@+module Valiant.Binary.ScientificSpec (spec) where++import Data.Scientific (Scientific, scientific)+import Valiant.Binary.Scientific ()+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Test.Hspec++spec :: Spec+spec = do+ describe "numeric round-trip" $ do+ it "zero" $ do+ roundTrip (0 :: Scientific) `shouldBe` Right 0++ it "positive integer (42)" $ do+ roundTrip (42 :: Scientific) `shouldBe` Right 42++ it "negative integer (-100)" $ do+ roundTrip (-100 :: Scientific) `shouldBe` Right (-100)++ it "small decimal (3.14)" $ do+ roundTrip (3.14 :: Scientific) `shouldBe` Right 3.14++ it "negative decimal (-0.001)" $ do+ roundTrip (-0.001 :: Scientific) `shouldBe` Right (-0.001)++ it "large integer (1000000)" $ do+ roundTrip (1000000 :: Scientific) `shouldBe` Right 1000000++ it "scientific notation (1.23e10)" $ do+ roundTrip (scientific 123 8) `shouldBe` Right (scientific 123 8)++ it "very small (0.0001)" $ do+ roundTrip (0.0001 :: Scientific) `shouldBe` Right 0.0001++ it "one" $ do+ roundTrip (1 :: Scientific) `shouldBe` Right 1++ it "ten thousand" $ do+ roundTrip (10000 :: Scientific) `shouldBe` Right 10000++ it "9999" $ do+ roundTrip (9999 :: Scientific) `shouldBe` Right 9999++ it "negative one" $ do+ roundTrip (-1 :: Scientific) `shouldBe` Right (-1)++roundTrip :: Scientific -> Either String Scientific+roundTrip = pgDecode . pgEncode
+ test/Valiant/Binary/UUIDSpec.hs view
@@ -0,0 +1,48 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+module Valiant.Binary.UUIDSpec (spec) where++import Data.ByteString qualified as BS+import Data.UUID.Types qualified as UUID+import PgWire.Binary.Types (PgDecode (..), PgEncode (..))+import Valiant.Binary.UUID ()+import Test.Hspec++spec :: Spec+spec = do+ describe "UUID round-trip" $ do+ it "nil UUID" $ do+ let u = UUID.nil+ pgDecode (pgEncode u) `shouldBe` Right u++ it "known UUID" $ do+ let Just u = UUID.fromString "550e8400-e29b-41d4-a716-446655440000"+ pgDecode (pgEncode u) `shouldBe` Right u++ it "another UUID" $ do+ let Just u = UUID.fromString "6ba7b810-9dad-11d1-80b4-00c04fd430c8"+ pgDecode (pgEncode u) `shouldBe` Right u++ describe "UUID encode" $ do+ it "produces exactly 16 bytes" $ do+ let Just u = UUID.fromString "550e8400-e29b-41d4-a716-446655440000"+ BS.length (pgEncode u) `shouldBe` 16++ it "nil UUID encodes to 16 zero bytes" $+ pgEncode UUID.nil `shouldBe` BS.replicate 16 0++ describe "UUID decode errors" $ do+ it "fails on wrong byte count (15 bytes)" $ do+ let result = pgDecode (BS.replicate 15 0) :: Either String UUID.UUID+ result `shouldSatisfy` isLeft++ it "fails on wrong byte count (17 bytes)" $ do+ let result = pgDecode (BS.replicate 17 0) :: Either String UUID.UUID+ result `shouldSatisfy` isLeft++ it "fails on empty input" $ do+ let result = pgDecode BS.empty :: Either String UUID.UUID+ result `shouldSatisfy` isLeft++isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False
+ test/Valiant/ErrorSpec.hs view
@@ -0,0 +1,143 @@+module Valiant.ErrorSpec (spec) where++import Control.Exception (try)+import Data.ByteString.Char8 qualified as BS8+import Valiant.Error+import PgWire.Connection (Connection, close, connectString, simpleQuery)+import PgWire.Error (PgWireError (..))+import PgWire.MockServer+import Test.Hspec++withMockConn :: MockConfig -> (Connection -> IO a) -> IO a+withMockConn 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++triggerError :: Connection -> IO PgWireError+triggerError conn = do+ result <- try (simpleQuery conn "TRIGGER ERROR")+ case result of+ Left err -> pure err+ Right _ -> error "Expected an error but got a result"++spec :: Spec+spec = do+ describe "sqlState" $ do+ it "extracts SQLSTATE from a query error" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "42P01" "relation does not exist"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ sqlState err `shouldBe` Just "42P01"++ it "returns Nothing for non-query errors" $ do+ sqlState (DecodeError "test") `shouldBe` Nothing+ sqlState PoolTimeout `shouldBe` Nothing+ sqlState PoolClosed `shouldBe` Nothing++ describe "constraint violation predicates" $ do+ it "detects unique violation (23505)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "23505" "duplicate key"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isUniqueViolation err `shouldBe` True+ isForeignKeyViolation err `shouldBe` False++ it "detects foreign key violation (23503)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "23503" "violates foreign key"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isForeignKeyViolation err `shouldBe` True+ isUniqueViolation err `shouldBe` False++ it "detects not-null violation (23502)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "23502" "null value"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isNotNullViolation err `shouldBe` True++ it "detects check violation (23514)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "23514" "check constraint"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isCheckViolation err `shouldBe` True++ describe "transaction error predicates" $ do+ it "detects serialization failure (40001)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "40001" "could not serialize"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isSerializationError err `shouldBe` True+ isDeadlockError err `shouldBe` False++ it "detects deadlock (40P01)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "40P01" "deadlock detected"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isDeadlockError err `shouldBe` True++ describe "other error predicates" $ do+ it "detects undefined table (42P01)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "42P01" "relation does not exist"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isUndefinedTable err `shouldBe` True++ it "detects syntax error (42601)" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "42601" "syntax error"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ isSyntaxError err `shouldBe` True++ describe "constraintViolation" $ do+ it "extracts UniqueViolation from 23505" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "23505" "duplicate key"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ case constraintViolation err of+ Just (UniqueViolation _) -> pure ()+ other -> expectationFailure $ "Expected UniqueViolation, got: " <> show other++ it "returns Nothing for non-constraint errors" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "42P01" "relation does not exist"+ }+ withMockConn cfg $ \conn -> do+ err <- triggerError conn+ constraintViolation err `shouldBe` Nothing++ describe "catchConstraintViolation" $ do+ it "catches and handles constraint violations" $ do+ let cfg = defaultMockConfig+ { mockQueryHandler = errorHandler "23505" "duplicate key"+ }+ withMockConn cfg $ \conn -> do+ result <- catchConstraintViolation+ (\_ cv -> case cv of+ UniqueViolation _ -> pure ("caught" :: String)+ _ -> pure "wrong type")+ (simpleQuery conn "INSERT" >> pure "no error")+ result `shouldBe` ("caught" :: String)
+ test/Valiant/FromRowSpec.hs view
@@ -0,0 +1,118 @@+{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}+module Valiant.FromRowSpec (spec) where++import Data.ByteString qualified as BS+import Data.Int (Int32, Int64)+import Data.Text (Text)+import Data.Vector qualified as V+import Valiant.Binary.Decode ()+import Valiant.Binary.Encode ()+import PgWire.Binary.Types (PgEncode (..))+import Valiant.FromRow+import Test.Hspec++-- Helper: encode a value and wrap as Just+enc :: (PgEncode a) => a -> Maybe BS.ByteString+enc = Just . pgEncode++spec :: Spec+spec = do+ describe "2-tuple FromRow" $ do+ it "decodes (Int32, Text)" $ do+ let row = V.fromList [enc (42 :: Int32), enc ("hello" :: Text)]+ fromRow row `shouldBe` Right (42 :: Int32, "hello" :: Text)++ it "decodes (Bool, Int64)" $ do+ let row = V.fromList [enc True, enc (999 :: Int64)]+ fromRow row `shouldBe` Right (True, 999 :: Int64)++ describe "3-tuple FromRow" $ do+ it "decodes (Int32, Text, Bool)" $ do+ let row = V.fromList [enc (1 :: Int32), enc ("alice" :: Text), enc True]+ fromRow row `shouldBe` Right (1 :: Int32, "alice" :: Text, True)++ it "decodes (Int64, Int32, Text)" $ do+ let row = V.fromList [enc (100 :: Int64), enc (42 :: Int32), enc ("test" :: Text)]+ fromRow row `shouldBe` Right (100 :: Int64, 42 :: Int32, "test" :: Text)++ describe "4-tuple FromRow" $ do+ it "decodes (Int32, Text, Bool, Int64)" $ do+ let row = V.fromList+ [ enc (1 :: Int32)+ , enc ("bob" :: Text)+ , enc False+ , enc (9999 :: Int64)+ ]+ fromRow row `shouldBe` Right (1 :: Int32, "bob" :: Text, False, 9999 :: Int64)++ describe "5-tuple FromRow" $ do+ it "decodes (Int32, Text, Int64, Bool, Double)" $ do+ let row = V.fromList+ [ enc (10 :: Int32)+ , enc ("carol" :: Text)+ , enc (500 :: Int64)+ , enc True+ , enc (3.14 :: Double)+ ]+ fromRow row `shouldBe` Right (10 :: Int32, "carol" :: Text, 500 :: Int64, True, 3.14 :: Double)++ describe "Column index out of range" $ do+ it "errors when row has fewer columns than expected (2-tuple with 1 column)" $ do+ let row = V.fromList [enc (42 :: Int32)]+ result = fromRow row :: Either String (Int32, Text)+ result `shouldSatisfy` isLeft++ it "errors when row has 0 columns for a 2-tuple" $ do+ let row = V.empty+ result = fromRow row :: Either String (Int32, Text)+ result `shouldSatisfy` isLeft++ it "errors when row has 2 columns for a 3-tuple" $ do+ let row = V.fromList [enc (1 :: Int32), enc ("x" :: Text)]+ result = fromRow row :: Either String (Int32, Text, Bool)+ result `shouldSatisfy` isLeft++ describe "NULL for non-nullable column" $ do+ it "errors when a non-Maybe column is NULL (first column)" $ do+ let row = V.fromList [Nothing, enc ("hello" :: Text)]+ result = fromRow row :: Either String (Int32, Text)+ result `shouldSatisfy` isLeft++ it "errors when a non-Maybe column is NULL (second column)" $ do+ let row = V.fromList [enc (42 :: Int32), Nothing]+ result = fromRow row :: Either String (Int32, Text)+ result `shouldSatisfy` isLeft++ it "error message mentions NULL" $ do+ let row = V.fromList [Nothing, enc ("hello" :: Text)]+ Left msg = fromRow row :: Either String (Int32, Text)+ msg `shouldSatisfy` \s -> "NULL" `isInfixOfStr` s++ describe "6-tuple FromRow" $ do+ it "decodes (Int32, Text, Bool, Int64, Double, Int32)" $ do+ let row = V.fromList+ [ enc (1 :: Int32)+ , enc ("test" :: Text)+ , enc True+ , enc (999 :: Int64)+ , enc (2.71 :: Double)+ , enc (7 :: Int32)+ ]+ fromRow row `shouldBe` Right (1 :: Int32, "test" :: Text, True, 999 :: Int64, 2.71 :: Double, 7 :: Int32)++-- Helpers+isLeft :: Either a b -> Bool+isLeft (Left _) = True+isLeft _ = False++isInfixOfStr :: String -> String -> Bool+isInfixOfStr needle haystack = any (isPrefixOfStr needle) (tails haystack)++isPrefixOfStr :: String -> String -> Bool+isPrefixOfStr [] _ = True+isPrefixOfStr _ [] = False+isPrefixOfStr (x : xs) (y : ys) = x == y && isPrefixOfStr xs ys++tails :: [a] -> [[a]]+tails [] = [[]]+tails xs@(_ : xs') = xs : tails xs'
+ test/Valiant/NamedParamsSpec.hs view
@@ -0,0 +1,180 @@+{-# LANGUAGE DeriveAnyClass #-}+{-# LANGUAGE DerivingStrategies #-}++module Valiant.NamedParamsSpec (spec) where++import Control.Exception (ErrorCall (..), evaluate, try)+import Data.Int (Int32)+import Data.Text (Text)+import Data.Text.Encoding qualified as TE+import Data.Vector qualified as V+import GHC.Generics (Generic)+import Valiant.NamedParams+import Valiant.Statement (Statement (..))+import PgWire.Binary.Types (PgEncode (..))+import Test.Hspec++-- Test record types+data SimpleParams = SimpleParams+ { userId :: Int32+ }+ deriving stock (Generic)+ deriving anyclass (ToNamedParams)++data MultiParams = MultiParams+ { orgId :: Int32+ , role :: Text+ , active :: Bool+ }+ deriving stock (Generic)+ deriving anyclass (ToNamedParams)++data NullableParams = NullableParams+ { name :: Text+ , email :: Maybe Text+ }+ deriving stock (Generic)+ deriving anyclass (ToNamedParams)++-- Record with fields that don't match any SQL param+data ExtraFieldParams = ExtraFieldParams+ { knownField :: Int32+ , unknownField :: Text+ }+ deriving stock (Generic)+ deriving anyclass (ToNamedParams)++spec :: Spec+spec = do+ describe "ToNamedParams (Generic)" $ do+ it "encodes a single-field record" $ do+ let pairs = toNamedParamList (SimpleParams 42)+ length pairs `shouldBe` 1+ case pairs of+ [(n, val)] -> do+ n `shouldBe` "userId"+ val `shouldBe` Just (pgEncode (42 :: Int32))+ _ -> expectationFailure "expected exactly one pair"++ it "encodes a multi-field record with correct names" $ do+ let pairs = toNamedParamList (MultiParams 1 "admin" True)+ names = map fst pairs+ names `shouldBe` ["orgId", "role", "active"]++ it "encodes a multi-field record with correct values" $ do+ let pairs = toNamedParamList (MultiParams 1 "admin" True)+ snd (pairs !! 0) `shouldBe` Just (pgEncode (1 :: Int32))+ snd (pairs !! 1) `shouldBe` Just (TE.encodeUtf8 "admin")+ snd (pairs !! 2) `shouldBe` Just (pgEncode True)++ it "encodes Nothing as NULL" $ do+ let pairs = toNamedParamList (NullableParams "Alice" Nothing)+ snd (pairs !! 0) `shouldBe` Just (TE.encodeUtf8 "Alice")+ snd (pairs !! 1) `shouldBe` Nothing++ it "encodes Just value correctly" $ do+ let pairs = toNamedParamList (NullableParams "Alice" (Just "alice@example.com"))+ snd (pairs !! 1) `shouldBe` Just (TE.encodeUtf8 "alice@example.com")++ describe "mkStatementNamed" $ do+ it "reorders fields to match positional parameter order" $ do+ let stmt = mkStatementNamed @MultiParams @()+ "SELECT 1 WHERE org_id = $1 AND active = $2 AND role = $3"+ [23, 16, 25]+ []+ ["orgId", "active", "role"]+ "test.sql"+ -- orgId=$1, active=$2, role=$3+ encoded = stmtEncode stmt (MultiParams 42 "admin" True)+ -- $1 = orgId = 42+ (V.!) encoded 0 `shouldBe` Just (pgEncode (42 :: Int32))+ -- $2 = active = True+ (V.!) encoded 1 `shouldBe` Just (pgEncode True)+ -- $3 = role = "admin"+ (V.!) encoded 2 `shouldBe` Just (TE.encodeUtf8 "admin")++ it "handles duplicate param names (record field used once)" $ do+ let stmt = mkStatementNamed @SimpleParams @()+ "SELECT 1 WHERE a = $1 OR b = $1"+ [23]+ []+ ["userId"]+ "test.sql"+ encoded = stmtEncode stmt (SimpleParams 99)+ V.length encoded `shouldBe` 1+ (V.!) encoded 0 `shouldBe` Just (pgEncode (99 :: Int32))++ it "preserves original field order in same-order case" $ do+ let stmt = mkStatementNamed @MultiParams @()+ "SELECT 1"+ [23, 25, 16]+ []+ ["orgId", "role", "active"]+ "test.sql"+ encoded = stmtEncode stmt (MultiParams 7 "viewer" False)+ (V.!) encoded 0 `shouldBe` Just (pgEncode (7 :: Int32))+ (V.!) encoded 1 `shouldBe` Just (TE.encodeUtf8 "viewer")+ (V.!) encoded 2 `shouldBe` Just (pgEncode False)++ it "handles nullable params in named mode" $ do+ let stmt = mkStatementNamed @NullableParams @()+ "INSERT INTO t (name, email) VALUES ($1, $2)"+ [25, 25]+ []+ ["name", "email"]+ "test.sql"+ encoded = stmtEncode stmt (NullableParams "Bob" Nothing)+ (V.!) encoded 0 `shouldBe` Just (TE.encodeUtf8 "Bob")+ (V.!) encoded 1 `shouldBe` Nothing++ describe "error messages" $ do+ it "reports missing record fields with helpful message" $ do+ let stmt = mkStatementNamed @SimpleParams @()+ "SELECT 1 WHERE org = $1 AND role = $2"+ [23, 25]+ []+ ["orgId", "role"] -- record only has userId, not orgId or role+ "test.sql"+ result <- try @ErrorCall $ evaluate (stmtEncode stmt (SimpleParams 1))+ case result of+ Left (ErrorCall msg) -> do+ msg `shouldContain` "named parameter mismatch"+ msg `shouldContain` "Missing fields"+ msg `shouldContain` ":orgId"+ msg `shouldContain` ":role"+ msg `shouldContain` "userId" -- listed as available+ Right _ -> expectationFailure "expected error for missing fields"++ it "reports extra record fields with helpful message" $ do+ let stmt = mkStatementNamed @ExtraFieldParams @()+ "SELECT 1 WHERE id = $1"+ [23]+ []+ ["knownField"] -- SQL only uses knownField, not unknownField+ "test.sql"+ result <- try @ErrorCall $ evaluate (stmtEncode stmt (ExtraFieldParams 1 "extra"))+ case result of+ Left (ErrorCall msg) -> do+ msg `shouldContain` "named parameter mismatch"+ msg `shouldContain` "Extra fields"+ msg `shouldContain` "unknownField"+ msg `shouldContain` "Hint"+ Right _ -> expectationFailure "expected error for extra fields"++ it "reports both missing and extra fields together" $ do+ let stmt = mkStatementNamed @SimpleParams @()+ "SELECT 1 WHERE name = $1"+ [25]+ []+ ["name"] -- SQL expects "name" but record has "userId"+ "test.sql"+ result <- try @ErrorCall $ evaluate (stmtEncode stmt (SimpleParams 1))+ case result of+ Left (ErrorCall msg) -> do+ msg `shouldContain` "Missing fields"+ msg `shouldContain` ":name"+ msg `shouldContain` "Extra fields"+ msg `shouldContain` "userId"+ Right _ -> expectationFailure "expected error for mismatched fields"++-- Use hspec's shouldContain (works on [Char] = String)
+ test/Valiant/TupleSpec.hs view
@@ -0,0 +1,119 @@+module Valiant.TupleSpec (spec) where++import Data.Int (Int32, Int64)+import Data.Text (Text)+import Valiant.FromRow (fromRow)+import Valiant.ToParams (toParams)+import Test.Hspec++-- | Round-trip each tuple size through ToParams (encode) then FromRow+-- (decode). Covers sizes 7-16; sizes 2-6 are covered by FromRowSpec.+-- The encoded form is @Vector (Maybe ByteString)@, which both classes+-- share, so toParams >>> fromRow should be Right . id.+spec :: Spec+spec = do+ describe "7-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t = (1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double, 4 :: Int32, "b" :: Text)+ fromRow (toParams t) `shouldBe` Right t++ describe "8-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t = (1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double, 4 :: Int32, "b" :: Text, False)+ fromRow (toParams t) `shouldBe` Right t++ describe "9-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64+ )+ fromRow (toParams t) `shouldBe` Right t++ describe "10-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ )+ fromRow (toParams t) `shouldBe` Right t++ describe "11-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ , 7 :: Int32+ )+ fromRow (toParams t) `shouldBe` Right t++ describe "12-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ , 7 :: Int32, "c" :: Text+ )+ fromRow (toParams t) `shouldBe` Right t++ describe "13-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ , 7 :: Int32, "c" :: Text, True+ )+ fromRow (toParams t) `shouldBe` Right t++ describe "14-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ , 7 :: Int32, "c" :: Text, True, 8 :: Int64+ )+ fromRow (toParams t) `shouldBe` Right t++ describe "15-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ , 7 :: Int32, "c" :: Text, True, 8 :: Int64, 9.75 :: Double+ )+ fromRow (toParams t) `shouldBe` Right t++ -- base provides Show and Eq only up to 15-tuples, so 16-tuple+ -- round trips are checked by destructuring the result instead of+ -- using shouldBe on the whole value.+ describe "16-tuple round trip" $+ it "encodes and decodes back to the original" $ do+ let t =+ ( 1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double+ , 4 :: Int32, "b" :: Text, False, 5 :: Int64, 6.25 :: Double+ , 7 :: Int32, "c" :: Text, True, 8 :: Int64, 9.75 :: Double+ , 10 :: Int32+ )+ case fromRow (toParams t) of+ Left err -> expectationFailure err+ Right (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) -> do+ (a, b, c, d, e, f, g, h) `shouldBe`+ (1 :: Int32, "a" :: Text, True, 2 :: Int64, 3.5 :: Double, 4 :: Int32, "b" :: Text, False)+ (i, j, k, l, m, n, o, p) `shouldBe`+ (5 :: Int64, 6.25 :: Double, 7 :: Int32, "c" :: Text, True, 8 :: Int64, 9.75 :: Double, 10 :: Int32)++ describe "16-tuple with NULLs via Maybe" $+ it "round trips with Nothing mixed in" $ do+ let t =+ ( Just (1 :: Int32), Nothing :: Maybe Text, True, 2 :: Int64, 3.5 :: Double+ , Nothing :: Maybe Int32, "b" :: Text, False, Just (5 :: Int64), 6.25 :: Double+ , 7 :: Int32, Nothing :: Maybe Text, True, 8 :: Int64, Nothing :: Maybe Double+ , 10 :: Int32+ )+ case fromRow (toParams t) of+ Left err -> expectationFailure err+ Right (a, b, c, d, e, f, g, h, i, j, k, l, m, n, o, p) -> do+ (a, b, c, d, e, f, g, h) `shouldBe`+ (Just (1 :: Int32), Nothing :: Maybe Text, True, 2 :: Int64, 3.5 :: Double, Nothing :: Maybe Int32, "b" :: Text, False)+ (i, j, k, l, m, n, o, p) `shouldBe`+ (Just (5 :: Int64), 6.25 :: Double, 7 :: Int32, Nothing :: Maybe Text, True, 8 :: Int64, Nothing :: Maybe Double, 10 :: Int32)
+ valiant.cabal view
@@ -0,0 +1,231 @@+cabal-version: 3.0+name: valiant+version: 0.1.0.0+synopsis: Compile-time checked SQL for Haskell, runtime library+description:+ Runtime library for valiant. Provides the @Statement@ type, binary format+ codecs, query execution functions, and re-exports from @pg-wire@ (the+ underlying PostgreSQL wire protocol driver).+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: runtime++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:+ Valiant+ Valiant.Advisory+ Valiant.Batch+ Valiant.Binary.Array+ Valiant.Binary.Composite+ Valiant.Binary.Decode+ Valiant.Binary.Encode+ Valiant.Binary.Enum+ Valiant.Binary.HStore+ Valiant.Binary.Inet+ Valiant.Binary.MacAddr+ Valiant.Binary.Interval+ Valiant.Binary.Point+ Valiant.Binary.Unbounded+ Valiant.Binary.JSON+ Valiant.Binary.Range+ Valiant.Binary.Scientific+ Valiant.Binary.UUID+ Valiant.Copy+ Valiant.Dynamic+ Valiant.Error+ Valiant.Execute+ Valiant.Fold+ Valiant.FromRow+ Valiant.FromRowFast+ Valiant.FromRowStrict+ Valiant.LargeObject+ Valiant.Logging+ Valiant.NamedParams+ Valiant.Notify+ Valiant.Pipeline+ Valiant.Statement+ Valiant.Streaming+ Valiant.ToParams+ Valiant.Transaction++ build-depends:+ -- Wire protocol driver+ , pg-wire ==0.1.*++ -- Core+ , base >=4.17 && <5+ , bytestring >=0.11 && <0.13+ , text >=2.0 && <2.2+ , containers >=0.6 && <0.8+ , hashable >=1.4 && <1.6+ , psqueues >=0.2.8 && <0.3+ , vector >=0.13 && <0.14+ , time >=1.12 && <1.15+ , deepseq >=1.4 && <1.6++ -- Binary codec helpers+ , scientific >=0.3 && <0.4+ , uuid-types >=1.0 && <1.1+ , aeson >=2.1 && <2.3++ -- Concurrency (for Streaming, Notify)+ , async >=2.2 && <2.3++ -- Network (for Cancel, Notify)+ , network >=3.1 && <3.3++test-suite valiant-test+ import: warnings+ type: exitcode-stdio-1.0+ hs-source-dirs: test+ main-is: Main.hs+ default-language: GHC2021+ default-extensions:+ OverloadedStrings+ ghc-options: -rtsopts "-with-rtsopts=-K8K"++ other-modules:+ Valiant.Binary.ArraySpec+ Valiant.Binary.CompositeSpec+ Valiant.Binary.DecodeSpec+ Valiant.Binary.EncodeSpec+ Valiant.Binary.HStoreSpec+ Valiant.Binary.InetSpec+ Valiant.Binary.MacAddrSpec+ Valiant.Binary.IntervalSpec+ Valiant.Binary.JSONSpec+ Valiant.Binary.NewtypeSpec+ Valiant.Binary.PropertyHedgehogSpec+ Valiant.Binary.RangeSpec+ Valiant.Binary.RefineSpec+ Valiant.Binary.ScientificSpec+ Valiant.Binary.UUIDSpec+ Valiant.ErrorSpec+ Valiant.FromRowSpec+ Valiant.NamedParamsSpec+ Valiant.TupleSpec++ build-depends:+ , aeson+ , base >=4.17 && <5+ , bytestring+ , containers+ , hedgehog >=1.2 && <1.6+ , hspec >=2.11 && <2.13+ , hspec-hedgehog >=0.1 && <0.3+ , valiant+ , pg-wire+ , pg-wire:mockserver+ , scientific+ , text+ , uuid-types+ , time+ , vector++test-suite valiant-integration+ import: warnings+ type: exitcode-stdio-1.0+ hs-source-dirs: integration+ main-is: Main.hs+ default-language: GHC2021+ default-extensions:+ OverloadedStrings+ ghc-options: -threaded -rtsopts "-with-rtsopts=-K8K"++ other-modules:+ ChaosSpec+ ConnectionSpec+ ConnectionFeaturesSpec+ CopySpec+ ExecuteSpec+ LargeObjectSpec+ NotifySpec+ PipelineSpec+ PoolSpec+ SoundnessSpec+ StreamingSpec+ TestSupport+ TransactionSpec++ build-depends:+ , async >=2.2 && <2.3+ , base >=4.17 && <5+ , bytestring+ , containers+ , hspec >=2.11 && <2.13+ , nothunks >=0.1 && <0.3+ , valiant+ , pg-wire+ , stm >=2.5 && <2.6+ , text+ , time >=1.12 && <1.15+ , vector++ -- Requires DATABASE_URL. Run with:+ -- eval $(scripts/pg-setup.sh)+ -- cabal test valiant-integration --test-show-details=direct++benchmark valiant-bench+ import: warnings+ type: exitcode-stdio-1.0+ hs-source-dirs: bench+ main-is: Main.hs+ default-language: GHC2021+ default-extensions:+ OverloadedStrings++ other-modules:+ BenchAsyncpgCompare+ BenchCodec+ BenchConcurrent+ BenchPool+ BenchQuery+ TestSupport++ build-depends:+ , async >=2.2 && <2.3+ , base >=4.17 && <5+ , bytestring+ , containers+ , criterion >=1.6 && <1.8+ , valiant+ , pg-wire+ , scientific+ , text+ , time+ , vector+ ghc-options: -threaded -rtsopts "-with-rtsopts=-I0 -A32m"