pg_embedded_setup_unpriv/cluster/
connection.rs

1//! Connection helpers for `TestCluster`, including metadata accessors and optional Diesel support.
2
3use camino::{Utf8Path, Utf8PathBuf};
4#[cfg(feature = "diesel-support")]
5use color_eyre::eyre::WrapErr;
6use postgres::{Client, NoTls};
7use postgresql_embedded::Settings;
8
9use crate::{TestBootstrapSettings, error::BootstrapResult};
10
11/// Escapes a SQL identifier by doubling embedded double quotes.
12///
13/// `PostgreSQL` identifiers are quoted with double quotes. Any embedded
14/// double quote must be escaped by doubling it.
15pub(crate) fn escape_identifier(name: &str) -> String { name.replace('"', "\"\"") }
16
17/// Creates a new `PostgreSQL` client connection from the given URL.
18///
19/// This is a shared helper for admin database connections used by both
20/// `TestClusterConnection` and `TemporaryDatabase`.
21pub(crate) fn connect_admin(url: &str) -> BootstrapResult<Client> {
22    Client::connect(url, NoTls).map_err(admin_connect_error)
23}
24
25/// Wraps a failed admin connection, keeping the driver's error as the source.
26///
27/// `tokio_postgres` displays a server-side failure as just `db error`, with
28/// the reason (such as `password authentication failed`, code `28P01`) in its
29/// source. The message therefore names that source too, and the error itself
30/// stays in the chain for callers that walk it.
31pub(crate) fn admin_connect_error(err: postgres::Error) -> crate::error::BootstrapError {
32    let reason = std::error::Error::source(&err)
33        .map_or_else(|| err.to_string(), |source| format!("{err}: {source}"));
34    let message = format!("failed to connect to admin database: {reason}");
35    crate::error::BootstrapError::from(color_eyre::Report::new(err).wrap_err(message))
36}
37
38/// Provides ergonomic accessors for connection-oriented cluster metadata.
39///
40/// # Examples
41/// ```no_run
42/// use pg_embedded_setup_unpriv::TestCluster;
43///
44/// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
45/// let cluster = TestCluster::new()?;
46/// let metadata = cluster.connection().metadata();
47/// assert_eq!(metadata.host(), "localhost");
48/// # Ok(())
49/// # }
50/// ```
51#[derive(Debug, Clone)]
52pub struct ConnectionMetadata {
53    settings: Settings,
54    pgpass_file: Utf8PathBuf,
55}
56
57impl ConnectionMetadata {
58    pub(crate) fn from_settings(settings: &TestBootstrapSettings) -> Self {
59        Self {
60            settings: settings.settings.clone(),
61            pgpass_file: settings.environment.pgpass_file.clone(),
62        }
63    }
64
65    /// Returns the configured database host.
66    #[must_use]
67    pub const fn host(&self) -> &str { self.settings.host.as_str() }
68
69    /// Returns the configured port.
70    #[must_use]
71    pub const fn port(&self) -> u16 { self.settings.port }
72
73    /// Returns the configured superuser name.
74    #[must_use]
75    pub const fn superuser(&self) -> &str { self.settings.username.as_str() }
76
77    /// Returns the generated superuser password.
78    #[must_use]
79    pub const fn password(&self) -> &str { self.settings.password.as_str() }
80
81    /// Returns the prepared `.pgpass` file path.
82    #[must_use]
83    pub fn pgpass_file(&self) -> &Utf8Path { self.pgpass_file.as_ref() }
84
85    /// Constructs a libpq-compatible URL for `database` using the underlying
86    /// `postgresql_embedded` helper.
87    #[must_use]
88    pub fn database_url(&self, database: &str) -> String { self.settings.url(database) }
89}
90
91/// Accessor for connection helpers derived from a
92/// [`TestCluster`](crate::TestCluster).
93///
94/// Enable the `diesel-support` feature to call the Diesel connection helper.
95///
96/// # Examples
97/// ```no_run
98/// use pg_embedded_setup_unpriv::TestCluster;
99///
100/// # fn main() -> pg_embedded_setup_unpriv::BootstrapResult<()> {
101/// let cluster = TestCluster::new()?;
102/// let url = cluster.connection().database_url("postgres");
103/// assert!(url.contains("postgresql://"));
104/// # Ok(())
105/// # }
106/// ```
107#[derive(Debug, Clone)]
108pub struct TestClusterConnection {
109    metadata: ConnectionMetadata,
110}
111
112impl TestClusterConnection {
113    pub(crate) fn new(settings: &TestBootstrapSettings) -> Self {
114        Self {
115            metadata: ConnectionMetadata::from_settings(settings),
116        }
117    }
118
119    /// Returns host metadata without exposing internal storage.
120    #[must_use]
121    pub const fn host(&self) -> &str { self.metadata.host() }
122
123    /// Returns the configured port.
124    #[must_use]
125    pub const fn port(&self) -> u16 { self.metadata.port() }
126
127    /// Returns the configured superuser account name.
128    #[must_use]
129    pub const fn superuser(&self) -> &str { self.metadata.superuser() }
130
131    /// Returns the generated password for the superuser.
132    #[must_use]
133    pub const fn password(&self) -> &str { self.metadata.password() }
134
135    /// Returns the `.pgpass` file prepared during bootstrap.
136    #[must_use]
137    pub fn pgpass_file(&self) -> &Utf8Path { self.metadata.pgpass_file() }
138
139    /// Provides an owned snapshot of the connection metadata.
140    #[must_use]
141    pub fn metadata(&self) -> ConnectionMetadata { self.metadata.clone() }
142
143    /// Builds a libpq-compatible database URL for `database`.
144    #[must_use]
145    pub fn database_url(&self, database: &str) -> String { self.metadata.database_url(database) }
146
147    /// Establishes a Diesel connection for the target `database`.
148    ///
149    /// # Errors
150    /// Returns a [`crate::error::BootstrapError`] when Diesel cannot connect.
151    #[cfg(feature = "diesel-support")]
152    pub fn diesel_connection(&self, database: &str) -> BootstrapResult<diesel::PgConnection> {
153        use diesel::Connection;
154
155        diesel::PgConnection::establish(&self.database_url(database))
156            .wrap_err(format!("failed to connect to {database} via Diesel"))
157            .map_err(crate::error::BootstrapError::from)
158    }
159
160    /// Connects to the `postgres` administration database.
161    pub(super) fn admin_client(&self) -> BootstrapResult<Client> {
162        connect_admin(&self.database_url("postgres"))
163    }
164}
165
166#[cfg(test)]
167mod tests {
168    //! Tests for cluster connection settings.
169    use std::time::Duration;
170
171    use postgresql_embedded::Settings;
172
173    use super::*;
174    use crate::{
175        CleanupMode,
176        TestBootstrapSettings,
177        bootstrap::{ExecutionMode, ExecutionPrivileges, TestBootstrapEnvironment},
178    };
179
180    fn sample_settings() -> TestBootstrapSettings {
181        let settings = Settings {
182            host: "127.0.0.1".into(),
183            port: 55_321,
184            username: "fixture_user".into(),
185            password: "fixture_pass".into(),
186            data_dir: "/tmp/cluster-data".into(),
187            installation_dir: "/tmp/cluster-install".into(),
188            ..Settings::default()
189        };
190
191        TestBootstrapSettings {
192            privileges: ExecutionPrivileges::Unprivileged,
193            execution_mode: ExecutionMode::InProcess,
194            settings,
195            environment: TestBootstrapEnvironment {
196                home: Utf8PathBuf::from("/tmp/home"),
197                xdg_cache_home: Utf8PathBuf::from("/tmp/home/cache"),
198                xdg_runtime_dir: Utf8PathBuf::from("/tmp/home/run"),
199                pgpass_file: Utf8PathBuf::from("/tmp/home/.pgpass"),
200                tz_dir: Some(Utf8PathBuf::from("/usr/share/zoneinfo")),
201                timezone: "UTC".into(),
202            },
203            worker_binary: None,
204            setup_timeout: Duration::from_secs(1),
205            start_timeout: Duration::from_secs(1),
206            shutdown_timeout: Duration::from_secs(1),
207            cleanup_mode: CleanupMode::default(),
208            binary_cache_dir: None,
209            extensions: None,
210            installed_extensions: Vec::new(),
211        }
212    }
213
214    #[test]
215    fn metadata_reflects_underlying_settings() {
216        let settings = sample_settings();
217        let connection = TestClusterConnection::new(&settings);
218        let metadata = connection.metadata();
219
220        assert_eq!(metadata.host(), "127.0.0.1");
221        assert_eq!(metadata.port(), 55_321);
222        assert_eq!(metadata.superuser(), "fixture_user");
223        assert_eq!(metadata.password(), "fixture_pass");
224        assert_eq!(metadata.pgpass_file(), Utf8Path::new("/tmp/home/.pgpass"));
225    }
226
227    #[test]
228    fn database_url_matches_postgresql_embedded() {
229        let settings = sample_settings();
230        let connection = TestClusterConnection::new(&settings);
231        let expected = settings.settings.url("postgres");
232
233        assert_eq!(connection.database_url("postgres"), expected);
234    }
235
236    /// Listens on a loopback port and serves one connection that rejects the
237    /// password, returning the address to connect to.
238    fn reject_password_once() -> std::io::Result<std::net::SocketAddr> {
239        let listener = std::net::TcpListener::bind("127.0.0.1:0")?;
240        let address = listener.local_addr()?;
241        std::thread::spawn(move || {
242            // A failure here leaves the client with a transport error, which
243            // the test's message assertion then reports.
244            let _served = answer_with_28p01(&listener);
245        });
246        Ok(address)
247    }
248
249    /// Reads one startup message and answers it with a server-side `28P01`
250    /// error, as a server rejecting the password does.
251    #[expect(
252        clippy::big_endian_bytes,
253        reason = "the PostgreSQL wire protocol frames every length in network byte order"
254    )]
255    fn answer_with_28p01(listener: &std::net::TcpListener) -> std::io::Result<()> {
256        use std::io::{Error, Read, Write};
257
258        let (mut stream, _) = listener.accept()?;
259        let mut length = [0_u8; 4];
260        stream.read_exact(&mut length)?;
261        let body_len = usize::try_from(u32::from_be_bytes(length))
262            .map_err(Error::other)?
263            .saturating_sub(4);
264        let mut body = vec![0_u8; body_len];
265        stream.read_exact(&mut body)?;
266        let mut fields = Vec::new();
267        for (code, value) in [
268            (b'S', "FATAL"),
269            (b'V', "FATAL"),
270            (b'C', "28P01"),
271            (b'M', "password authentication failed for user \"postgres\""),
272        ] {
273            fields.push(code);
274            fields.extend_from_slice(value.as_bytes());
275            fields.push(0);
276        }
277        fields.push(0);
278        let reply_len = u32::try_from(fields.len() + 4).map_err(Error::other)?;
279        let mut reply = vec![b'E'];
280        reply.extend_from_slice(&reply_len.to_be_bytes());
281        reply.extend_from_slice(&fields);
282        stream.write_all(&reply)
283    }
284
285    /// A server-side failure, which `tokio_postgres` displays as just
286    /// `db error`, reaches the message with its cause and stays in the chain.
287    #[test]
288    fn a_rejected_password_names_its_cause() {
289        let address = reject_password_once().expect("a loopback port");
290        let url = format!("postgresql://postgres:wrong@{address}/postgres");
291
292        let err = connect_admin(&url)
293            .err()
294            .expect("the server rejects the password");
295
296        let message = err.to_string();
297        assert!(
298            message.contains("db error: FATAL: password authentication failed"),
299            "the nested cause must reach the message: {message}"
300        );
301        let report = err.into_report();
302        let driver = report
303            .chain()
304            .find_map(|cause| cause.downcast_ref::<postgres::Error>())
305            .expect("the postgres error stays in the chain");
306        let code = driver.code().map(postgres::error::SqlState::code);
307        assert_eq!(
308            code,
309            Some("28P01"),
310            "the server's code survives: {report:?}"
311        );
312    }
313
314    /// A failed admin connection keeps the driver's error in the chain and
315    /// names its cause, so a caller sees why rather than a bare `db error`.
316    #[test]
317    fn a_failed_admin_connection_keeps_its_source() {
318        // Port 1 on loopback refuses at once; no server is involved.
319        let err = connect_admin("postgresql://postgres:x@127.0.0.1:1/postgres")
320            .err()
321            .expect("nothing listens on port 1");
322
323        let message = err.to_string();
324        let report = err.into_report();
325        let chain_has_driver_error = report
326            .chain()
327            .any(|cause| cause.downcast_ref::<postgres::Error>().is_some());
328        assert!(chain_has_driver_error, "the source was dropped: {report:?}");
329        assert!(
330            message.starts_with("failed to connect to admin database: "),
331            "{message}"
332        );
333        assert!(
334            message.to_lowercase().contains("refused"),
335            "the cause must reach the message: {message}"
336        );
337    }
338}
pg_embedded_setup_unpriv/cluster/connection.rs

pg_embedded_setup_unpriv/cluster/
connection.rs
