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