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