Skip to main content

pg_embedded_setup_unpriv/bootstrap/
mod.rs

1//! Bootstraps embedded `PostgreSQL` while adapting to the caller's privileges.
2//!
3//! Provides [`bootstrap_for_tests`] so suites can retrieve structured settings and
4//! prepared environment variables without reimplementing bootstrap orchestration.
5mod env;
6mod env_types;
7mod mode;
8mod prepare;
9
10#[cfg(test)]
11use std::sync::{Arc, Mutex, OnceLock};
12use std::time::Duration;
13
14use color_eyre::eyre::{Context, eyre};
15pub use env::{TestBootstrapEnvironment, find_timezone_dir};
16pub use mode::{ExecutionMode, ExecutionPrivileges, detect_execution_privileges};
17pub(crate) use mode::{root_privilege_drop_supported, unsupported_root_privilege_drop_error};
18use postgresql_embedded::Settings;
19use serde::{Deserialize, Serialize};
20
21use self::{
22    env::{shutdown_timeout_from_env, worker_binary_from_env},
23    mode::determine_execution_mode,
24    prepare::prepare_bootstrap,
25};
26use crate::{
27    PgEnvCfg,
28    error::{BootstrapResult, Result as CrateResult},
29};
30
31const DEFAULT_SETUP_TIMEOUT: Duration = Duration::from_secs(180);
32const DEFAULT_START_TIMEOUT: Duration = Duration::from_secs(60);
33
34#[cfg(test)]
35type SetupOnlyLifecycleHook =
36    Arc<dyn Fn(TestBootstrapSettings) -> BootstrapResult<()> + Send + Sync>;
37
38#[cfg(test)]
39static SETUP_ONLY_LIFECYCLE_HOOK: OnceLock<Mutex<Option<SetupOnlyLifecycleHook>>> = OnceLock::new();
40
41#[derive(Clone, Copy)]
42enum BootstrapKind {
43    Default,
44    Test,
45}
46
47/// Controls cleanup behaviour when a cluster is dropped.
48#[derive(Debug, Clone, Copy, Eq, PartialEq, Serialize, Deserialize, Default)]
49pub enum CleanupMode {
50    /// Remove only the data directory.
51    #[default]
52    DataOnly,
53    /// Remove both the data and installation directories.
54    Full,
55    /// Skip cleanup entirely (useful for debugging).
56    None,
57}
58
59/// Structured settings returned from [`bootstrap_for_tests`].
60#[derive(Debug, Clone)]
61pub struct TestBootstrapSettings {
62    /// Privilege level detected for the current process.
63    pub privileges: ExecutionPrivileges,
64    /// Strategy for executing `PostgreSQL` lifecycle commands.
65    pub execution_mode: ExecutionMode,
66    /// `PostgreSQL` configuration prepared for the embedded instance.
67    pub settings: Settings,
68    /// Environment variables required to exercise the embedded instance.
69    pub environment: TestBootstrapEnvironment,
70    /// Optional path to the helper binary used for subprocess execution.
71    pub worker_binary: Option<camino::Utf8PathBuf>,
72    /// Maximum time to allow the worker to complete the setup phase.
73    pub setup_timeout: Duration,
74    /// Maximum time to allow the worker to complete the start phase.
75    pub start_timeout: Duration,
76    /// Grace period granted to `PostgreSQL` during drop before teardown proceeds regardless.
77    pub shutdown_timeout: Duration,
78    /// Controls cleanup behaviour when the cluster drops.
79    pub cleanup_mode: CleanupMode,
80    /// Optional override for the binary cache directory.
81    ///
82    /// When set, `TestCluster` uses this directory instead of the default
83    /// resolved from environment variables.
84    pub binary_cache_dir: Option<camino::Utf8PathBuf>,
85}
86
87/// Bootstraps an embedded `PostgreSQL` instance, downloads the distribution,
88/// and initializes the data directory via `initdb`.
89///
90/// The server is **not** started — the resulting installation is ready for
91/// subsequent use by [`TestCluster`](crate::TestCluster) or other tools.
92///
93/// The function honours the following environment variables when present:
94/// - `PG_RUNTIME_DIR`: Overrides the `PostgreSQL` installation directory.
95/// - `PG_DATA_DIR`: Overrides the data directory used for initialization.
96/// - `PG_SUPERUSER`: Sets the superuser account name.
97/// - `PG_PASSWORD`: Supplies the superuser password.
98///
99/// When executed as `root` on Unix platforms the runtime drops privileges to the `nobody` user
100/// and prepares the filesystem on that user's behalf. Unprivileged executions reuse the current
101/// user identity. The function returns a [`crate::Error`] describing failures encountered during
102/// bootstrap or setup.
103///
104/// # Examples
105/// ```no_run
106/// use pg_embedded_setup_unpriv::run;
107///
108/// fn main() -> Result<(), pg_embedded_setup_unpriv::Error> {
109///     run()?;
110///     Ok(())
111/// }
112/// ```
113///
114/// # Errors
115/// Returns an error when bootstrap preparation fails, when the `PostgreSQL`
116/// distribution cannot be downloaded, or when `initdb` fails.
117pub fn run() -> CrateResult<()> {
118    let bootstrap = orchestrate_bootstrap(BootstrapKind::Default)?;
119    run_setup_only_lifecycle(bootstrap)
120}
121
122/// Bootstraps `PostgreSQL` for integration tests and surfaces the prepared settings.
123///
124/// # Examples
125/// ```no_run
126/// use pg_embedded_setup_unpriv::bootstrap_for_tests;
127/// use temp_env::with_vars;
128///
129/// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
130/// let bootstrap = bootstrap_for_tests()?;
131/// with_vars(
132///     bootstrap.environment.to_env(),
133///     || -> pg_embedded_setup_unpriv::BootstrapResult<()> {
134///         // Launch application logic that relies on `bootstrap.settings` here.
135///         Ok(())
136///     },
137/// )?;
138/// # Ok(())
139/// # }
140/// ```
141///
142/// # Errors
143/// Returns an error when bootstrap preparation fails or when subprocess orchestration
144/// cannot be configured.
145pub fn bootstrap_for_tests() -> BootstrapResult<TestBootstrapSettings> {
146    orchestrate_bootstrap(BootstrapKind::Test)
147}
148
149fn orchestrate_bootstrap(kind: BootstrapKind) -> BootstrapResult<TestBootstrapSettings> {
150    install_color_eyre();
151    if matches!(kind, BootstrapKind::Test) {
152        validate_backend_selection()?;
153    }
154
155    let privileges = detect_execution_privileges();
156    let cfg = PgEnvCfg::load().context("failed to load configuration via OrthoConfig")?;
157    let settings = match kind {
158        BootstrapKind::Default => cfg.to_settings()?,
159        BootstrapKind::Test => cfg.to_settings_for_tests()?,
160    };
161    let worker_binary = worker_binary_from_env(privileges)?;
162    let execution_mode = determine_execution_mode(privileges, worker_binary.as_ref())?;
163    let shutdown_timeout = shutdown_timeout_from_env()?;
164    let prepared = prepare_bootstrap(privileges, settings, &cfg)?;
165
166    Ok(TestBootstrapSettings {
167        privileges,
168        execution_mode,
169        settings: prepared.settings,
170        environment: prepared.environment,
171        worker_binary,
172        setup_timeout: DEFAULT_SETUP_TIMEOUT,
173        start_timeout: DEFAULT_START_TIMEOUT,
174        shutdown_timeout,
175        cleanup_mode: CleanupMode::default(),
176        binary_cache_dir: cfg.binary_cache_dir,
177    })
178}
179
180fn install_color_eyre() {
181    if let Err(err) = color_eyre::install() {
182        tracing::debug!("color_eyre already installed: {err}");
183    }
184}
185
186fn validate_backend_selection() -> BootstrapResult<()> {
187    let raw = std::env::var_os("PG_TEST_BACKEND").unwrap_or_default();
188    let value = raw.to_string_lossy();
189    let trimmed = value.trim();
190    if trimmed.is_empty() || trimmed == "postgresql_embedded" {
191        return Ok(());
192    }
193    Err(eyre!(
194        "SKIP-TEST-CLUSTER: unsupported PG_TEST_BACKEND '{trimmed}'; supported backends: \
195         postgresql_embedded"
196    )
197    .into())
198}
199
200/// Executes the setup-only lifecycle for CLI invocations.
201///
202/// Keeping this as a dedicated helper keeps the public `run()` flow linear:
203/// bootstrap preparation first, then setup lifecycle execution.
204fn run_setup_only_lifecycle(bootstrap: TestBootstrapSettings) -> CrateResult<()> {
205    #[cfg(test)]
206    {
207        let installed_setup_hook = SETUP_ONLY_LIFECYCLE_HOOK
208            .get_or_init(|| Mutex::new(None))
209            .lock()
210            .unwrap_or_else(std::sync::PoisonError::into_inner)
211            .clone();
212        if let Some(setup_hook) = installed_setup_hook {
213            return setup_hook(bootstrap).map_err(Into::into);
214        }
215    }
216
217    crate::cluster::setup_postgres_only(bootstrap)?;
218    Ok(())
219}
220
221#[cfg(test)]
222pub(crate) struct SetupOnlyLifecycleHookGuard;
223
224#[cfg(test)]
225impl Drop for SetupOnlyLifecycleHookGuard {
226    fn drop(&mut self) {
227        let mut slot = SETUP_ONLY_LIFECYCLE_HOOK
228            .get_or_init(|| Mutex::new(None))
229            .lock()
230            .unwrap_or_else(std::sync::PoisonError::into_inner);
231        slot.take();
232    }
233}
234
235#[cfg(test)]
236pub(crate) fn install_setup_only_lifecycle_hook<F>(hook: F) -> SetupOnlyLifecycleHookGuard
237where
238    F: Fn(TestBootstrapSettings) -> BootstrapResult<()> + Send + Sync + 'static,
239{
240    let mut slot = SETUP_ONLY_LIFECYCLE_HOOK
241        .get_or_init(|| Mutex::new(None))
242        .lock()
243        .unwrap_or_else(std::sync::PoisonError::into_inner);
244    *slot = Some(Arc::new(hook));
245    SetupOnlyLifecycleHookGuard
246}
247
248#[cfg(test)]
249mod mod_tests;