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