Skip to main content

pg_embedded_setup_unpriv/cluster/
mod.rs

1//! RAII wrapper that boots an embedded `PostgreSQL` instance for tests.
2//!
3//! The cluster starts during [`TestCluster::new`] and shuts down automatically when the
4//! value drops out of scope.
5//!
6//! # Synchronous API
7//!
8//! Use [`TestCluster::new`] from synchronous contexts or when you want the cluster to
9//! own its own Tokio runtime:
10//!
11//! ```no_run
12//! use pg_embedded_setup_unpriv::TestCluster;
13//!
14//! # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
15//! let cluster = TestCluster::new()?;
16//! let url = cluster.settings().url("my_database");
17//! // Perform test database work here.
18//! drop(cluster); // `PostgreSQL` stops automatically.
19//!
20//! # Ok(())
21//! # }
22//! ```
23//!
24//! # Async API
25//!
26//! When running within an existing async runtime (e.g., `#[tokio::test]`), use
27//! [`TestCluster::start_async`] to avoid the "Cannot start a runtime from within a
28//! runtime" panic that occurs when nesting Tokio runtimes:
29//!
30//! ```ignore
31//! use pg_embedded_setup_unpriv::TestCluster;
32//!
33//! #[tokio::test]
34//! async fn test_with_embedded_postgres() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
35//!     let cluster = TestCluster::start_async().await?;
36//!     let url = cluster.settings().url("my_database");
37//!     // ... async database operations ...
38//!     cluster.stop_async().await?;
39//!     Ok(())
40//! }
41//! ```
42//!
43//! The async API requires the `async-api` feature flag:
44//!
45//! ```toml
46//! [dependencies]
47//! pg-embedded-setup-unpriv = { version = "...", features = ["async-api"] }
48//! ```
49
50#[cfg(feature = "async-api")]
51mod async_api;
52mod cache_integration;
53mod cleanup;
54mod connection;
55mod delegation;
56mod guard;
57mod handle;
58mod installation;
59mod lifecycle;
60mod lifecycle_template;
61pub(crate) mod panic_utils;
62mod runtime;
63mod runtime_mode;
64mod shutdown;
65#[cfg(any(unix, windows))]
66mod shutdown_hook;
67#[cfg(all(
68    any(unix, windows),
69    any(doc, test, feature = "cluster-unit-tests", feature = "dev-worker")
70))]
71pub use self::shutdown_hook::{
72    PostmasterPid,
73    PostmasterProcess,
74    postmaster_process_is_running,
75    process_is_running,
76    read_postmaster_pid,
77    read_postmaster_process,
78};
79#[cfg(all(test, feature = "loom-tests"))]
80mod lifecycle_loom_tests;
81mod startup;
82mod temporary_database;
83mod worker_invoker;
84mod worker_operation;
85
86use std::ops::Deref;
87
88use tracing::info_span;
89
90pub(crate) use self::startup::setup_postgres_only;
91#[cfg(any(doc, test, feature = "cluster-unit-tests", feature = "dev-worker"))]
92pub use self::worker_invoker::WorkerInvoker;
93#[doc(hidden)]
94pub use self::worker_operation::WorkerOperation;
95pub use self::{
96    connection::{ConnectionMetadata, TestClusterConnection},
97    guard::ClusterGuard,
98    handle::ClusterHandle,
99    lifecycle::DatabaseName,
100    temporary_database::TemporaryDatabase,
101};
102use self::{
103    runtime::build_runtime,
104    runtime_mode::ClusterRuntime,
105    startup::{cache_config_from_bootstrap, start_postgres},
106};
107use crate::{
108    bootstrap_for_tests,
109    env::ScopedEnv,
110    error::BootstrapResult,
111    observability::LOG_TARGET,
112};
113
114/// Embedded `PostgreSQL` instance whose lifecycle follows Rust's drop semantics.
115///
116/// `TestCluster` combines a [`ClusterHandle`] (for cluster access) with a
117/// [`ClusterGuard`] (for lifecycle management). For most use cases, this
118/// combined type is the simplest option.
119///
120/// # Send-Safe Patterns
121///
122/// `TestCluster` is `!Send` because it contains environment guards that must
123/// be dropped on the creating thread. For patterns requiring `Send` (such as
124/// `OnceLock` or rstest timeouts), use [`new_split()`](Self::new_split) to
125/// obtain a `Send`-safe [`ClusterHandle`]:
126///
127/// ```no_run
128/// use std::sync::OnceLock;
129///
130/// use pg_embedded_setup_unpriv::{ClusterHandle, TestCluster};
131///
132/// static SHARED: OnceLock<ClusterHandle> = OnceLock::new();
133///
134/// fn shared_cluster() -> &'static ClusterHandle {
135///     SHARED.get_or_init(|| {
136///         let (handle, guard) = TestCluster::new_split().expect("cluster bootstrap failed");
137///         handle
138///             .register_shutdown_on_exit()
139///             .expect("shutdown hook registration failed");
140///         std::mem::forget(guard);
141///         handle
142///     })
143/// }
144/// ```
145#[derive(Debug)]
146pub struct TestCluster {
147    /// Send-safe handle providing cluster access.
148    pub(crate) handle: ClusterHandle,
149    /// Lifecycle guard managing shutdown and environment restoration.
150    pub(crate) guard: ClusterGuard,
151}
152
153impl TestCluster {
154    /// Boots a `PostgreSQL` instance configured by [`bootstrap_for_tests`].
155    ///
156    /// The constructor blocks until the underlying server process is running and returns an
157    /// error when startup fails.
158    ///
159    /// # Errors
160    /// Returns an error if the bootstrap configuration cannot be prepared or if starting the
161    /// embedded cluster fails.
162    pub fn new() -> BootstrapResult<Self> {
163        let (handle, guard) = Self::new_split()?;
164        Ok(Self { handle, guard })
165    }
166
167    /// Boots a `PostgreSQL` instance and returns a separate handle and guard.
168    ///
169    /// This constructor is useful for patterns requiring `Send`, such as shared
170    /// cluster fixtures with [`OnceLock`](std::sync::OnceLock) or rstest fixtures
171    /// with timeouts.
172    ///
173    /// # Returns
174    ///
175    /// A tuple of:
176    /// - [`ClusterHandle`]: `Send + Sync` handle for accessing the cluster
177    /// - [`ClusterGuard`]: `!Send` guard managing shutdown and environment
178    ///
179    /// # Errors
180    ///
181    /// Returns an error if the bootstrap configuration cannot be prepared or if
182    /// starting the embedded cluster fails.
183    ///
184    /// # Examples
185    ///
186    /// ## Shared Cluster with `OnceLock`
187    ///
188    /// For shared clusters that run for the entire process lifetime, register
189    /// the shutdown hook and forget the guard to prevent shutdown on drop:
190    ///
191    /// ```no_run
192    /// use std::sync::OnceLock;
193    ///
194    /// use pg_embedded_setup_unpriv::{ClusterHandle, TestCluster};
195    ///
196    /// static SHARED: OnceLock<ClusterHandle> = OnceLock::new();
197    ///
198    /// fn shared_cluster() -> &'static ClusterHandle {
199    ///     SHARED.get_or_init(|| {
200    ///         let (handle, guard) = TestCluster::new_split().expect("cluster bootstrap failed");
201    ///         handle
202    ///             .register_shutdown_on_exit()
203    ///             .expect("shutdown hook registration failed");
204    ///         std::mem::forget(guard);
205    ///         handle
206    ///     })
207    /// }
208    /// ```
209    ///
210    /// **Warning**: Dropping the guard shuts down the cluster. Do not use the
211    /// handle after the guard has been dropped unless the guard was forgotten.
212    pub fn new_split() -> BootstrapResult<(ClusterHandle, ClusterGuard)> {
213        let span = info_span!(target: LOG_TARGET, "test_cluster");
214        // Resolve cache directory BEFORE applying test environment.
215        // Otherwise, the test sandbox's XDG_CACHE_HOME would be used.
216        let (runtime, env_vars, env_guard, outcome) = {
217            let _entered = span.enter();
218            let initial_bootstrap = bootstrap_for_tests()?;
219            let cache_config = cache_config_from_bootstrap(&initial_bootstrap);
220            let runtime = build_runtime()?;
221            let env_vars = initial_bootstrap.environment.to_env();
222            let env_guard = ScopedEnv::apply(&env_vars);
223            let outcome = start_postgres(&runtime, initial_bootstrap, &env_vars, &cache_config)?;
224            (runtime, env_vars, env_guard, outcome)
225        };
226
227        let handle = ClusterHandle::new(outcome.bootstrap.clone());
228        let guard = ClusterGuard {
229            runtime: ClusterRuntime::Sync(runtime),
230            postgres: outcome.postgres,
231            bootstrap: outcome.bootstrap,
232            is_managed_via_worker: outcome.is_managed_via_worker,
233            env_vars,
234            worker_guard: None,
235            _env_guard: env_guard,
236            _cluster_span: span,
237        };
238
239        Ok((handle, guard))
240    }
241
242    /// Extends the cluster lifetime to cover additional scoped environment guards.
243    ///
244    /// Primarily used by fixtures that need to ensure `PG_EMBEDDED_WORKER` remains set for the
245    /// duration of the cluster lifetime.
246    #[doc(hidden)]
247    #[must_use]
248    pub fn with_worker_guard(self, worker_guard: Option<ScopedEnv>) -> Self {
249        Self {
250            handle: self.handle,
251            guard: self.guard.with_worker_guard(worker_guard),
252        }
253    }
254}
255
256/// Provides transparent access to [`ClusterHandle`] methods.
257///
258/// This allows `TestCluster` to be used interchangeably with `ClusterHandle`
259/// for all read-only operations like `settings()`, `connection()`, etc.
260impl Deref for TestCluster {
261    type Target = ClusterHandle;
262
263    fn deref(&self) -> &Self::Target { &self.handle }
264}
265
266// Note: TestCluster does NOT implement Drop because the ClusterGuard handles shutdown.
267// When TestCluster drops, its _guard field drops, which triggers ClusterGuard::Drop.
268
269#[cfg(test)]
270mod mod_tests;
271
272#[cfg(all(test, feature = "cluster-unit-tests"))]
273mod drop_logging_tests;
274
275#[cfg(all(test, not(feature = "cluster-unit-tests")))]
276#[path = "../../tests/test_cluster.rs"]
277mod test_cluster_tests;