diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,10 @@ +# Changelog + +## 0.1.0.0 + +- Initial release +- Core functionality: `with`, `withConfig`, `start`, `stop` +- Hasql integration via `connectionSettings` +- initdb caching support +- Copy-on-write support for macOS and Linux +- Snapshot support diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2025, Nadeem Bitar + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * 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. + + * 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. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,429 @@ +# ephemeral-pg + +A modern Haskell library for creating temporary PostgreSQL databases for testing. + +## Features + +- Native hasql integration +- initdb caching for fast startup +- Copy-on-write support (macOS/Linux) +- Filesystem snapshots +- No shell injection vulnerabilities +- Type-safe configuration + +## Requirements + +- GHC 9.6+ +- PostgreSQL 14+ + +## Installation + +Add to your `cabal` file: + +```cabal +build-depends: + ephemeral-pg +``` + +Or with Stack, add to your `package.yaml`: + +```yaml +dependencies: + - ephemeral-pg +``` + +## Quick Start + +```haskell +import EphemeralPg qualified as Pg +import Hasql.Connection qualified as Connection + +main :: IO () +main = do + result <- Pg.with \db -> do + Right conn <- Connection.acquire (Pg.connectionSettings db) + -- Use the connection... + Connection.release conn + case result of + Left err -> putStrLn $ "Error: " <> Pg.renderStartError err + Right () -> putStrLn "Success!" +``` + +## Usage Examples + +### Basic Usage + +The simplest way to use ephemeral-pg is with the `with` function, which creates +a temporary database, runs your action, and cleans up automatically: + +```haskell +import EphemeralPg qualified as Pg +import Hasql.Connection qualified as Connection +import Hasql.Session qualified as Session +import Hasql.Statement qualified as Statement + +testQuery :: IO () +testQuery = do + result <- Pg.with \db -> do + Right conn <- Connection.acquire (Pg.connectionSettings db) + + -- Run a simple query + result <- Session.run (Session.statement () selectOne) conn + + Connection.release conn + pure result + + case result of + Left err -> putStrLn $ "Startup error: " <> Pg.renderStartError err + Right (Left sessionErr) -> putStrLn $ "Query error: " <> show sessionErr + Right (Right value) -> putStrLn $ "Result: " <> show value + where + selectOne = Statement.Statement "SELECT 1" mempty decoder True + decoder = Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4)) +``` + +### Custom Configuration + +You can customize the database configuration: + +```haskell +import EphemeralPg qualified as Pg +import EphemeralPg.Config qualified as Config + +customTest :: IO () +customTest = do + let config = Pg.defaultConfig + { Config.configDatabaseName = "mytest" + , Config.configPostgresSettings = + [ ("log_statement", "'all'") + , ("log_min_duration_statement", "0") + ] + } + + result <- Pg.withConfig config \db -> do + -- Database name is "mytest" + -- All statements are logged + pure () + + case result of + Left err -> putStrLn $ Pg.renderStartError err + Right () -> putStrLn "Done!" +``` + +### Using Verbose Configuration + +For debugging, use the pre-configured verbose settings: + +```haskell +import EphemeralPg qualified as Pg + +debugTest :: IO () +debugTest = do + result <- Pg.withConfig Pg.verboseConfig \db -> do + -- All statements logged with duration + pure () + pure () +``` + +### Using auto_explain + +For query plan analysis: + +```haskell +import EphemeralPg qualified as Pg + +analyzeQueries :: IO () +analyzeQueries = do + result <- Pg.withConfig Pg.autoExplainConfig \db -> do + -- Query plans automatically logged for slow queries + pure () + pure () +``` + +### Caching for Fast Startup + +For faster test suite execution, use caching. The first run initializes the +cache, and subsequent runs copy from it (using CoW if available): + +```haskell +import EphemeralPg qualified as Pg + +cachedTest :: IO () +cachedTest = do + -- First call: ~2s (runs initdb and caches result) + -- Subsequent calls: ~200ms (copies from cache) + result <- Pg.withCached \db -> do + -- Use the database... + pure () + pure () +``` + +### Manual Lifecycle Management + +For more control, use `start` and `stop` directly: + +```haskell +import EphemeralPg qualified as Pg +import Control.Exception (bracket) + +manualLifecycle :: IO () +manualLifecycle = do + result <- Pg.start Pg.defaultConfig + case result of + Left err -> putStrLn $ Pg.renderStartError err + Right db -> do + -- Use the database... + putStrLn $ "Database running on port: " <> show (Pg.port db) + + -- Clean up when done + Pg.stop db +``` + +Or with bracket for exception safety: + +```haskell +import EphemeralPg qualified as Pg +import Control.Exception (bracket) + +bracketExample :: IO () +bracketExample = do + result <- Pg.start Pg.defaultConfig + case result of + Left err -> putStrLn $ Pg.renderStartError err + Right db -> + bracket (pure db) Pg.stop \db' -> do + -- Use the database... + pure () +``` + +### Restarting the Database + +You can restart the database to apply configuration changes or test recovery: + +```haskell +import EphemeralPg qualified as Pg + +restartTest :: IO () +restartTest = do + result <- Pg.with \db -> do + let port1 = Pg.port db + + -- Restart the server + restartResult <- Pg.restart db + case restartResult of + Left err -> fail $ "Restart failed: " <> show err + Right db' -> do + -- Server restarted, data preserved + let port2 = Pg.port db' + -- Port remains the same + pure (port1 == port2) + + case result of + Left err -> putStrLn $ Pg.renderStartError err + Right same -> putStrLn $ "Ports same: " <> show same +``` + +### Snapshots + +Create filesystem snapshots for fast test isolation: + +```haskell +import EphemeralPg qualified as Pg +import EphemeralPg.Snapshot qualified as Snapshot + +snapshotTest :: IO () +snapshotTest = do + result <- Pg.with \db -> do + -- Set up initial state + -- ... create tables, insert data ... + + -- Create a snapshot + Right snapshot <- Snapshot.createSnapshot db + + -- Run a destructive test + -- ... delete everything ... + + -- Restore to the snapshot + Right () <- Snapshot.restoreSnapshot snapshot db + + -- Data is restored + + -- Clean up snapshot when done + Snapshot.deleteSnapshot snapshot + + pure () + + pure () +``` + +### Dump and Restore + +Export and import database contents: + +```haskell +import EphemeralPg qualified as Pg +import EphemeralPg.Dump qualified as Dump + +dumpTest :: IO () +dumpTest = do + result <- Pg.with \db -> do + -- Set up data + -- ... + + -- Dump to file + Right () <- Dump.dump db "/tmp/backup.sql" + + pure () + + -- Later, restore to a new database + result2 <- Pg.with \db2 -> do + Right () <- Dump.restore db2 "/tmp/backup.sql" + -- Data is now in db2 + pure () + + pure () +``` + +## Configuration Reference + +### Config Fields + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `configDatabaseName` | `Text` | `"test"` | Name of the database to create | +| `configUser` | `Text` | Current user | PostgreSQL username | +| `configPassword` | `Text` | `""` | PostgreSQL password (empty = trust auth) | +| `configPort` | `Last Int` | Auto-assigned | Port number (finds free port if not specified) | +| `configDataDirectory` | `DirectoryConfig` | Temporary | Data directory location | +| `configSocketDirectory` | `DirectoryConfig` | Temporary | Unix socket directory | +| `configTemporaryRoot` | `Last FilePath` | System temp | Root for temporary directories | +| `configPostgresSettings` | `[(Text, Text)]` | Optimized defaults | postgresql.conf settings | +| `configInitDbArgs` | `[Text]` | `[]` | Additional initdb arguments | + +### Pre-configured Configs + +- `defaultConfig` - Basic configuration with optimized defaults +- `verboseConfig` - All statements logged with duration +- `autoExplainConfig` - Query plans logged for slow queries + +### PostgreSQL Settings Defaults + +The default configuration includes these optimized settings for testing: + +``` +shared_buffers = 12MB +fsync = off +synchronous_commit = off +full_page_writes = off +log_min_duration_statement = 0 +log_connections = on +log_disconnections = on +random_page_cost = 1.0 +``` + +## API Reference + +### Core Functions + +```haskell +-- Bracket-style (recommended) +with :: (Database -> IO a) -> IO (Either StartError a) +withConfig :: Config -> (Database -> IO a) -> IO (Either StartError a) + +-- With caching +withCached :: (Database -> IO a) -> IO (Either StartError a) + +-- Manual lifecycle +start :: Config -> IO (Either StartError Database) +stop :: Database -> IO () +restart :: Database -> IO (Either StartError Database) +``` + +### Database Handle + +```haskell +connectionSettings :: Database -> Settings -- hasql Settings +connectionString :: Database -> Text -- libpq connection string +dataDirectory :: Database -> FilePath +socketDirectory :: Database -> FilePath +port :: Database -> Int +databaseName :: Database -> Text +user :: Database -> Text +``` + +### Cache Management + +```haskell +clearCache :: IO () -- Clear current user's cache +clearAllCaches :: IO () -- Clear all cached data +``` + +## Benchmarks + +Measured with [tasty-bench](https://hackage.haskell.org/package/tasty-bench) in wall-clock mode on Apple Silicon. + +| Benchmark | Time | vs baseline | +|---|---|---| +| **Lifecycle** | | | +| `defaultConfig` | 796 ms | -- | +| `defaultConfig <> verboseConfig` | 812 ms | 1.02x | +| `defaultConfig <> autoExplainConfig 100` | 769 ms | 0.97x | +| **Caching** | | | +| `withConfig` (uncached) | 755 ms | -- | +| `withCached` | 434 ms | 0.58x | +| **Snapshot** | | | +| lifecycle only | 761 ms | -- | +| lifecycle + `createSnapshot` | 1.13 s | 1.48x | +| lifecycle + `createSnapshot` + `restoreSnapshot` | 1.39 s | 1.83x | +| **Connection** | | | +| hasql `acquire` + `release` | 1.64 ms | -- | + +Highlights: +- Cached startup is ~42% faster than uncached (CoW copy vs full initdb) +- `verboseConfig` and `autoExplainConfig` add negligible overhead +- Snapshot create adds ~365 ms, restore adds ~265 ms +- Connection acquire/release is ~1.6 ms once the database is running + +## Comparison with tmp-postgres + +ephemeral-pg is a modern replacement for tmp-postgres, addressing several issues: + +| Feature | tmp-postgres | ephemeral-pg | +|---------|-------------|--------------| +| Shell injection | Vulnerable | Safe (typed-process) | +| Windows support | Partial | Unix-only (explicit) | +| Socket path length | Can fail | Validated | +| GHC version | 8.0+ | 9.6+ | +| hasql integration | Via options | Native | +| initdb caching | Yes | Yes (improved) | +| Copy-on-write | Yes | Yes | +| PostgreSQL version | 9.3+ | 14+ | + +## Troubleshooting + +### "initdb not found" + +Ensure PostgreSQL binaries are in your PATH: + +```bash +# macOS (Homebrew) +export PATH="/opt/homebrew/opt/postgresql@14/bin:$PATH" + +# Linux (Debian/Ubuntu) +export PATH="/usr/lib/postgresql/14/bin:$PATH" +``` + +### Socket path too long + +Unix sockets have a path length limit (~104 bytes). ephemeral-pg uses short +paths to avoid this, but if you specify a custom socket directory, ensure +the full path is under 90 characters. + +### Permission denied + +Ensure you have write access to the temporary directory. By default, the +system temp directory is used. + +## License + +BSD-3-Clause diff --git a/ephemeral-pg.cabal b/ephemeral-pg.cabal new file mode 100644 --- /dev/null +++ b/ephemeral-pg.cabal @@ -0,0 +1,97 @@ +cabal-version: 3.0 +name: ephemeral-pg +version: 0.1.0.0 +synopsis: Temporary PostgreSQL databases for testing +description: + A modern library for creating temporary PostgreSQL instances for testing. + Features include initdb caching, copy-on-write support, and native hasql integration. + +license: BSD-3-Clause +license-file: LICENSE +author: Nadeem Bitar +maintainer: Nadeem Bitar +category: Database, Testing +build-type: Simple +extra-doc-files: + CHANGELOG.md + README.md + +source-repository head + type: git + location: https://github.com/shinzui/ephemeral-pg + +common warnings + ghc-options: + -Wall + -Wcompat + -Widentities + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -Wpartial-fields + -Wredundant-constraints + -Wunused-packages + +library + import: warnings + default-language: GHC2021 + hs-source-dirs: src + exposed-modules: + EphemeralPg + EphemeralPg.Config + EphemeralPg.Database + EphemeralPg.Dump + EphemeralPg.Error + EphemeralPg.Snapshot + + other-modules: + EphemeralPg.Internal.Cache + EphemeralPg.Internal.CopyOnWrite + EphemeralPg.Internal.Directory + EphemeralPg.Internal.Except + EphemeralPg.Internal.Port + EphemeralPg.Process + EphemeralPg.Process.CreateDb + EphemeralPg.Process.InitDb + EphemeralPg.Process.Postgres + + build-depends: + base >=4.18 && <5, + bytestring >=0.11 && <0.13, + directory >=1.3 && <1.4, + filepath >=1.4 && <1.6, + hashable >=1.4 && <1.6, + hasql >=1.10 && <1.11, + network >=3.1 && <3.3, + process >=1.6 && <1.8, + temporary >=1.3 && <1.4, + text >=2.0 && <2.2, + transformers >=0.5 && <0.7, + typed-process >=0.2.12 && <0.3, + unix >=2.8 && <2.9, + + default-extensions: + BlockArguments + DeriveAnyClass + DerivingStrategies + LambdaCase + OverloadedStrings + RecordWildCards + StrictData + +test-suite ephemeral-pg-test + import: warnings + default-language: GHC2021 + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + build-depends: + QuickCheck >=2.14 && <2.16, + base, + ephemeral-pg, + hasql, + hspec >=2.11 && <2.12, + text, + + default-extensions: + BlockArguments + OverloadedStrings diff --git a/src/EphemeralPg.hs b/src/EphemeralPg.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg.hs @@ -0,0 +1,489 @@ +-- | Temporary PostgreSQL databases for testing. +-- +-- This module provides functions for creating isolated, temporary PostgreSQL +-- instances for testing purposes. Databases are automatically cleaned up +-- when they go out of scope. +-- +-- = Quick Start +-- +-- @ +-- import EphemeralPg qualified as Pg +-- import Hasql.Connection qualified as Connection +-- +-- main :: IO () +-- main = do +-- result <- Pg.'with' $ \\db -> do +-- Right conn <- Connection.acquire (Pg.'connectionSettings' db) +-- -- Use the connection... +-- Connection.release conn +-- case result of +-- Left err -> putStrLn $ "Error: " <> show err +-- Right () -> putStrLn "Success!" +-- @ +-- +-- = Custom Configuration +-- +-- @ +-- import EphemeralPg qualified as Pg +-- +-- main :: IO () +-- main = do +-- let config = Pg.'defaultConfig' { Pg.configDatabaseName = "testdb" } +-- Pg.'withConfig' config $ \\db -> do +-- -- Use the database... +-- pure () +-- @ +module EphemeralPg + ( -- * Database Handle + Database, + connectionSettings, + connectionString, + dataDirectory, + socketDirectory, + port, + databaseName, + user, + + -- * Lifecycle Management + with, + withConfig, + withCached, + start, + startCached, + stop, + restart, + + -- * Configuration + Config (..), + DirectoryConfig (..), + ShutdownMode (..), + defaultConfig, + verboseConfig, + autoExplainConfig, + + -- * Cache Management + CacheConfig (..), + defaultCacheConfig, + clearCache, + clearAllCaches, + + -- * Errors + StartError (..), + StopError (..), + renderStartError, + renderStopError, + ) +where + +import Control.Exception (mask, onException) +import Control.Monad (when) +import Control.Monad.IO.Class (liftIO) +import Data.Monoid (Last (..)) +import Data.Text (Text) +import Data.Word (Word16) +import EphemeralPg.Config + ( Config (..), + DirectoryConfig (..), + ShutdownMode (..), + autoExplainConfig, + defaultConfig, + defaultShutdownTimeoutSeconds, + verboseConfig, + ) +import EphemeralPg.Database + ( Database (..), + connectionSettings, + connectionString, + dataDirectory, + databaseName, + port, + socketDirectory, + user, + ) +import EphemeralPg.Error + ( StartError (..), + StopError (..), + renderStartError, + renderStopError, + ) +import EphemeralPg.Internal.Cache + ( CacheConfig (..), + CacheKey, + cleanupRuntimeFiles, + clearAllCaches, + clearCache, + createCache, + defaultCacheConfig, + getCacheKey, + isCached, + restoreFromCache, + ) +import EphemeralPg.Internal.Directory + ( createTempDataDirectory, + createTempSocketDirectory, + removeDirectoryIfExists, + resolveDirectory, + retryRemoveDirectory, + ) +import EphemeralPg.Internal.Except (liftE, onError, runStartup) +import EphemeralPg.Internal.Port (findFreePort) +import EphemeralPg.Process (getCurrentUser) +import EphemeralPg.Process.CreateDb (runCreateDb) +import EphemeralPg.Process.InitDb (runInitDb, writePostgresConf) +import EphemeralPg.Process.Postgres (startPostgres, stopPostgres) + +-- | Create a temporary database with default configuration, run an action, then clean up. +-- +-- This is the recommended way to use ephemeral-pg. The database is +-- guaranteed to be stopped and cleaned up even if an exception is thrown. +-- +-- @ +-- result <- 'with' $ \\db -> do +-- conn <- Connection.acquire ('connectionSettings' db) +-- -- Use the connection... +-- @ +with :: (Database -> IO a) -> IO (Either StartError a) +with = withConfig defaultConfig + +-- | Like 'with' but with custom configuration. +-- +-- @ +-- let config = 'defaultConfig' { configDatabaseName = "testdb" } +-- 'withConfig' config $ \\db -> do +-- -- Use the database... +-- @ +withConfig :: Config -> (Database -> IO a) -> IO (Either StartError a) +withConfig config action = mask $ \restore -> do + result <- start config + case result of + Left err -> pure (Left err) + Right db -> do + a <- restore (action db) `onException` stop db + stop db + pure (Right a) + +-- | Start a temporary database. +-- +-- You are responsible for calling 'stop' when done. Prefer 'with' or +-- 'withConfig' when possible. +-- +-- @ +-- db <- 'start' 'defaultConfig' +-- case db of +-- Right database -> do +-- -- Use database... +-- 'stop' database +-- Left err -> handleError err +-- @ +start :: Config -> IO (Either StartError Database) +start config = runStartup $ do + let mTempRoot = getLast (configTemporaryRoot config) + + -- Create data directory + (dataDir, dataDirIsTemp) <- + liftE $ + resolveDirectory + (configDataDirectory config) + mTempRoot + "data" + createTempDataDirectory + + -- Create socket directory + (socketDir, socketDirIsTemp) <- + liftE + ( resolveDirectory + (configSocketDirectory config) + mTempRoot + "socket" + createTempSocketDirectory + ) + `onError` when dataDirIsTemp (removeDirectoryIfExists dataDir) + + -- Get port + p <- + liftE (getPort config) + `onError` cleanup dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Get username + username <- liftIO $ getUsername config + + -- Run initdb + () <- + liftE (runInitDb config dataDir) + `onError` cleanup dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Start postgres + pgProcess <- + liftE (startPostgres config dataDir socketDir p username) + `onError` cleanup dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Create database + let dbName = configDatabaseName config + () <- + liftE (runCreateDb config socketDir p username dbName) + `onError` do + _ <- stopPostgres pgProcess ShutdownImmediate 5 + cleanup dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Build cleanup action + let cleanupAction = do + when dataDirIsTemp $ do + -- Use retry to handle pg_stat race + _ <- retryRemoveDirectory dataDir 5 100000 + pure () + when socketDirIsTemp $ + removeDirectoryIfExists socketDir + + pure $ + Database + { dbDataDirectory = dataDir, + dbSocketDirectory = socketDir, + dbPort = p, + dbDatabaseName = dbName, + dbUser = username, + dbPassword = configPassword config, + dbProcess = pgProcess, + dbCleanup = cleanupAction, + dbDataDirIsTemp = dataDirIsTemp, + dbSocketDirIsTemp = socketDirIsTemp + } + where + cleanup :: Bool -> FilePath -> Bool -> FilePath -> IO () + cleanup dataDirIsTemp dataDir socketDirIsTemp socketDir = do + when dataDirIsTemp $ removeDirectoryIfExists dataDir + when socketDirIsTemp $ removeDirectoryIfExists socketDir + +-- | Get port from config or find a free one. +getPort :: Config -> IO (Either StartError Word16) +getPort config = case getLast (configPort config) of + Just p -> pure $ Right p + Nothing -> findFreePort + +-- | Get username from config or current user. +getUsername :: Config -> IO Text +getUsername config = case configUser config of + "" -> getCurrentUser + u -> pure u + +-- | Stop a database and clean up resources. +-- +-- This sends SIGTERM to postgres and waits for graceful shutdown, +-- then removes temporary directories. +-- +-- Safe to call multiple times (subsequent calls are no-ops). +stop :: Database -> IO () +stop db = do + let mode = ShutdownGraceful + let timeoutSecs = defaultShutdownTimeoutSeconds + + -- Stop postgres + _ <- stopPostgres (dbProcess db) mode timeoutSecs + + -- Run cleanup (removes temp directories) + dbCleanup db + +-- | Restart a database. +-- +-- This stops the postgres server and starts it again, returning a new +-- 'Database' handle with the updated process information. +-- +-- The data directory and all database contents are preserved. +-- This is useful for testing scenarios that require a server restart, +-- such as configuration changes that require a restart to take effect. +-- +-- @ +-- db <- 'start' 'defaultConfig' +-- case db of +-- Right database -> do +-- -- Use database... +-- newDb <- 'restart' database +-- case newDb of +-- Right database' -> -- Use restarted database... +-- Left err -> handleError err +-- Left err -> handleError err +-- @ +restart :: Database -> IO (Either StartError Database) +restart db = runStartup $ do + -- Stop postgres gracefully + liftIO $ do + _ <- stopPostgres (dbProcess db) ShutdownGraceful defaultShutdownTimeoutSeconds + pure () + + -- Start postgres again with the same configuration + newProcess <- + liftE $ + startPostgres + defaultConfig + (dbDataDirectory db) + (dbSocketDirectory db) + (dbPort db) + (dbUser db) + + pure $ db {dbProcess = newProcess} + +-- | Like 'with' but uses initdb caching for faster startup. +-- +-- The first invocation runs initdb and caches the result. +-- Subsequent invocations copy from the cache (using CoW if available). +-- +-- @ +-- result <- 'withCached' $ \\db -> do +-- conn <- Connection.acquire ('connectionSettings' db) +-- -- Use the connection... +-- @ +withCached :: (Database -> IO a) -> IO (Either StartError a) +withCached = withCachedConfig defaultConfig defaultCacheConfig + +-- | Like 'withCached' but with custom configuration. +withCachedConfig :: Config -> CacheConfig -> (Database -> IO a) -> IO (Either StartError a) +withCachedConfig config cacheConfig action = mask $ \restore -> do + result <- startCached config cacheConfig + case result of + Left err -> pure (Left err) + Right db -> do + a <- restore (action db) `onException` stop db + stop db + pure (Right a) + +-- | Start a temporary database using initdb caching. +-- +-- If caching is enabled and a cache exists, the data directory is copied +-- from the cache. Otherwise, initdb is run and the result is cached. +startCached :: Config -> CacheConfig -> IO (Either StartError Database) +startCached config cacheConfig + | not (cacheConfigEnabled cacheConfig) = start config + | otherwise = do + -- Get cache key + keyResult <- getCacheKey config + case keyResult of + Left _err -> + -- Can't determine cache key, fall back to non-cached start + start config + Right cacheKey -> do + -- Check if cache exists + cached <- isCached cacheKey (cacheConfigRoot cacheConfig) + if cached + then startFromCache config cacheConfig cacheKey + else startAndCache config cacheConfig cacheKey + +-- | Start from an existing cache. +startFromCache :: Config -> CacheConfig -> CacheKey -> IO (Either StartError Database) +startFromCache config cacheConfig cacheKey = do + let mTempRoot = getLast (configTemporaryRoot config) + + -- Create temporary data directory (to get the path) + dataResult <- createTempDataDirectory mTempRoot + case dataResult of + Left err -> pure $ Left err + Right (dataDir, dataDirIsTemp) -> do + -- Remove the directory so cp can create it fresh + -- (otherwise cp -cR creates nested directories on macOS) + removeDirectoryIfExists dataDir + + -- Restore from cache + restoreResult <- restoreFromCache cacheKey dataDir (cacheConfigRoot cacheConfig) + case restoreResult of + Left _err -> do + -- Cache restore failed, fall back to non-cached start + removeDirectoryIfExists dataDir + start config + Right () -> do + -- Clean up any runtime files from the cache (postmaster.pid, etc.) + cleanupRuntimeFiles dataDir + + -- Write postgresql.conf (cache doesn't include our custom settings) + writePostgresConf config dataDir + + -- Continue with normal startup from the restored data directory + continueStartup config dataDir dataDirIsTemp + +-- | Start normally and cache the result. +-- Cache is created after initdb but before postgres starts. +startAndCache :: Config -> CacheConfig -> CacheKey -> IO (Either StartError Database) +startAndCache config cacheConfig cacheKey = do + let mTempRoot = getLast (configTemporaryRoot config) + + -- Create data directory + dataResult <- createTempDataDirectory mTempRoot + case dataResult of + Left err -> pure $ Left err + Right (dataDir, dataDirIsTemp) -> do + -- Run initdb + initResult <- runInitDb config dataDir + case initResult of + Left err -> do + when dataDirIsTemp $ removeDirectoryIfExists dataDir + pure $ Left err + Right () -> do + -- Cache the data directory NOW (before postgres starts) + -- This ensures the cache contains only clean initdb output + when dataDirIsTemp $ do + _ <- createCache cacheKey dataDir (cacheConfigRoot cacheConfig) + pure () + + -- Continue with normal startup from the initialized data directory + continueStartup config dataDir dataDirIsTemp + +-- | Continue startup from an existing data directory. +continueStartup :: Config -> FilePath -> Bool -> IO (Either StartError Database) +continueStartup config dataDir dataDirIsTemp = runStartup $ do + let mTempRoot = getLast (configTemporaryRoot config) + + -- Create socket directory + (socketDir, socketDirIsTemp) <- + liftE + ( resolveDirectory + (configSocketDirectory config) + mTempRoot + "socket" + createTempSocketDirectory + ) + `onError` when dataDirIsTemp (removeDirectoryIfExists dataDir) + + -- Get port + p <- + liftE (getPort config) + `onError` cleanupDirs dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Get username + username <- liftIO $ getUsername config + + -- Start postgres (initdb already done) + pgProcess <- + liftE (startPostgres config dataDir socketDir p username) + `onError` cleanupDirs dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Create database + let dbName = configDatabaseName config + () <- + liftE (runCreateDb config socketDir p username dbName) + `onError` do + _ <- stopPostgres pgProcess ShutdownImmediate 5 + cleanupDirs dataDirIsTemp dataDir socketDirIsTemp socketDir + + -- Build cleanup action + let cleanupAction = do + when dataDirIsTemp $ do + _ <- retryRemoveDirectory dataDir 5 100000 + pure () + when socketDirIsTemp $ + removeDirectoryIfExists socketDir + + pure $ + Database + { dbDataDirectory = dataDir, + dbSocketDirectory = socketDir, + dbPort = p, + dbDatabaseName = dbName, + dbUser = username, + dbPassword = configPassword config, + dbProcess = pgProcess, + dbCleanup = cleanupAction, + dbDataDirIsTemp = dataDirIsTemp, + dbSocketDirIsTemp = socketDirIsTemp + } + where + cleanupDirs :: Bool -> FilePath -> Bool -> FilePath -> IO () + cleanupDirs dataDirIsTemp' dataDir' socketDirIsTemp' socketDir' = do + when dataDirIsTemp' $ removeDirectoryIfExists dataDir' + when socketDirIsTemp' $ removeDirectoryIfExists socketDir' diff --git a/src/EphemeralPg/Config.hs b/src/EphemeralPg/Config.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Config.hs @@ -0,0 +1,271 @@ +-- | Configuration for temporary PostgreSQL databases. +-- +-- This module provides the 'Config' type and predefined configurations +-- for common use cases. +-- +-- Configurations can be combined using the 'Semigroup' instance: +-- +-- @ +-- myConfig = 'defaultConfig' <> mempty { configDatabaseName = "testdb" } +-- @ +module EphemeralPg.Config + ( -- * Configuration type + Config (..), + DirectoryConfig (..), + ShutdownMode (..), + + -- * Predefined configurations + defaultConfig, + verboseConfig, + autoExplainConfig, + + -- * Default values + defaultPostgresSettings, + defaultInitDbArgs, + defaultConnectionTimeoutSeconds, + defaultShutdownTimeoutSeconds, + + -- * Socket path limits + maxSocketPathLength, + ) +where + +import Data.Monoid (Last (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Word (Word16) +import System.IO (Handle) + +-- | Configuration for a directory. +data DirectoryConfig + = -- | Create a temporary directory that is deleted on cleanup (default) + DirectoryTemporary + | -- | Use a specific path that is NOT deleted on cleanup + DirectoryPermanent FilePath + deriving stock (Eq, Show) + +-- | How to shut down the PostgreSQL server. +data ShutdownMode + = -- | Send SIGTERM, wait for clean shutdown (default) + ShutdownGraceful + | -- | Send SIGINT for fast shutdown (rolls back active transactions) + ShutdownFast + | -- | Send SIGQUIT for immediate shutdown + ShutdownImmediate + deriving stock (Eq, Show) + +-- | Configuration for a temporary PostgreSQL database. +-- +-- Combine configurations using '<>'. Later values take precedence +-- for 'Last' fields. +data Config = Config + { -- | Port to listen on. 'Nothing' means auto-select a free port. + configPort :: Last Word16, + -- | Database name to create. + configDatabaseName :: Text, + -- | Username for the database. + configUser :: Text, + -- | Optional password. + configPassword :: Maybe Text, + -- | Data directory configuration. + configDataDirectory :: DirectoryConfig, + -- | Socket directory configuration. + configSocketDirectory :: DirectoryConfig, + -- | Root directory for temporary files. + configTemporaryRoot :: Last FilePath, + -- | postgresql.conf settings. + configPostgresSettings :: [(Text, Text)], + -- | Additional arguments for initdb. + configInitDbArgs :: [Text], + -- | Additional arguments for postgres server. + configPostgresArgs :: [Text], + -- | Additional arguments for createdb. + configCreateDbArgs :: [Text], + -- | Timeout waiting for postgres to accept connections (seconds). + configConnectionTimeoutSeconds :: Last Int, + -- | Timeout waiting for graceful shutdown (seconds). + configShutdownTimeoutSeconds :: Last Int, + -- | How to shut down postgres. + configShutdownMode :: Last ShutdownMode, + -- | Handle for postgres stdout (Nothing = discard). + configStdout :: Last (Maybe Handle), + -- | Handle for postgres stderr (Nothing = discard). + configStderr :: Last (Maybe Handle) + } + deriving stock (Show) + +instance Semigroup Config where + a <> b = + Config + { configPort = configPort a <> configPort b, + configDatabaseName = + if configDatabaseName b == "" + then configDatabaseName a + else configDatabaseName b, + configUser = + if configUser b == "" + then configUser a + else configUser b, + configPassword = configPassword b <|> configPassword a, + configDataDirectory = combineDir (configDataDirectory a) (configDataDirectory b), + configSocketDirectory = combineDir (configSocketDirectory a) (configSocketDirectory b), + configTemporaryRoot = configTemporaryRoot a <> configTemporaryRoot b, + configPostgresSettings = configPostgresSettings a <> configPostgresSettings b, + configInitDbArgs = configInitDbArgs a <> configInitDbArgs b, + configPostgresArgs = configPostgresArgs a <> configPostgresArgs b, + configCreateDbArgs = configCreateDbArgs a <> configCreateDbArgs b, + configConnectionTimeoutSeconds = configConnectionTimeoutSeconds a <> configConnectionTimeoutSeconds b, + configShutdownTimeoutSeconds = configShutdownTimeoutSeconds a <> configShutdownTimeoutSeconds b, + configShutdownMode = configShutdownMode a <> configShutdownMode b, + configStdout = configStdout a <> configStdout b, + configStderr = configStderr a <> configStderr b + } + where + combineDir DirectoryTemporary d = d + combineDir d DirectoryTemporary = d + combineDir _ d = d + + (<|>) :: Maybe a -> Maybe a -> Maybe a + (<|>) Nothing x = x + (<|>) x _ = x + +instance Monoid Config where + mempty = + Config + { configPort = Last Nothing, + configDatabaseName = "", + configUser = "", + configPassword = Nothing, + configDataDirectory = DirectoryTemporary, + configSocketDirectory = DirectoryTemporary, + configTemporaryRoot = Last Nothing, + configPostgresSettings = [], + configInitDbArgs = [], + configPostgresArgs = [], + configCreateDbArgs = [], + configConnectionTimeoutSeconds = Last Nothing, + configShutdownTimeoutSeconds = Last Nothing, + configShutdownMode = Last Nothing, + configStdout = Last Nothing, + configStderr = Last Nothing + } + +-- | Maximum socket path length for Unix domain sockets. +-- +-- This is 104 bytes on macOS and 108 bytes on Linux. +-- We use the more conservative macOS limit. +maxSocketPathLength :: Int +maxSocketPathLength = 104 + +-- | Default connection timeout in seconds. +defaultConnectionTimeoutSeconds :: Int +defaultConnectionTimeoutSeconds = 60 + +-- | Default shutdown timeout in seconds. +defaultShutdownTimeoutSeconds :: Int +defaultShutdownTimeoutSeconds = 30 + +-- | Default PostgreSQL settings optimized for testing. +-- +-- These settings disable durability features for maximum performance +-- and are NOT suitable for production use. +defaultPostgresSettings :: [(Text, Text)] +defaultPostgresSettings = + [ ("shared_buffers", "'12MB'"), + ("fsync", "'off'"), + ("synchronous_commit", "'off'"), + ("full_page_writes", "'off'"), + ("random_page_cost", "'1.0'"), + ("log_min_messages", "'PANIC'"), + ("log_min_error_statement", "'PANIC'"), + ("log_statement", "'none'"), + ("client_min_messages", "'ERROR'"), + ("wal_level", "'minimal'"), + ("max_wal_senders", "0"), + ("archive_mode", "'off'"), + -- Listen on localhost for IPv4 + ("listen_addresses", "'127.0.0.1'") + ] + +-- | Default initdb arguments. +defaultInitDbArgs :: [Text] +defaultInitDbArgs = + [ "--no-sync", + "--encoding=UTF8", + "--no-locale", + "--auth=trust" + ] + +-- | High-performance configuration suitable for testing. +-- +-- This configuration: +-- +-- * Auto-selects a free port +-- * Uses temporary directories (auto-cleaned) +-- * Disables durability features for speed +-- * Uses the current system user +-- * Creates a database named "postgres" +defaultConfig :: Config +defaultConfig = + Config + { configPort = Last Nothing, + configDatabaseName = "postgres", + configUser = "", -- Will use current user + configPassword = Nothing, + configDataDirectory = DirectoryTemporary, + configSocketDirectory = DirectoryTemporary, + configTemporaryRoot = Last Nothing, + configPostgresSettings = defaultPostgresSettings, + configInitDbArgs = defaultInitDbArgs, + configPostgresArgs = [], + configCreateDbArgs = [], + configConnectionTimeoutSeconds = Last (Just defaultConnectionTimeoutSeconds), + configShutdownTimeoutSeconds = Last (Just defaultShutdownTimeoutSeconds), + configShutdownMode = Last (Just ShutdownGraceful), + configStdout = Last (Just Nothing), -- Discard by default + configStderr = Last (Just Nothing) -- Discard by default + } + +-- | Configuration with verbose PostgreSQL logging. +-- +-- Combine with 'defaultConfig': +-- +-- @ +-- myConfig = 'defaultConfig' <> 'verboseConfig' +-- @ +verboseConfig :: Config +verboseConfig = + mempty + { configPostgresSettings = + [ ("log_min_duration_statement", "'0'"), + ("log_min_messages", "'DEBUG1'"), + ("log_min_error_statement", "'DEBUG1'"), + ("log_checkpoints", "'on'"), + ("log_connections", "'on'"), + ("log_disconnections", "'on'"), + ("log_lock_waits", "'on'"), + ("log_temp_files", "'0'"), + ("log_autovacuum_min_duration", "'0'"), + ("log_line_prefix", "'%t [%p]: '") + ] + } + +-- | Configuration that enables auto_explain for query analysis. +-- +-- Logs query plans for queries exceeding the specified duration. +-- +-- @ +-- myConfig = 'defaultConfig' <> 'autoExplainConfig' 100 -- Log queries > 100ms +-- @ +autoExplainConfig :: Int -> Config +autoExplainConfig minDurationMs = + mempty + { configPostgresSettings = + [ ("shared_preload_libraries", "'auto_explain'"), + ("auto_explain.log_min_duration", "'" <> ms <> "'"), + ("auto_explain.log_analyze", "'on'") + ] + } + where + ms :: Text + ms = T.pack (show minDurationMs) <> "ms" diff --git a/src/EphemeralPg/Database.hs b/src/EphemeralPg/Database.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Database.hs @@ -0,0 +1,134 @@ +-- | Database handle and connection utilities. +-- +-- This module provides the 'Database' type representing a running +-- PostgreSQL instance, and functions to extract connection information. +module EphemeralPg.Database + ( -- * Database handle + Database (..), + PostgresProcess (..), + + -- * Connection information + connectionSettings, + connectionString, + + -- * Accessors + dataDirectory, + socketDirectory, + port, + databaseName, + user, + ) +where + +import Data.Text (Text) +import Data.Text qualified as T +import Data.Word (Word16) +import Hasql.Connection.Settings qualified as Settings +import System.Posix.Types (CPid) +import System.Process.Typed (Process) + +-- | Handle to a running PostgreSQL server process. +data PostgresProcess = PostgresProcess + { -- | The typed-process handle + postgresProcess :: Process () () (), + -- | Process ID for signal handling + postgresPid :: CPid + } + +instance Show PostgresProcess where + show PostgresProcess {postgresPid} = + "PostgresProcess { pid = " <> show postgresPid <> " }" + +-- | A running PostgreSQL database instance. +-- +-- Obtain via 'EphemeralPg.start' or 'EphemeralPg.with'. The database is +-- ready to accept connections immediately after creation. +data Database = Database + { -- | Path to the PostgreSQL data directory + dbDataDirectory :: FilePath, + -- | Path to the Unix socket directory + dbSocketDirectory :: FilePath, + -- | Port number the server is listening on + dbPort :: Word16, + -- | Name of the database + dbDatabaseName :: Text, + -- | Username for connections + dbUser :: Text, + -- | Optional password + dbPassword :: Maybe Text, + -- | Handle to the running postgres process + dbProcess :: PostgresProcess, + -- | Cleanup action (removes temp directories if needed) + dbCleanup :: IO (), + -- | Whether data directory is temporary + dbDataDirIsTemp :: Bool, + -- | Whether socket directory is temporary + dbSocketDirIsTemp :: Bool + } + +instance Show Database where + show db = + "Database { port = " + <> show (dbPort db) + <> ", database = " + <> show (dbDatabaseName db) + <> ", user = " + <> show (dbUser db) + <> ", dataDir = " + <> show (dbDataDirectory db) + <> " }" + +-- | Get hasql connection settings for this database. +-- +-- Example: +-- +-- @ +-- db <- 'EphemeralPg.start' 'EphemeralPg.defaultConfig' +-- case db of +-- Right database -> do +-- Right conn <- Connection.acquire ('connectionSettings' database) +-- -- use connection... +-- Left err -> handleError err +-- @ +connectionSettings :: Database -> Settings.Settings +connectionSettings db = + mconcat + [ Settings.hostAndPort (T.pack $ dbSocketDirectory db) (dbPort db), + Settings.dbname (dbDatabaseName db), + Settings.user (dbUser db), + maybe mempty Settings.password (dbPassword db) + ] + +-- | Get a libpq-compatible connection string. +-- +-- Format: @host=/path/to/socket port=5432 dbname=postgres user=username@ +-- +-- Useful for interoperating with other PostgreSQL libraries. +connectionString :: Database -> Text +connectionString db = + T.unwords + [ "host=" <> T.pack (dbSocketDirectory db), + "port=" <> T.pack (show $ dbPort db), + "dbname=" <> dbDatabaseName db, + "user=" <> dbUser db + ] + +-- | Get the data directory path. +dataDirectory :: Database -> FilePath +dataDirectory = dbDataDirectory + +-- | Get the socket directory path. +socketDirectory :: Database -> FilePath +socketDirectory = dbSocketDirectory + +-- | Get the port number. +port :: Database -> Word16 +port = dbPort + +-- | Get the database name. +databaseName :: Database -> Text +databaseName = dbDatabaseName + +-- | Get the username. +user :: Database -> Text +user = dbUser diff --git a/src/EphemeralPg/Dump.hs b/src/EphemeralPg/Dump.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Dump.hs @@ -0,0 +1,292 @@ +-- | Database dump and restore functionality. +-- +-- This module provides wrappers around pg_dump and psql for +-- creating and restoring SQL dumps of ephemeral databases. +-- +-- = Example Usage +-- +-- @ +-- result <- Pg.'with' $ \\db -> do +-- -- Set up some test data +-- runMigrations db +-- +-- -- Create a dump +-- dumpResult <- Pg.'dump' db \"/path/to/dump.sql\" +-- +-- -- Later, restore the dump to a new database +-- Pg.'with' $ \\newDb -> do +-- Pg.'restore' newDb \"/path/to/dump.sql\" +-- @ +module EphemeralPg.Dump + ( -- * Dump Operations + dump, + dumpToText, + + -- * Restore Operations + restore, + restoreFromText, + + -- * Configuration + DumpFormat (..), + DumpOptions (..), + defaultDumpOptions, + ) +where + +import Control.Exception (SomeException, try) +import Data.ByteString.Lazy qualified as LBS +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding qualified as T +import EphemeralPg.Database (Database (..)) +import EphemeralPg.Process (findExecutable) +import System.Exit (ExitCode (..)) +import System.Process.Typed + ( byteStringOutput, + proc, + readProcess, + setEnv, + setStderr, + setStdout, + ) + +-- | Output format for pg_dump. +data DumpFormat + = -- | Plain SQL text format + DumpPlain + | -- | Custom archive format (for pg_restore) + DumpCustom + | -- | Directory format + DumpDirectory + | -- | Tar archive format + DumpTar + deriving stock (Eq, Show) + +-- | Options for pg_dump. +data DumpOptions = DumpOptions + { -- | Output format (default: DumpPlain) + dumpFormat :: DumpFormat, + -- | Include CREATE DATABASE command + dumpCreateDb :: Bool, + -- | Schema only (no data) + dumpSchemaOnly :: Bool, + -- | Data only (no schema) + dumpDataOnly :: Bool, + -- | Include specific tables (empty = all) + dumpTables :: [Text], + -- | Exclude specific tables + dumpExcludeTables :: [Text] + } + deriving stock (Eq, Show) + +-- | Default dump options (plain SQL, all tables, schema + data). +defaultDumpOptions :: DumpOptions +defaultDumpOptions = + DumpOptions + { dumpFormat = DumpPlain, + dumpCreateDb = False, + dumpSchemaOnly = False, + dumpDataOnly = False, + dumpTables = [], + dumpExcludeTables = [] + } + +-- | Dump a database to a file. +dump :: Database -> FilePath -> DumpOptions -> IO (Either Text ()) +dump db outputPath opts = do + mPgDump <- findExecutable "pg_dump" + case mPgDump of + Nothing -> pure $ Left "pg_dump not found in PATH" + Just pgDumpPath -> do + let args = buildDumpArgs db opts <> ["-f", T.pack outputPath] + env = buildEnv db + processConfig = + proc pgDumpPath (map T.unpack args) + & setStdout byteStringOutput + & setStderr byteStringOutput + & setEnv env + result <- try $ readProcess processConfig + case result of + Left (ex :: SomeException) -> + pure $ Left $ "pg_dump failed: " <> T.pack (show ex) + Right (exitCode, _stdout, stderr) -> + case exitCode of + ExitSuccess -> pure $ Right () + ExitFailure code -> + pure $ + Left $ + "pg_dump failed with code " + <> T.pack (show code) + <> ": " + <> T.decodeUtf8Lenient (LBS.toStrict stderr) + where + (&) = flip ($) + +-- | Dump a database and return the dump as text. +-- +-- Only works with DumpPlain format. +dumpToText :: Database -> DumpOptions -> IO (Either Text Text) +dumpToText db opts = do + if dumpFormat opts /= DumpPlain + then pure $ Left "dumpToText only works with DumpPlain format" + else do + mPgDump <- findExecutable "pg_dump" + case mPgDump of + Nothing -> pure $ Left "pg_dump not found in PATH" + Just pgDumpPath -> do + let args = buildDumpArgs db opts + env = buildEnv db + processConfig = + proc pgDumpPath (map T.unpack args) + & setStdout byteStringOutput + & setStderr byteStringOutput + & setEnv env + result <- try $ readProcess processConfig + case result of + Left (ex :: SomeException) -> + pure $ Left $ "pg_dump failed: " <> T.pack (show ex) + Right (exitCode, stdout, stderr) -> + case exitCode of + ExitSuccess -> pure $ Right $ T.decodeUtf8Lenient $ LBS.toStrict stdout + ExitFailure code -> + pure $ + Left $ + "pg_dump failed with code " + <> T.pack (show code) + <> ": " + <> T.decodeUtf8Lenient (LBS.toStrict stderr) + where + (&) = flip ($) + +-- | Restore a database from a dump file. +restore :: Database -> FilePath -> IO (Either Text ()) +restore db inputPath = do + mPsql <- findExecutable "psql" + case mPsql of + Nothing -> pure $ Left "psql not found in PATH" + Just psqlPath -> do + let args = + [ "-h", + T.pack $ dbSocketDirectory db, + "-p", + T.pack $ show $ dbPort db, + "-U", + dbUser db, + "-d", + dbDatabaseName db, + "-f", + T.pack inputPath, + "-v", + "ON_ERROR_STOP=1" + ] + env = buildEnv db + processConfig = + proc psqlPath (map T.unpack args) + & setStdout byteStringOutput + & setStderr byteStringOutput + & setEnv env + result <- try $ readProcess processConfig + case result of + Left (ex :: SomeException) -> + pure $ Left $ "psql restore failed: " <> T.pack (show ex) + Right (exitCode, _stdout, stderr) -> + case exitCode of + ExitSuccess -> pure $ Right () + ExitFailure code -> + pure $ + Left $ + "psql restore failed with code " + <> T.pack (show code) + <> ": " + <> T.decodeUtf8Lenient (LBS.toStrict stderr) + where + (&) = flip ($) + +-- | Restore a database from SQL text. +restoreFromText :: Database -> Text -> IO (Either Text ()) +restoreFromText db sqlText = do + mPsql <- findExecutable "psql" + case mPsql of + Nothing -> pure $ Left "psql not found in PATH" + Just psqlPath -> do + let args = + [ "-h", + T.pack $ dbSocketDirectory db, + "-p", + T.pack $ show $ dbPort db, + "-U", + dbUser db, + "-d", + dbDatabaseName db, + "-c", + sqlText, + "-v", + "ON_ERROR_STOP=1" + ] + env = buildEnv db + processConfig = + proc psqlPath (map T.unpack args) + & setStdout byteStringOutput + & setStderr byteStringOutput + & setEnv env + result <- try $ readProcess processConfig + case result of + Left (ex :: SomeException) -> + pure $ Left $ "psql restore failed: " <> T.pack (show ex) + Right (exitCode, _stdout, stderr) -> + case exitCode of + ExitSuccess -> pure $ Right () + ExitFailure code -> + pure $ + Left $ + "psql restore failed with code " + <> T.pack (show code) + <> ": " + <> T.decodeUtf8Lenient (LBS.toStrict stderr) + where + (&) = flip ($) + +-- | Build pg_dump command line arguments. +buildDumpArgs :: Database -> DumpOptions -> [Text] +buildDumpArgs db opts = + [ "-h", + T.pack $ dbSocketDirectory db, + "-p", + T.pack $ show $ dbPort db, + "-U", + dbUser db, + "-d", + dbDatabaseName db + ] + <> formatArg + <> createDbArg + <> schemaOnlyArg + <> dataOnlyArg + <> tableArgs + <> excludeTableArgs + where + formatArg = case dumpFormat opts of + DumpPlain -> ["-Fp"] + DumpCustom -> ["-Fc"] + DumpDirectory -> ["-Fd"] + DumpTar -> ["-Ft"] + + createDbArg = + if dumpCreateDb opts then ["-C"] else [] + + schemaOnlyArg = + if dumpSchemaOnly opts then ["-s"] else [] + + dataOnlyArg = + if dumpDataOnly opts then ["-a"] else [] + + tableArgs = + concatMap (\t -> ["-t", t]) (dumpTables opts) + + excludeTableArgs = + concatMap (\t -> ["-T", t]) (dumpExcludeTables opts) + +-- | Build environment variables for pg commands. +buildEnv :: Database -> [(String, String)] +buildEnv db = + [("PGPASSWORD", T.unpack pw) | Just pw <- [dbPassword db]] diff --git a/src/EphemeralPg/Error.hs b/src/EphemeralPg/Error.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Error.hs @@ -0,0 +1,249 @@ +{-# LANGUAGE NoFieldSelectors #-} +{-# OPTIONS_GHC -Wno-partial-fields #-} + +-- | Error types for ephemeral-pg. +-- +-- This module provides a rich error hierarchy for diagnosing failures +-- when starting or stopping temporary PostgreSQL databases. +-- +-- Note: We use record syntax with sum types for structured error data. +-- The NoFieldSelectors extension prevents generation of partial field +-- accessor functions. +module EphemeralPg.Error + ( -- * Top-level errors + StartError (..), + StopError (..), + + -- * Specific error types + InitDbError (..), + PostgresError (..), + CreateDbError (..), + ConfigError (..), + ResourceError (..), + TimeoutError (..), + + -- * Error rendering + renderStartError, + renderStopError, + ) +where + +import Control.Exception (Exception) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Word (Word16) +import System.Exit (ExitCode (..)) + +-- | Top-level error when starting a database. +data StartError + = -- | initdb failed to initialize the data directory + InitDbError InitDbError + | -- | PostgreSQL server failed to start + PostgresStartError PostgresError + | -- | createdb failed to create the database + CreateDbError CreateDbError + | -- | Configuration is invalid + ConfigError ConfigError + | -- | System resource error (permissions, disk space, etc.) + ResourceError ResourceError + | -- | Timeout waiting for database to be ready + TimeoutError TimeoutError + deriving stock (Eq, Show) + deriving anyclass (Exception) + +-- | Errors when stopping a database. +data StopError + = -- | Failed to send shutdown signal + ShutdownSignalFailed Int Text + | -- | Graceful shutdown timed out, had to force kill + ShutdownTimedOut Int + | -- | Failed to clean up directories + CleanupFailed FilePath Text + deriving stock (Eq, Show) + deriving anyclass (Exception) + +-- | Errors from initdb execution. +data InitDbError + = -- | initdb executable not found in PATH + InitDbNotFound + | -- | PostgreSQL version too old + InitDbVersionMismatch + { initDbExpected :: Text, + initDbActual :: Text + } + | -- | initdb process failed + InitDbFailed + { initDbExitCode :: ExitCode, + initDbStdout :: Text, + initDbStderr :: Text, + initDbCommand :: Text + } + | -- | Cannot write to data directory + InitDbPermissionDenied FilePath + deriving stock (Eq, Show) + +-- | Errors from PostgreSQL server operations. +data PostgresError + = -- | postgres executable not found in PATH + PostgresNotFound + | -- | Server process exited unexpectedly + PostgresStartFailed + { pgExitCode :: ExitCode, + pgStdout :: Text, + pgStderr :: Text, + pgCommand :: Text + } + | -- | Could not connect to verify server is running + PostgresConnectionFailed Text + | -- | Server version doesn't match requirements + PostgresVersionMismatch + { pgExpected :: Text, + pgActual :: Text + } + deriving stock (Eq, Show) + +-- | Errors from createdb execution. +data CreateDbError + = -- | createdb executable not found in PATH + CreateDbNotFound + | -- | createdb process failed + CreateDbFailed + { createDbExitCode :: ExitCode, + createDbStdout :: Text, + createDbStderr :: Text, + createDbCommand :: Text, + createDbName :: Text + } + | -- | Database with this name already exists + DatabaseAlreadyExists Text + deriving stock (Eq, Show) + +-- | Configuration validation errors. +data ConfigError + = -- | Port number is invalid or unavailable + InvalidPort Word16 Text + | -- | Database name is invalid + InvalidDatabaseName Text Text + | -- | Username is invalid + InvalidUser Text Text + | -- | Data directory path is invalid + InvalidDataDirectory FilePath Text + | -- | Socket directory path is invalid + InvalidSocketDirectory FilePath Text + | -- | Socket path exceeds Unix domain socket limit + SocketPathTooLong + { socketPath :: FilePath, + socketPathLength :: Int, + socketPathMaxLength :: Int + } + | -- | Configuration options conflict with each other + ConflictingConfig Text + deriving stock (Eq, Show) + +-- | System resource errors. +data ResourceError + = -- | Cannot access a required path + PermissionDenied FilePath Text + | -- | Cannot create a directory + DirectoryCreationFailed FilePath Text + | -- | Directory should be empty but isn't + DirectoryNotEmpty FilePath + | -- | Not enough disk space + DiskSpaceExhausted FilePath + | -- | Too many open files + FileDescriptorLimit + | -- | Requested port is already in use + PortInUse Word16 + | -- | Could not find a free port + PortAllocationFailed Text + deriving stock (Eq, Show) + +-- | Timeout errors. +data TimeoutError + = -- | Timed out waiting for PostgreSQL to accept connections + ConnectionTimeout + { timeoutDurationSeconds :: Int, + timeoutHost :: Text, + timeoutPort :: Word16 + } + | -- | Timed out waiting for graceful shutdown + ShutdownTimeout + { shutdownTimeoutSeconds :: Int, + shutdownPid :: Int + } + deriving stock (Eq, Show) + +-- | Render a 'StartError' for display to users. +renderStartError :: StartError -> Text +renderStartError = \case + InitDbError InitDbNotFound -> + "initdb not found. Please ensure PostgreSQL is installed and initdb is in your PATH.\n" + <> "On Debian/Ubuntu, you may need to add /usr/lib/postgresql//bin to PATH." + InitDbError (InitDbVersionMismatch expected actual) -> + "PostgreSQL version mismatch: expected " <> expected <> " but found " <> actual + InitDbError (InitDbFailed _code _stdout stderr cmd) -> + "initdb failed:\n Command: " <> cmd <> "\n Error: " <> stderr + InitDbError (InitDbPermissionDenied path) -> + "Permission denied: cannot write to " <> T.pack path + PostgresStartError PostgresNotFound -> + "postgres not found. Please ensure PostgreSQL is installed and postgres is in your PATH." + PostgresStartError (PostgresStartFailed _code _stdout stderr cmd) -> + "PostgreSQL server failed to start:\n Command: " <> cmd <> "\n Error: " <> stderr + PostgresStartError (PostgresConnectionFailed msg) -> + "Could not connect to PostgreSQL: " <> msg + PostgresStartError (PostgresVersionMismatch expected actual) -> + "PostgreSQL version mismatch: expected " <> expected <> " but found " <> actual + CreateDbError CreateDbNotFound -> + "createdb not found. Please ensure PostgreSQL is installed and createdb is in your PATH." + CreateDbError (CreateDbFailed _code _stdout stderr cmd name) -> + "createdb failed for database '" <> name <> "':\n Command: " <> cmd <> "\n Error: " <> stderr + CreateDbError (DatabaseAlreadyExists name) -> + "Database '" <> name <> "' already exists" + ConfigError (InvalidPort port msg) -> + "Invalid port " <> T.pack (show port) <> ": " <> msg + ConfigError (InvalidDatabaseName name msg) -> + "Invalid database name '" <> name <> "': " <> msg + ConfigError (InvalidUser name msg) -> + "Invalid username '" <> name <> "': " <> msg + ConfigError (InvalidDataDirectory path msg) -> + "Invalid data directory '" <> T.pack path <> "': " <> msg + ConfigError (InvalidSocketDirectory path msg) -> + "Invalid socket directory '" <> T.pack path <> "': " <> msg + ConfigError (SocketPathTooLong path len maxLen) -> + "Socket path too long (" + <> T.pack (show len) + <> " > " + <> T.pack (show maxLen) + <> " bytes): " + <> T.pack path + <> "\nTry using a shorter temporary directory path." + ConfigError (ConflictingConfig msg) -> + "Conflicting configuration: " <> msg + ResourceError (PermissionDenied path msg) -> + "Permission denied for '" <> T.pack path <> "': " <> msg + ResourceError (DirectoryCreationFailed path msg) -> + "Failed to create directory '" <> T.pack path <> "': " <> msg + ResourceError (DirectoryNotEmpty path) -> + "Directory not empty: " <> T.pack path + ResourceError (DiskSpaceExhausted path) -> + "Disk space exhausted at " <> T.pack path + ResourceError FileDescriptorLimit -> + "Too many open files. Try increasing your system's file descriptor limit." + ResourceError (PortInUse port) -> + "Port " <> T.pack (show port) <> " is already in use" + ResourceError (PortAllocationFailed msg) -> + "Could not allocate a free port: " <> msg + TimeoutError (ConnectionTimeout secs _host port) -> + "Timed out after " <> T.pack (show secs) <> "s waiting for PostgreSQL on port " <> T.pack (show port) + TimeoutError (ShutdownTimeout secs pid) -> + "Timed out after " <> T.pack (show secs) <> "s waiting for process " <> T.pack (show pid) <> " to shut down" + +-- | Render a 'StopError' for display to users. +renderStopError :: StopError -> Text +renderStopError = \case + ShutdownSignalFailed pid msg -> + "Failed to send shutdown signal to process " <> T.pack (show pid) <> ": " <> msg + ShutdownTimedOut secs -> + "Graceful shutdown timed out after " <> T.pack (show secs) <> "s, process was killed" + CleanupFailed path msg -> + "Failed to clean up '" <> T.pack path <> "': " <> msg diff --git a/src/EphemeralPg/Internal/Cache.hs b/src/EphemeralPg/Internal/Cache.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Internal/Cache.hs @@ -0,0 +1,260 @@ +-- | initdb caching for fast PostgreSQL startup. +-- +-- This module provides caching of initialized PostgreSQL data directories. +-- The first startup runs initdb, subsequent startups copy from the cache. +-- +-- Cache structure: +-- +-- @ +-- ~/.ephemeral-pg/ +-- cache/ +-- \-\/ +-- data/ -- Cached initdb output +-- metadata.json -- Cache metadata +-- @ +module EphemeralPg.Internal.Cache + ( -- * Cache operations + CacheKey (..), + CacheConfig (..), + defaultCacheConfig, + getCacheKey, + getCacheDirectory, + isCached, + createCache, + restoreFromCache, + clearCache, + clearAllCaches, + cleanupRuntimeFiles, + + -- * Cache directory management + ensureCacheDirectory, + getCacheRoot, + ) +where + +import Control.Exception (SomeException, try) +import Data.ByteString.Lazy qualified as LBS +import Data.Hashable (hash) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding qualified as T +import EphemeralPg.Config (Config (..)) +import EphemeralPg.Internal.CopyOnWrite + ( CowCapability (..), + copyDirectory, + detectCowCapability, + ) +import System.Directory + ( XdgDirectory (XdgCache), + createDirectoryIfMissing, + doesDirectoryExist, + doesFileExist, + getXdgDirectory, + listDirectory, + removeDirectoryRecursive, + removeFile, + ) +import System.Exit (ExitCode (..)) +import System.FilePath (()) +import System.Process.Typed (byteStringOutput, proc, readProcess, setStderr, setStdout) + +-- | A cache key uniquely identifies a cached initdb cluster. +data CacheKey = CacheKey + { -- | PostgreSQL major version (e.g., "17") + cacheKeyPgVersion :: Text, + -- | Hash of configuration that affects initdb output + cacheKeyConfigHash :: Text + } + deriving stock (Eq, Show) + +-- | Configuration for the cache system. +data CacheConfig = CacheConfig + { -- | Root directory for cache (default: ~/.ephemeral-pg) + cacheConfigRoot :: Maybe FilePath, + -- | Copy-on-write capability (detected if Nothing) + cacheConfigCow :: Maybe CowCapability, + -- | Whether to use caching at all + cacheConfigEnabled :: Bool + } + deriving stock (Eq, Show) + +-- | Default cache configuration. +defaultCacheConfig :: CacheConfig +defaultCacheConfig = + CacheConfig + { cacheConfigRoot = Nothing, + cacheConfigCow = Nothing, + cacheConfigEnabled = True + } + +-- | Get the cache root directory. +-- +-- Uses XDG_CACHE_HOME if available, otherwise ~/.cache/ephemeral-pg +getCacheRoot :: Maybe FilePath -> IO FilePath +getCacheRoot mRoot = case mRoot of + Just root -> pure root + Nothing -> getXdgDirectory XdgCache "ephemeral-pg" + +-- | Ensure the cache directory structure exists. +ensureCacheDirectory :: Maybe FilePath -> IO FilePath +ensureCacheDirectory mRoot = do + root <- getCacheRoot mRoot + let cacheDir = root "cache" + createDirectoryIfMissing True cacheDir + pure cacheDir + +-- | Get the PostgreSQL version. +getPostgresVersion :: IO (Either Text Text) +getPostgresVersion = do + result <- try $ readProcess config + case result of + Left (ex :: SomeException) -> + pure $ Left $ "Failed to get postgres version: " <> T.pack (show ex) + Right (exitCode, stdout, _stderr) -> + case exitCode of + ExitSuccess -> do + let versionLine = T.decodeUtf8Lenient $ LBS.toStrict stdout + pure $ Right $ extractMajorVersion versionLine + ExitFailure code -> + pure $ Left $ "postgres --version failed with code " <> T.pack (show code) + where + config = + proc "postgres" ["--version"] + & setStdout byteStringOutput + & setStderr byteStringOutput + + (&) = flip ($) + + -- Extract major version from "postgres (PostgreSQL) 17.7" + extractMajorVersion :: Text -> Text + extractMajorVersion line = + let parts = T.words line + -- Find the version number (last word that contains a digit) + versionPart = case filter (T.any (`elem` ['0' .. '9'])) parts of + [] -> "unknown" + vs -> last vs + -- Take just the major version (before first dot) + majorVersion = T.takeWhile (/= '.') versionPart + in majorVersion + +-- | Generate a cache key for the given configuration. +getCacheKey :: Config -> IO (Either Text CacheKey) +getCacheKey config = do + versionResult <- getPostgresVersion + case versionResult of + Left err -> pure $ Left err + Right version -> do + -- Hash the configuration elements that affect initdb output + let configStr = + T.unlines + [ T.unwords $ configInitDbArgs config, + T.unlines $ map (\(k, v) -> k <> "=" <> v) $ configPostgresSettings config, + configUser config + ] + let configHash = T.pack $ show $ abs $ hash (T.unpack configStr) + pure $ + Right $ + CacheKey + { cacheKeyPgVersion = version, + cacheKeyConfigHash = configHash + } + +-- | Get the directory path for a given cache key. +getCacheDirectory :: CacheKey -> Maybe FilePath -> IO FilePath +getCacheDirectory key mRoot = do + cacheDir <- ensureCacheDirectory mRoot + let keyDir = T.unpack (cacheKeyPgVersion key) <> "-" <> T.unpack (cacheKeyConfigHash key) + pure $ cacheDir keyDir + +-- | Check if a cache exists for the given key. +isCached :: CacheKey -> Maybe FilePath -> IO Bool +isCached key mRoot = do + dir <- getCacheDirectory key mRoot + let dataDir = dir "data" + doesDirectoryExist dataDir + +-- | Create a cache from an initialized data directory. +createCache :: CacheKey -> FilePath -> Maybe FilePath -> IO (Either Text ()) +createCache key srcDataDir mRoot = do + dir <- getCacheDirectory key mRoot + createDirectoryIfMissing True dir + let dstDataDir = dir "data" + + -- Detect CoW capability + cowCapability <- detectCowCapability dir + + -- Copy the data directory to the cache + copyDirectory cowCapability srcDataDir dstDataDir + +-- | Restore from cache to a new data directory. +restoreFromCache :: CacheKey -> FilePath -> Maybe FilePath -> IO (Either Text ()) +restoreFromCache key dstDataDir mRoot = do + dir <- getCacheDirectory key mRoot + let srcDataDir = dir "data" + + exists <- doesDirectoryExist srcDataDir + if not exists + then pure $ Left $ "Cache not found: " <> T.pack srcDataDir + else do + -- Detect CoW capability + cowCapability <- detectCowCapability dir + + -- Copy from cache to destination + copyDirectory cowCapability srcDataDir dstDataDir + +-- | Clear the cache for a specific key. +clearCache :: CacheKey -> Maybe FilePath -> IO (Either Text ()) +clearCache key mRoot = do + dir <- getCacheDirectory key mRoot + exists <- doesDirectoryExist dir + if not exists + then pure $ Right () + else do + result <- try $ removeDirectoryRecursive dir + case result of + Left (ex :: SomeException) -> + pure $ Left $ "Failed to clear cache: " <> T.pack (show ex) + Right () -> + pure $ Right () + +-- | Clear all caches. +clearAllCaches :: Maybe FilePath -> IO (Either Text ()) +clearAllCaches mRoot = do + cacheDir <- ensureCacheDirectory mRoot + result <- try $ do + entries <- listDirectory cacheDir + mapM_ (\e -> removeDirectoryRecursive (cacheDir e)) entries + case result of + Left (ex :: SomeException) -> + pure $ Left $ "Failed to clear all caches: " <> T.pack (show ex) + Right () -> + pure $ Right () + +-- | Clean up PostgreSQL runtime files from a data directory. +-- +-- This removes files that are created when postgres is running +-- and should not be present in a fresh data directory (or cache). +-- Files removed: +-- - postmaster.pid (postgres process ID file) +-- - postmaster.opts (command line options) +cleanupRuntimeFiles :: FilePath -> IO () +cleanupRuntimeFiles dataDir = do + let runtimeFiles = + [ dataDir "postmaster.pid", + dataDir "postmaster.opts" + ] + mapM_ removeIfExists runtimeFiles + where + removeIfExists :: FilePath -> IO () + removeIfExists path = do + exists <- doesFileExist path + if exists + then removeFile path `catch_` pure () + else pure () + + catch_ :: IO a -> IO a -> IO a + catch_ action fallback = do + result <- try @SomeException action + case result of + Left _ -> fallback + Right a -> pure a diff --git a/src/EphemeralPg/Internal/CopyOnWrite.hs b/src/EphemeralPg/Internal/CopyOnWrite.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Internal/CopyOnWrite.hs @@ -0,0 +1,136 @@ +-- | Copy-on-write detection and operations. +-- +-- This module provides copy-on-write support for efficient directory copying. +-- On macOS, this uses @clonefile()@ via @cp -c@. On Linux with Btrfs/XFS, +-- this uses @cp --reflink=always@. +-- +-- Detection is done by actually attempting a clone operation on a test file, +-- so it works correctly regardless of filesystem type. +module EphemeralPg.Internal.CopyOnWrite + ( -- * Capability detection + CowCapability (..), + CowMethod (..), + detectCowCapability, + + -- * Directory operations + copyDirectory, + copyDirectoryCoW, + copyDirectoryRegular, + ) +where + +import Control.Exception (SomeException, try) +import Data.Text (Text) +import Data.Text qualified as T +import System.Directory (removeFile) +import System.Exit (ExitCode (..)) +import System.FilePath (()) +import System.IO.Temp (withSystemTempDirectory) +import System.Info (os) +import System.Process.Typed (nullStream, proc, runProcess, runProcess_, setStderr) + +-- | Copy-on-write capability for a filesystem. +data CowCapability + = -- | CoW is supported with the given method + CowSupported CowMethod + | -- | CoW is not supported, with reason + CowNotSupported Text + deriving stock (Eq, Show) + +-- | The method used for copy-on-write. +data CowMethod + = -- | macOS clonefile via cp -c + CowClonefile + | -- | Linux reflink via cp --reflink + CowReflink + deriving stock (Eq, Show) + +-- | Detect copy-on-write capability for a path. +-- +-- This tests the actual filesystem by attempting a clone operation. +-- The test is performed once and the result can be cached. +-- +-- On macOS: Tests @cp -c@ (uses clonefile syscall) +-- On Linux: Tests @cp --reflink=always@ +detectCowCapability :: FilePath -> IO CowCapability +detectCowCapability _testDir = do + withSystemTempDirectory "cow-test-" $ \tmpDir -> do + let srcFile = tmpDir "src" + let dstFile = tmpDir "dst" + + -- Create a small test file + writeFile srcFile "test" + + -- Try the platform-specific clone command + result <- tryClone srcFile dstFile + + -- Clean up + removeFile srcFile `catch_` pure () + removeFile dstFile `catch_` pure () + + pure result + where + catch_ :: IO a -> IO a -> IO a + catch_ action fallback = do + result <- try @SomeException action + case result of + Left _ -> fallback + Right a -> pure a + + tryClone :: FilePath -> FilePath -> IO CowCapability + tryClone src dst = case os of + "darwin" -> tryMacOSClone src dst + "linux" -> tryLinuxReflink src dst + _ -> pure $ CowNotSupported $ "Unsupported OS: " <> T.pack os + + tryMacOSClone :: FilePath -> FilePath -> IO CowCapability + tryMacOSClone src dst = do + -- cp -c uses clonefile() on macOS + exitCode <- runProcess $ setStderr nullStream $ proc "cp" ["-c", src, dst] + case exitCode of + ExitSuccess -> pure $ CowSupported CowClonefile + ExitFailure _ -> pure $ CowNotSupported "clonefile not supported on this filesystem" + + tryLinuxReflink :: FilePath -> FilePath -> IO CowCapability + tryLinuxReflink src dst = do + -- cp --reflink=always fails if reflink is not supported + exitCode <- runProcess $ setStderr nullStream $ proc "cp" ["--reflink=always", src, dst] + case exitCode of + ExitSuccess -> pure $ CowSupported CowReflink + ExitFailure _ -> pure $ CowNotSupported "reflink not supported on this filesystem" + +-- | Copy a directory, using copy-on-write if the capability is available. +copyDirectory :: CowCapability -> FilePath -> FilePath -> IO (Either Text ()) +copyDirectory capability src dst = case capability of + CowSupported method -> copyDirectoryCoW method src dst + CowNotSupported _ -> copyDirectoryRegular src dst + +-- | Copy a directory using copy-on-write. +copyDirectoryCoW :: CowMethod -> FilePath -> FilePath -> IO (Either Text ()) +copyDirectoryCoW method src dst = do + result <- try $ runCopy method + case result of + Left (_ :: SomeException) -> + -- Fall back to regular copy on failure + copyDirectoryRegular src dst + Right () -> + pure $ Right () + where + runCopy :: CowMethod -> IO () + runCopy = \case + CowClonefile -> + -- macOS: use cp -cR for recursive copy-on-write + runProcess_ $ proc "cp" ["-cR", src, dst] + CowReflink -> + -- Linux: use cp --reflink=auto -R (auto falls back gracefully) + runProcess_ $ proc "cp" ["--reflink=auto", "-R", src, dst] + +-- | Copy a directory using regular (non-CoW) copy. +copyDirectoryRegular :: FilePath -> FilePath -> IO (Either Text ()) +copyDirectoryRegular src dst = do + result <- try $ runProcess_ $ proc "cp" ["-R", src, dst] + case result of + Left (ex :: SomeException) -> + pure $ Left $ T.pack $ show ex + Right () -> + pure $ Right () diff --git a/src/EphemeralPg/Internal/Directory.hs b/src/EphemeralPg/Internal/Directory.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Internal/Directory.hs @@ -0,0 +1,147 @@ +-- | Directory management for temporary PostgreSQL instances. +module EphemeralPg.Internal.Directory + ( -- * Directory creation + createTempDataDirectory, + createTempSocketDirectory, + resolveDirectory, + + -- * Socket path validation + validateSocketPath, + estimateSocketPathLength, + + -- * Cleanup + removeDirectoryIfExists, + retryRemoveDirectory, + ) +where + +import Control.Concurrent (threadDelay) +import Control.Exception (SomeException, catch, try) +import Control.Monad (when) +import Data.Text (Text) +import Data.Text qualified as T +import EphemeralPg.Config (DirectoryConfig (..), maxSocketPathLength) +import EphemeralPg.Error (ConfigError (..), ResourceError (..), StartError (..)) +import System.Directory + ( createDirectoryIfMissing, + doesDirectoryExist, + getTemporaryDirectory, + removeDirectoryRecursive, + ) +import System.IO.Temp (createTempDirectory) + +-- | Create a temporary data directory. +-- +-- Returns the path and whether it should be cleaned up. +createTempDataDirectory :: Maybe FilePath -> IO (Either StartError (FilePath, Bool)) +createTempDataDirectory mTempRoot = do + tmpRoot <- maybe getTemporaryDirectory pure mTempRoot + result <- try $ createTempDirectory tmpRoot "ephpg-data-" + case result of + Left (e :: SomeException) -> + pure $ Left $ ResourceError $ DirectoryCreationFailed tmpRoot (T.pack $ show e) + Right path -> + pure $ Right (path, True) + +-- | Create a temporary socket directory. +-- +-- Uses short names to avoid Unix socket path length limits. +-- Returns the path and whether it should be cleaned up. +createTempSocketDirectory :: Maybe FilePath -> IO (Either StartError (FilePath, Bool)) +createTempSocketDirectory mTempRoot = do + tmpRoot <- maybe getTemporaryDirectory pure mTempRoot + -- Use very short prefix to minimize socket path length + result <- try $ createTempDirectory tmpRoot "pg-" + case result of + Left (e :: SomeException) -> + pure $ Left $ ResourceError $ DirectoryCreationFailed tmpRoot (T.pack $ show e) + Right path -> do + -- Validate the socket path length + case validateSocketPath path of + Left err -> do + -- Clean up the directory we just created + removeDirectoryIfExists path + pure $ Left err + Right () -> + pure $ Right (path, True) + +-- | Resolve a directory configuration to an actual path. +-- +-- Returns the path and whether it should be cleaned up. +resolveDirectory :: + DirectoryConfig -> + -- | Temp root + Maybe FilePath -> + -- | Directory purpose (for error messages) + Text -> + -- | Creator + (Maybe FilePath -> IO (Either StartError (FilePath, Bool))) -> + IO (Either StartError (FilePath, Bool)) +resolveDirectory config mTempRoot _purpose creator = + case config of + DirectoryTemporary -> + creator mTempRoot + DirectoryPermanent path -> do + result <- try $ createDirectoryIfMissing True path + case result of + Left (e :: SomeException) -> + pure $ Left $ ResourceError $ DirectoryCreationFailed path (T.pack $ show e) + Right () -> + pure $ Right (path, False) + +-- | Validate that a socket path won't exceed Unix limits. +-- +-- The socket file will be named @.s.PGSQL.@ which adds about 15 characters. +validateSocketPath :: FilePath -> Either StartError () +validateSocketPath socketDir = do + let estimatedLen = estimateSocketPathLength socketDir + when (estimatedLen > maxSocketPathLength) $ + Left $ + ConfigError $ + SocketPathTooLong + { socketPath = socketDir, + socketPathLength = estimatedLen, + socketPathMaxLength = maxSocketPathLength + } + +-- | Estimate the full socket path length. +-- +-- PostgreSQL socket files are named @.s.PGSQL.@. +-- Assuming max port 65535, the suffix is at most 17 characters. +estimateSocketPathLength :: FilePath -> Int +estimateSocketPathLength socketDir = + length socketDir + 1 + 17 -- +1 for path separator, +17 for ".s.PGSQL.65535" + +-- | Remove a directory if it exists, ignoring errors. +removeDirectoryIfExists :: FilePath -> IO () +removeDirectoryIfExists path = do + exists <- doesDirectoryExist path + when exists $ + removeDirectoryRecursive path `catch` \(_ :: SomeException) -> pure () + +-- | Remove a directory with retries. +-- +-- This handles the race condition where PostgreSQL may still be writing +-- files (like pg_stat) during shutdown. +retryRemoveDirectory :: FilePath -> Int -> Int -> IO (Either Text ()) +retryRemoveDirectory path maxRetries delayMicros = go maxRetries + where + go 0 = do + result <- try $ removeDirectoryRecursive path + case result of + Left (e :: SomeException) -> + pure $ Left $ "Failed after " <> T.pack (show maxRetries) <> " retries: " <> T.pack (show e) + Right () -> + pure $ Right () + go n = do + result <- try $ removeDirectoryRecursive path + case result of + Left (_ :: SomeException) -> do + -- Wait and retry + threadDelayMicros delayMicros + go (n - 1) + Right () -> + pure $ Right () + + threadDelayMicros :: Int -> IO () + threadDelayMicros = threadDelay diff --git a/src/EphemeralPg/Internal/Except.hs b/src/EphemeralPg/Internal/Except.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Internal/Except.hs @@ -0,0 +1,67 @@ +-- | ExceptT-based helpers for startup operations. +-- +-- This module provides a cleaner way to sequence operations that can fail, +-- replacing deeply nested case statements with 'ExceptT' do-notation. +module EphemeralPg.Internal.Except + ( -- * Startup monad + Startup, + runStartup, + + -- * Lifting operations + liftE, + liftMaybe, + + -- * Error handling + onError, + ) +where + +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Trans.Except (ExceptT (..), catchE, runExceptT, throwE) +import EphemeralPg.Error (StartError) + +-- | The startup monad: 'ExceptT' 'StartError' 'IO'. +-- +-- This provides: +-- +-- * Automatic short-circuiting on errors (like @Either@, but in @IO@) +-- * Clean @do@ notation for sequencing fallible operations +-- * @liftIO@ for embedding pure @IO@ actions +type Startup = ExceptT StartError IO + +-- | Run a startup computation, returning @Either StartError a@. +runStartup :: Startup a -> IO (Either StartError a) +runStartup = runExceptT + +-- | Lift an @IO (Either StartError a)@ into 'Startup'. +-- +-- This is the most common lifting operation, used for functions that +-- already return @Either StartError@. +-- +-- @ +-- (dataDir, isTemp) <- liftE $ createTempDataDirectory mRoot +-- @ +liftE :: IO (Either StartError a) -> Startup a +liftE = ExceptT + +-- | Lift a 'Maybe' into 'Startup', using the given error if 'Nothing'. +-- +-- @ +-- postgresPath <- liftMaybe (PostgresStartError PostgresNotFound) +-- =<< liftIO (findExecutable "postgres") +-- @ +liftMaybe :: StartError -> Maybe a -> Startup a +liftMaybe err = maybe (throwE err) pure + +-- | Run a cleanup action if the computation fails, then re-throw the error. +-- +-- This is used to ensure resources are cleaned up on failure: +-- +-- @ +-- (socketDir, socketDirIsTemp) <- liftE (createTempSocketDirectory mRoot) +-- \`onError\` when dataDirIsTemp (removeDirectoryIfExists dataDir) +-- @ +onError :: Startup a -> IO () -> Startup a +onError action cleanup = catchE action $ \err -> do + liftIO cleanup + throwE err diff --git a/src/EphemeralPg/Internal/Port.hs b/src/EphemeralPg/Internal/Port.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Internal/Port.hs @@ -0,0 +1,59 @@ +-- | Port allocation for temporary PostgreSQL instances. +module EphemeralPg.Internal.Port + ( -- * Port allocation + findFreePort, + isPortAvailable, + ) +where + +import Control.Exception (SomeException, bracket, try) +import Data.Text qualified as T +import Data.Word (Word16) +import EphemeralPg.Error (ResourceError (..), StartError (..)) +import Network.Socket + ( Family (AF_INET), + SockAddr (SockAddrInet), + Socket, + SocketType (Stream), + bind, + close, + socket, + socketPort, + tupleToHostAddress, + ) + +-- | Find a free port by binding to port 0 and checking what the OS assigns. +findFreePort :: IO (Either StartError Word16) +findFreePort = do + result <- try $ bracket openSocket close getAssignedPort + case result of + Left (e :: SomeException) -> + pure $ Left $ ResourceError $ PortAllocationFailed $ T.pack $ show e + Right port -> + pure $ Right port + where + openSocket :: IO Socket + openSocket = do + sock <- socket AF_INET Stream 0 + -- Bind to localhost on port 0 (OS assigns a free port) + bind sock (SockAddrInet 0 (tupleToHostAddress (127, 0, 0, 1))) + pure sock + + getAssignedPort :: Socket -> IO Word16 + getAssignedPort sock = do + port <- socketPort sock + pure $ fromIntegral port + +-- | Check if a specific port is available. +isPortAvailable :: Word16 -> IO Bool +isPortAvailable port = do + result <- try $ bracket openSocket close (const $ pure ()) + case result of + Left (_ :: SomeException) -> pure False + Right () -> pure True + where + openSocket :: IO Socket + openSocket = do + sock <- socket AF_INET Stream 0 + bind sock (SockAddrInet (fromIntegral port) (tupleToHostAddress (127, 0, 0, 1))) + pure sock diff --git a/src/EphemeralPg/Process.hs b/src/EphemeralPg/Process.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Process.hs @@ -0,0 +1,59 @@ +-- | Common process utilities for PostgreSQL command execution. +module EphemeralPg.Process + ( -- * Process execution + runProcessCapture, + + -- * Utilities + findExecutable, + getCurrentUser, + ) +where + +import Control.Exception (SomeException, try) +import Data.ByteString.Lazy qualified as LBS +import Data.Function ((&)) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding qualified as T +import System.Directory qualified as Dir +import System.Exit (ExitCode (..)) +import System.Posix.User (getEffectiveUserName) +import System.Process.Typed + ( byteStringOutput, + proc, + readProcess, + setStderr, + setStdout, + ) + +-- | Run a process and capture its output. +runProcessCapture :: + -- | Executable + FilePath -> + -- | Arguments + [String] -> + -- | (exit code, stdout, stderr) + IO (ExitCode, Text, Text) +runProcessCapture exe args = do + let config = + proc exe args + & setStdout byteStringOutput + & setStderr byteStringOutput + (exitCode, stdout, stderr) <- readProcess config + pure + ( exitCode, + T.decodeUtf8Lenient $ LBS.toStrict stdout, + T.decodeUtf8Lenient $ LBS.toStrict stderr + ) + +-- | Find an executable in PATH. +findExecutable :: String -> IO (Maybe FilePath) +findExecutable = Dir.findExecutable + +-- | Get the current effective username. +getCurrentUser :: IO Text +getCurrentUser = do + result <- try getEffectiveUserName + case result of + Left (_ :: SomeException) -> pure "postgres" + Right name -> pure $ T.pack name diff --git a/src/EphemeralPg/Process/CreateDb.hs b/src/EphemeralPg/Process/CreateDb.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Process/CreateDb.hs @@ -0,0 +1,69 @@ +-- | createdb process execution. +module EphemeralPg.Process.CreateDb + ( runCreateDb, + ) +where + +import Data.Text (Text) +import Data.Text qualified as T +import Data.Word (Word16) +import EphemeralPg.Config (Config (..)) +import EphemeralPg.Error (CreateDbError (..), StartError (..)) +import EphemeralPg.Process (findExecutable, runProcessCapture) +import System.Exit (ExitCode (..)) + +-- | Run createdb to create a database. +runCreateDb :: + Config -> + -- | Socket directory + FilePath -> + -- | Port + Word16 -> + -- | Username + Text -> + -- | Database name + Text -> + IO (Either StartError ()) +runCreateDb config socketDir port username dbName = do + -- The "postgres" database is created by initdb, so we don't need to create it + if dbName == "postgres" + then pure $ Right () + else do + -- Find createdb executable + mCreateDb <- findExecutable "createdb" + case mCreateDb of + Nothing -> pure $ Left $ CreateDbError CreateDbNotFound + Just createDbPath -> do + -- Build arguments + let args = buildCreateDbArgs config socketDir port username dbName + + -- Run createdb + (exitCode, stdout, stderr) <- runProcessCapture createDbPath (map T.unpack args) + + case exitCode of + ExitSuccess -> pure $ Right () + ExitFailure _ -> + -- Check if it's because the database already exists + if "already exists" `T.isInfixOf` stderr + then pure $ Left $ CreateDbError $ DatabaseAlreadyExists dbName + else + pure $ + Left $ + CreateDbError $ + CreateDbFailed + { createDbExitCode = exitCode, + createDbStdout = stdout, + createDbStderr = stderr, + createDbCommand = T.unwords (T.pack createDbPath : args), + createDbName = dbName + } + +-- | Build createdb command line arguments. +buildCreateDbArgs :: Config -> FilePath -> Word16 -> Text -> Text -> [Text] +buildCreateDbArgs config socketDir port username dbName = + [ "--host=" <> T.pack socketDir, + "--port=" <> T.pack (show port), + "--username=" <> username, + dbName + ] + <> configCreateDbArgs config diff --git a/src/EphemeralPg/Process/InitDb.hs b/src/EphemeralPg/Process/InitDb.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Process/InitDb.hs @@ -0,0 +1,72 @@ +-- | initdb process execution. +module EphemeralPg.Process.InitDb + ( -- * Initialization + runInitDb, + + -- * Configuration file generation + writePostgresConf, + ) +where + +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.IO qualified as T +import EphemeralPg.Config (Config (..)) +import EphemeralPg.Error (InitDbError (..), StartError (..)) +import EphemeralPg.Process (findExecutable, getCurrentUser, runProcessCapture) +import System.Exit (ExitCode (..)) +import System.FilePath (()) + +-- | Run initdb to initialize a PostgreSQL data directory. +runInitDb :: Config -> FilePath -> IO (Either StartError ()) +runInitDb config dataDir = do + -- Find initdb executable + mInitDb <- findExecutable "initdb" + case mInitDb of + Nothing -> pure $ Left $ InitDbError InitDbNotFound + Just initDbPath -> do + -- Get username + username <- case configUser config of + "" -> getCurrentUser + u -> pure u + + -- Build arguments + let args = buildInitDbArgs config dataDir username + + -- Run initdb + (exitCode, stdout, stderr) <- runProcessCapture initDbPath (map T.unpack args) + + case exitCode of + ExitSuccess -> do + -- Write postgresql.conf with our settings + writePostgresConf config dataDir + pure $ Right () + ExitFailure _ -> + pure $ + Left $ + InitDbError $ + InitDbFailed + { initDbExitCode = exitCode, + initDbStdout = stdout, + initDbStderr = stderr, + initDbCommand = T.unwords (T.pack initDbPath : args) + } + +-- | Build initdb command line arguments. +buildInitDbArgs :: Config -> FilePath -> Text -> [Text] +buildInitDbArgs config dataDir username = + [ "--pgdata=" <> T.pack dataDir, + "--username=" <> username + ] + <> configInitDbArgs config + +-- | Write postgresql.conf with the configured settings. +writePostgresConf :: Config -> FilePath -> IO () +writePostgresConf config dataDir = do + let confPath = dataDir "postgresql.conf" + let settings = configPostgresSettings config + let content = T.unlines $ map formatSetting settings + T.appendFile confPath ("\n# ephemeral-pg settings\n" <> content) + where + formatSetting :: (Text, Text) -> Text + formatSetting (key, value) = key <> " = " <> value diff --git a/src/EphemeralPg/Process/Postgres.hs b/src/EphemeralPg/Process/Postgres.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Process/Postgres.hs @@ -0,0 +1,210 @@ +-- | PostgreSQL server process management. +module EphemeralPg.Process.Postgres + ( -- * Server lifecycle + startPostgres, + stopPostgres, + waitForPostgres, + + -- * Types + PostgresProcess (..), + ) +where + +import Control.Concurrent (threadDelay) +import Control.Exception (SomeException, mask_, try) +import Control.Monad.IO.Class (liftIO) +import Control.Monad.Trans.Except (throwE) +import Data.Function ((&)) +import Data.Monoid (Last (..)) +import Data.Text (Text) +import Data.Text qualified as T +import Data.Word (Word16) +import EphemeralPg.Config + ( Config (..), + ShutdownMode (..), + defaultConnectionTimeoutSeconds, + ) +import EphemeralPg.Database (PostgresProcess (..)) +import EphemeralPg.Error + ( PostgresError (..), + StartError (..), + StopError (..), + TimeoutError (..), + ) +import EphemeralPg.Internal.Except (liftE, liftMaybe, onError, runStartup) +import EphemeralPg.Process (findExecutable) +import System.Exit (ExitCode (..)) +import System.Posix.Signals (sigINT, sigKILL, sigQUIT, sigTERM, signalProcess) +import System.Posix.Types (CPid (..)) +import System.Process (getPid) +import System.Process.Typed + ( nullStream, + proc, + runProcess, + setCreateGroup, + setStderr, + setStdin, + setStdout, + startProcess, + unsafeProcessHandle, + waitExitCode, + ) +import System.Timeout (timeout) + +-- | Start the PostgreSQL server. +startPostgres :: + Config -> + -- | Data directory + FilePath -> + -- | Socket directory + FilePath -> + -- | Port + Word16 -> + -- | Username + Text -> + IO (Either StartError PostgresProcess) +startPostgres config dataDir socketDir port username = runStartup $ do + -- Find postgres executable + postgresPath <- + liftMaybe (PostgresStartError PostgresNotFound) + =<< liftIO (findExecutable "postgres") + + -- Build arguments and process config + let args = buildPostgresArgs config dataDir socketDir port username + let processConfig = + proc postgresPath (map T.unpack args) + & setStdin nullStream + & setStdout nullStream + & setStderr nullStream + & setCreateGroup True -- So we can signal the whole group + + -- Start the process + process <- + liftIO (try $ startProcess processConfig) >>= \case + Left (ex :: SomeException) -> + throwE $ + PostgresStartError $ + PostgresStartFailed + { pgExitCode = ExitFailure 1, + pgStdout = "", + pgStderr = T.pack $ show ex, + pgCommand = T.unwords (T.pack postgresPath : args) + } + Right p -> pure p + + -- Get the PID + let pHandle = unsafeProcessHandle process + pid <- + liftMaybe + ( PostgresStartError $ + PostgresStartFailed + { pgExitCode = ExitFailure 1, + pgStdout = "", + pgStderr = "Could not get process ID", + pgCommand = T.unwords (T.pack postgresPath : args) + } + ) + =<< liftIO (getPid pHandle) + + let postgresProcess = + PostgresProcess + { postgresProcess = process, + postgresPid = CPid $ fromIntegral pid + } + + -- Wait for the server to be ready + let timeoutSecs = + maybe defaultConnectionTimeoutSeconds id $ + getLast (configConnectionTimeoutSeconds config) + + liftE (waitForPostgres socketDir port timeoutSecs) + `onError` do + -- Kill the server since it didn't start properly + _ <- stopPostgres postgresProcess ShutdownImmediate 5 + pure () + + pure postgresProcess + +-- | Build postgres command line arguments. +buildPostgresArgs :: Config -> FilePath -> FilePath -> Word16 -> Text -> [Text] +buildPostgresArgs config dataDir socketDir port _username = + [ "-D", + T.pack dataDir, + "-k", + T.pack socketDir, + "-p", + T.pack (show port), + "-h", + "127.0.0.1" -- Also listen on TCP for debugging + ] + <> configPostgresArgs config + +-- | Wait for PostgreSQL to accept connections. +waitForPostgres :: FilePath -> Word16 -> Int -> IO (Either StartError ()) +waitForPostgres socketDir port timeoutSecs = do + let deadline = timeoutSecs * 1000000 -- Convert to microseconds + result <- timeout deadline waitLoop + case result of + Nothing -> + pure $ + Left $ + TimeoutError $ + ConnectionTimeout + { timeoutDurationSeconds = timeoutSecs, + timeoutHost = T.pack socketDir, + timeoutPort = port + } + Just () -> pure $ Right () + where + waitLoop :: IO () + waitLoop = do + -- Try to connect using pg_isready + mPgIsReady <- findExecutable "pg_isready" + case mPgIsReady of + Nothing -> do + -- Fall back to just waiting and hoping + threadDelay 1000000 -- 1 second + pure () + Just pgIsReadyPath -> do + let args = + [ "-h", + socketDir, + "-p", + show port, + "-t", + "1" -- 1 second timeout per attempt + ] + result <- try @SomeException $ runProcess $ proc pgIsReadyPath args + case result of + Right ExitSuccess -> pure () + _ -> do + threadDelay 100000 -- 100ms between attempts + waitLoop + +-- | Stop the PostgreSQL server. +stopPostgres :: PostgresProcess -> ShutdownMode -> Int -> IO (Maybe StopError) +stopPostgres PostgresProcess {..} mode timeoutSecs = mask_ $ do + let signal = case mode of + ShutdownGraceful -> sigTERM + ShutdownFast -> sigINT + ShutdownImmediate -> sigQUIT + + -- Send the signal + result <- try $ signalProcess signal postgresPid + case result of + Left (_ :: SomeException) -> + -- Process might already be dead + pure Nothing + Right () -> do + -- Wait for the process to exit with timeout + let deadline = timeoutSecs * 1000000 + exitResult <- timeout deadline $ waitExitCode postgresProcess + + case exitResult of + Just _ -> pure Nothing -- Exited normally + Nothing -> do + -- Timeout: force kill + _ <- try @SomeException $ signalProcess sigKILL postgresPid + -- Wait a bit more for the forced kill + _ <- timeout 5000000 $ waitExitCode postgresProcess + pure $ Just $ ShutdownTimedOut timeoutSecs diff --git a/src/EphemeralPg/Snapshot.hs b/src/EphemeralPg/Snapshot.hs new file mode 100644 --- /dev/null +++ b/src/EphemeralPg/Snapshot.hs @@ -0,0 +1,173 @@ +-- | Database snapshot functionality for test isolation. +-- +-- This module provides snapshot and restore capabilities for ephemeral +-- PostgreSQL databases. Snapshots capture the entire database state +-- and can be restored to reset the database between tests. +-- +-- = Example Usage +-- +-- @ +-- result <- Pg.'with' $ \\db -> do +-- -- Initial setup +-- runMigrations db +-- +-- -- Take a snapshot after setup +-- snapshot <- Pg.'createSnapshot' db +-- +-- -- Run test 1 +-- insertTestData db +-- runTest1 db +-- +-- -- Restore to clean state +-- Pg.'restoreSnapshot' snapshot db +-- +-- -- Run test 2 with clean state +-- insertTestData db +-- runTest2 db +-- @ +module EphemeralPg.Snapshot + ( -- * Snapshot Types + Snapshot (..), + + -- * Snapshot Operations + createSnapshot, + restoreSnapshot, + deleteSnapshot, + ) +where + +import Data.Text (Text) +import Data.Text qualified as T +import EphemeralPg.Config (ShutdownMode (..), defaultConfig) +import EphemeralPg.Database (Database (..), PostgresProcess (..)) +import EphemeralPg.Internal.Cache (cleanupRuntimeFiles) +import EphemeralPg.Internal.CopyOnWrite + ( copyDirectory, + detectCowCapability, + ) +import EphemeralPg.Internal.Directory (removeDirectoryIfExists) +import EphemeralPg.Process.Postgres (startPostgres, stopPostgres) +import System.Directory (doesDirectoryExist) +import System.IO.Temp (createTempDirectory, getCanonicalTemporaryDirectory) + +-- | A snapshot of database state. +data Snapshot = Snapshot + { -- | Path to the snapshot data directory + snapshotPath :: FilePath, + -- | Whether to delete the snapshot when done + snapshotTemporary :: Bool + } + deriving stock (Eq, Show) + +-- | Create a snapshot of the current database state. +-- +-- This stops postgres, copies the data directory using copy-on-write +-- if available, then restarts postgres. +-- +-- Note: This is a relatively expensive operation and should be used +-- sparingly (e.g., once after setup, not between every test). +createSnapshot :: Database -> IO (Either Text Snapshot) +createSnapshot db = do + -- Stop postgres gracefully to ensure data is flushed + stopResult <- stopPostgres (dbProcess db) ShutdownGraceful 30 + case stopResult of + Just err -> do + -- Restart postgres and return error + _ <- restartPostgres db + pure $ Left $ "Failed to stop postgres for snapshot: " <> T.pack (show err) + Nothing -> do + -- Create snapshot directory + tmpDir <- getCanonicalTemporaryDirectory + snapshotDir <- createTempDirectory tmpDir "ephpg-snap-" + + -- Detect CoW capability + cowCapability <- detectCowCapability snapshotDir + + -- Copy data directory to snapshot + copyResult <- copyDirectory cowCapability (dbDataDirectory db) snapshotDir + + case copyResult of + Left err -> do + removeDirectoryIfExists snapshotDir + _ <- restartPostgres db + pure $ Left $ "Failed to copy data directory: " <> err + Right () -> do + -- Clean up runtime files from the snapshot + cleanupRuntimeFiles snapshotDir + + -- Restart postgres + restartResult <- restartPostgres db + case restartResult of + Left err -> do + removeDirectoryIfExists snapshotDir + pure $ Left err + Right _newProcess -> + -- Note: We don't update the Database record with the new process + -- since Database is immutable. The caller should handle this. + pure $ + Right $ + Snapshot + { snapshotPath = snapshotDir, + snapshotTemporary = True + } + +-- | Restore a database from a snapshot. +-- +-- This stops postgres, replaces the data directory with the snapshot, +-- then restarts postgres. The original database content is lost. +restoreSnapshot :: Snapshot -> Database -> IO (Either Text ()) +restoreSnapshot snapshot db = do + -- Verify snapshot exists + exists <- doesDirectoryExist (snapshotPath snapshot) + if not exists + then pure $ Left $ "Snapshot not found: " <> T.pack (snapshotPath snapshot) + else do + -- Stop postgres + stopResult <- stopPostgres (dbProcess db) ShutdownGraceful 30 + case stopResult of + Just err -> do + _ <- restartPostgres db + pure $ Left $ "Failed to stop postgres for restore: " <> T.pack (show err) + Nothing -> do + -- Remove current data directory + removeDirectoryIfExists (dbDataDirectory db) + + -- Detect CoW capability + cowCapability <- detectCowCapability (snapshotPath snapshot) + + -- Copy snapshot to data directory + copyResult <- copyDirectory cowCapability (snapshotPath snapshot) (dbDataDirectory db) + + case copyResult of + Left err -> do + -- This is bad - data directory is in inconsistent state + pure $ Left $ "Failed to restore snapshot, database may be corrupted: " <> err + Right () -> do + -- Restart postgres + restartResult <- restartPostgres db + case restartResult of + Left err -> pure $ Left err + Right _newProcess -> pure $ Right () + +-- | Delete a snapshot. +-- +-- Only deletes temporary snapshots. Permanent snapshots are left alone. +deleteSnapshot :: Snapshot -> IO () +deleteSnapshot snapshot = + if snapshotTemporary snapshot + then removeDirectoryIfExists (snapshotPath snapshot) + else pure () + +-- | Restart postgres after a snapshot/restore operation. +restartPostgres :: Database -> IO (Either Text PostgresProcess) +restartPostgres db = do + result <- + startPostgres + defaultConfig + (dbDataDirectory db) + (dbSocketDirectory db) + (dbPort db) + (dbUser db) + case result of + Left err -> pure $ Left $ "Failed to restart postgres: " <> T.pack (show err) + Right newProcess -> pure $ Right newProcess diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,114 @@ +module Main where + +import Data.Text qualified as T +import EphemeralPg qualified as Pg +import EphemeralPg.Config qualified as Config +import Hasql.Connection qualified as Connection +import Test.Hspec +import Test.QuickCheck + +main :: IO () +main = hspec $ do + describe "EphemeralPg" $ do + it "can start and stop a database" $ do + result <- Pg.with $ \db -> do + -- Just verify we can get connection settings + let _settings = Pg.connectionSettings db + pure () + result `shouldSatisfy` isRight + + it "can connect to the database" $ do + result <- Pg.with $ \db -> do + connResult <- Connection.acquire (Pg.connectionSettings db) + case connResult of + Left err -> fail $ "Connection failed: " <> show err + Right conn -> do + Connection.release conn + pure () + result `shouldSatisfy` isRight + + it "can restart a database" $ do + result <- Pg.with $ \db -> do + -- Get initial port + let port1 = Pg.port db + + -- Restart the database + restartResult <- Pg.restart db + case restartResult of + Left err -> fail $ "Restart failed: " <> show err + Right db' -> do + -- Verify we can still connect after restart + connResult <- Connection.acquire (Pg.connectionSettings db') + case connResult of + Left err -> fail $ "Connection after restart failed: " <> show err + Right conn -> do + Connection.release conn + -- Port should be the same + Pg.port db' `shouldBe` port1 + result `shouldSatisfy` isRight + + describe "EphemeralPg caching" $ do + it "can start with caching enabled" $ do + result <- Pg.withCached $ \db -> do + connResult <- Connection.acquire (Pg.connectionSettings db) + case connResult of + Left err -> fail $ "Connection failed: " <> show err + Right conn -> do + Connection.release conn + pure () + result `shouldSatisfy` isRight + + it "second cached startup is faster" $ do + -- First run creates the cache + result1 <- Pg.withCached $ \db -> do + let _settings = Pg.connectionSettings db + pure () + result1 `shouldSatisfy` isRight + + -- Second run should use the cache + result2 <- Pg.withCached $ \db -> do + connResult <- Connection.acquire (Pg.connectionSettings db) + case connResult of + Left err -> fail $ "Connection failed: " <> show err + Right conn -> do + Connection.release conn + pure () + result2 `shouldSatisfy` isRight + + describe "Config" $ do + it "satisfies left identity (mempty <> x = x)" $ + property $ \dbName -> + let x = Config.defaultConfig {Config.configDatabaseName = T.pack dbName} + in Config.configDatabaseName (mempty <> x) == Config.configDatabaseName x + + it "satisfies right identity (x <> mempty = x)" $ + property $ \dbName -> + let x = Config.defaultConfig {Config.configDatabaseName = T.pack dbName} + in Config.configDatabaseName (x <> mempty) == Config.configDatabaseName x + + it "satisfies associativity ((x <> y) <> z = x <> (y <> z))" $ + property $ \(n1, n2, n3) -> + let x = mempty {Config.configDatabaseName = T.pack n1} + y = mempty {Config.configDatabaseName = T.pack n2} + z = mempty {Config.configDatabaseName = T.pack n3} + in Config.configDatabaseName ((x <> y) <> z) + == Config.configDatabaseName (x <> (y <> z)) + + it "combines postgres settings" $ do + let c1 = mempty {Config.configPostgresSettings = [("a", "1")]} + c2 = mempty {Config.configPostgresSettings = [("b", "2")]} + combined = c1 <> c2 + Config.configPostgresSettings combined `shouldBe` [("a", "1"), ("b", "2")] + + describe "Socket path validation" $ do + it "rejects paths that are too long" $ do + -- Unix socket path limit is typically 104-108 bytes + -- PostgreSQL adds ".s.PGSQL." (~17 chars) to the path + let longPath = replicate 200 'x' + -- This should fail during startup due to socket path length + -- We test this indirectly through the config + length longPath `shouldSatisfy` (> Config.maxSocketPathLength) + +isRight :: Either a b -> Bool +isRight (Right _) = True +isRight (Left _) = False