improv_storage_sql/
pg.rs

1//! Postgres backend (synchronous `postgres` crate).
2//!
3//! Credentials are **out of band**: `connect_postgres` takes a full connection
4//! string that the caller assembled at connect time (see `conn::open`, which
5//! injects the password from an env var). This module never logs a connection
6//! string — a URI can carry a secret, so errors are redacted.
7
8use crate::backend::{Backend, Cell, Table};
9use crate::{Result, SqlError};
10use postgres::config::{Host, SslMode};
11use postgres::types::Type;
12use postgres::{Client, Config, NoTls};
13
14/// Open a Postgres connection from a full libpq/URI connection string.
15///
16/// The string may embed a password; **do not log it**. Prefer building it via
17/// `conn::open`, which pulls the secret from the environment. On failure the
18/// error is redacted so no secret leaks into logs.
19///
20/// **No TLS is built in.** To keep the password off the wire in cleartext the
21/// connection is refused (fail closed, before any network I/O) unless every
22/// host is loopback or a Unix socket, or the string explicitly says
23/// `sslmode=disable`. `sslmode=require` is refused too, since it cannot be
24/// honoured. Reach a remote server over an SSH tunnel or a local proxy.
25pub fn connect_postgres(conn_str: &str) -> Result<Backend> {
26    let config: Config = conn_str
27        .parse()
28        .map_err(|_| SqlError::Connect("invalid postgres connection string (redacted)".into()))?;
29    check_pg_transport(&config)?;
30    let client = config
31        .connect(NoTls)
32        .map_err(|e| SqlError::Connect(format!("postgres connect failed: {}", redact_pg(e))))?;
33    Ok(Backend::Postgres(client))
34}
35
36/// Refuse a connection that would send credentials in cleartext off-host, or
37/// that asks for TLS this build cannot provide. See [`connect_postgres`].
38fn check_pg_transport(config: &Config) -> Result<()> {
39    let no_tls = |why: &str| {
40        Err(SqlError::Connect(format!(
41            "postgres: TLS isn't built into this build of Improv; {why}"
42        )))
43    };
44    let mode = config.get_ssl_mode();
45    if mode == SslMode::Require {
46        return no_tls("sslmode=require cannot be honoured (use an SSH tunnel to a loopback port)");
47    }
48    let host_is_local = |h: &Host| match h {
49        Host::Tcp(name) => {
50            name.eq_ignore_ascii_case("localhost")
51                || name
52                    .parse::<std::net::IpAddr>()
53                    .is_ok_and(|ip| ip.is_loopback())
54        }
55        #[cfg(unix)]
56        Host::Unix(_) => true,
57    };
58    let local = config.get_hosts().iter().all(host_is_local)
59        && config.get_hostaddrs().iter().all(|ip| ip.is_loopback());
60    if !local && mode != SslMode::Disable {
61        return no_tls(
62            "refusing a non-loopback host without an explicit sslmode=disable \
63             (the password would cross the network in cleartext)",
64        );
65    }
66    Ok(())
67}
68
69/// Best-effort scrub of a driver error so a password embedded in a connection
70/// string never reaches logs. We keep only the error's category, not its text.
71fn redact_pg(e: postgres::Error) -> String {
72    // postgres::Error's Display can echo back the DSN on some failures; return a
73    // fixed string rather than risk leaking it.
74    if e.as_db_error().is_some() {
75        "database error (redacted)".to_string()
76    } else {
77        "connection error (redacted)".to_string()
78    }
79}
80
81pub(crate) fn pg_query(client: &mut Client, sql: &str) -> Result<Table> {
82    // Run the query inside a READ ONLY transaction so even a data-modifying
83    // CTE (`WITH x AS (DELETE ...) SELECT ...`) is refused by the server. The
84    // transaction is rolled back on drop if anything below fails.
85    let pg_err =
86        |what: &str, e| SqlError::Import(format!("postgres {what} failed: {}", redact_pg(e)));
87    let mut tx = client
88        .build_transaction()
89        .read_only(true)
90        .start()
91        .map_err(|e| pg_err("begin read only", e))?;
92    // Prepare once: the statement carries the column names even for an empty
93    // result.
94    let stmt = tx.prepare(sql).map_err(|e| pg_err("prepare", e))?;
95    let rows = tx.query(&stmt, &[]).map_err(|e| pg_err("query", e))?;
96    let columns: Vec<String> = stmt
97        .columns()
98        .iter()
99        .map(|c| c.name().to_string())
100        .collect();
101    tx.commit().map_err(|e| pg_err("commit", e))?;
102
103    let mut out_rows = Vec::with_capacity(rows.len());
104    for row in &rows {
105        let mut cells = Vec::with_capacity(row.len());
106        for i in 0..row.len() {
107            cells.push(pg_cell(row, i)?);
108        }
109        out_rows.push(cells);
110    }
111    Ok(Table {
112        columns,
113        rows: out_rows,
114    })
115}
116
117/// Map one Postgres cell to a [`Cell`], dispatching on the column's SQL type.
118/// Numeric-family types → Int/Float, bool → Bool, everything textual → Text.
119fn pg_cell(row: &postgres::Row, i: usize) -> Result<Cell> {
120    let ty = row.columns()[i].type_();
121    let mapped = match *ty {
122        Type::BOOL => row.try_get::<_, Option<bool>>(i).map(|o| match o {
123            Some(b) => Cell::Bool(b),
124            None => Cell::Null,
125        }),
126        Type::INT2 => row
127            .try_get::<_, Option<i16>>(i)
128            .map(|o| o.map(|v| Cell::Int(v as i64)).unwrap_or(Cell::Null)),
129        Type::INT4 => row
130            .try_get::<_, Option<i32>>(i)
131            .map(|o| o.map(|v| Cell::Int(v as i64)).unwrap_or(Cell::Null)),
132        Type::INT8 => row
133            .try_get::<_, Option<i64>>(i)
134            .map(|o| o.map(Cell::Int).unwrap_or(Cell::Null)),
135        Type::FLOAT4 => row
136            .try_get::<_, Option<f32>>(i)
137            .map(|o| o.map(|v| Cell::Float(v as f64)).unwrap_or(Cell::Null)),
138        Type::FLOAT8 => row
139            .try_get::<_, Option<f64>>(i)
140            .map(|o| o.map(Cell::Float).unwrap_or(Cell::Null)),
141        // NUMERIC arrives in its binary wire form; decode it to Float.
142        Type::NUMERIC => row
143            .try_get::<_, Option<PgNumeric>>(i)
144            .map(|o| o.map(|n| Cell::Float(n.0)).unwrap_or(Cell::Null)),
145        // Text-family (TEXT/VARCHAR/BPCHAR/NAME/…): read as String.
146        _ => row
147            .try_get::<_, Option<String>>(i)
148            .map(|o| o.map(Cell::Text).unwrap_or(Cell::Null)),
149    };
150    mapped.map_err(|e| {
151        SqlError::Import(format!(
152            "postgres column {i} ({ty}) unmapped: {}",
153            redact_str(&e.to_string())
154        ))
155    })
156}
157
158/// A Postgres `NUMERIC` read as `f64` (precision beyond f64 is lost, which is
159/// what the model's numeric cells hold anyway).
160struct PgNumeric(f64);
161
162impl<'a> postgres::types::FromSql<'a> for PgNumeric {
163    /// Binary layout: ndigits i16, weight i16, sign u16, dscale u16, then
164    /// `ndigits` base-10000 digits (i16), the first worth `10000^weight`.
165    fn from_sql(
166        _ty: &Type,
167        raw: &'a [u8],
168    ) -> std::result::Result<Self, Box<dyn std::error::Error + Sync + Send>> {
169        let word =
170            |k: usize| -> std::result::Result<i16, Box<dyn std::error::Error + Sync + Send>> {
171                raw.get(2 * k..2 * k + 2)
172                    .map(|b| i16::from_be_bytes([b[0], b[1]]))
173                    .ok_or_else(|| "truncated numeric".into())
174            };
175        let (ndigits, weight, sign) = (word(0)?, word(1)?, word(2)? as u16);
176        if sign == 0xC000 {
177            return Ok(PgNumeric(f64::NAN));
178        }
179        let mut v = 0.0;
180        for k in 0..ndigits.max(0) as usize {
181            v += f64::from(word(4 + k)?) * 10000f64.powi(i32::from(weight) - k as i32);
182        }
183        Ok(PgNumeric(if sign == 0x4000 { -v } else { v }))
184    }
185
186    fn accepts(ty: &Type) -> bool {
187        *ty == Type::NUMERIC
188    }
189}
190
191pub(crate) fn pg_execute(client: &mut Client, sql: &str, params: &[Cell]) -> Result<()> {
192    // Bind each Cell as a &dyn ToSql. Build owned adapters, then borrow them.
193    let owned: Vec<PgParam> = params.iter().map(PgParam::from_cell).collect();
194    let bound: Vec<&(dyn postgres::types::ToSql + Sync)> = owned
195        .iter()
196        .map(|p| p as &(dyn postgres::types::ToSql + Sync))
197        .collect();
198    client
199        .execute(sql, &bound)
200        .map_err(|e| SqlError::Export(format!("postgres execute failed: {}", redact_pg(e))))?;
201    Ok(())
202}
203
204/// Owning parameter adapter so we can hand `&dyn ToSql` to the driver.
205#[derive(Debug)]
206enum PgParam {
207    Null,
208    Int(i64),
209    Float(f64),
210    Text(String),
211    Bool(bool),
212}
213
214impl PgParam {
215    fn from_cell(c: &Cell) -> PgParam {
216        match c {
217            Cell::Null => PgParam::Null,
218            Cell::Int(i) => PgParam::Int(*i),
219            Cell::Float(f) => PgParam::Float(*f),
220            Cell::Text(t) => PgParam::Text(t.clone()),
221            Cell::Bool(b) => PgParam::Bool(*b),
222        }
223    }
224}
225
226impl postgres::types::ToSql for PgParam {
227    fn to_sql(
228        &self,
229        ty: &Type,
230        out: &mut postgres::types::private::BytesMut,
231    ) -> std::result::Result<postgres::types::IsNull, Box<dyn std::error::Error + Sync + Send>>
232    {
233        match self {
234            PgParam::Null => Ok(postgres::types::IsNull::Yes),
235            PgParam::Int(i) => i.to_sql(ty, out),
236            PgParam::Float(f) => f.to_sql(ty, out),
237            PgParam::Text(t) => t.to_sql(ty, out),
238            PgParam::Bool(b) => b.to_sql(ty, out),
239        }
240    }
241
242    fn accepts(_ty: &Type) -> bool {
243        // Accept whatever the target column is; the concrete to_sql above will
244        // error if the runtime type truly can't encode.
245        true
246    }
247
248    postgres::types::to_sql_checked!();
249}
250
251/// Redact anything that looks like it could carry a secret from a free-form
252/// string before it reaches an error/log. Cheap and conservative.
253fn redact_str(s: &str) -> String {
254    if s.contains("password") || s.contains("://") {
255        "(redacted)".to_string()
256    } else {
257        s.to_string()
258    }
259}
260
261#[cfg(test)]
262#[cfg_attr(coverage_nightly, coverage(off))]
263mod tests {
264    use super::*;
265
266    fn transport(dsn: &str) -> Result<()> {
267        check_pg_transport(&dsn.parse().unwrap())
268    }
269
270    #[test]
271    fn transport_allows_loopback_and_sockets() {
272        for ok in [
273            "host=localhost user=u",
274            "host=127.0.0.1 user=u",
275            "host=::1 user=u",
276            "postgresql://u@localhost:5433/db",
277            "host=db.example.com user=u sslmode=disable",
278        ] {
279            assert!(transport(ok).is_ok(), "{ok:?}");
280        }
281        // A socket directory is a `Host::Unix` only on unix; elsewhere the
282        // driver parses it as a TCP host name, which is not loopback.
283        #[cfg(unix)]
284        assert!(transport("host=/var/run/postgresql user=u").is_ok());
285    }
286
287    #[test]
288    fn transport_refuses_cleartext_remote_and_unhonourable_tls() {
289        for bad in [
290            "host=db.example.com user=u",
291            "host=10.0.0.5 user=u sslmode=prefer",
292            "postgresql://u@db.example.com/db",
293            "host=localhost,db.example.com user=u",
294            "host=localhost hostaddr=10.0.0.5 user=u",
295            "host=localhost user=u sslmode=require",
296        ] {
297            let err = transport(bad).unwrap_err().to_string();
298            assert!(err.contains("TLS isn't built in"), "{bad:?}: {err}");
299        }
300        // Refused before any network I/O, through the public entry point.
301        assert!(connect_postgres("host=db.example.com user=u").is_err());
302        assert!(connect_postgres("not a dsn ===").is_err());
303    }
304
305    #[test]
306    fn redact_str_scrubs_secrets() {
307        assert_eq!(redact_str("password=x"), "(redacted)");
308        assert_eq!(redact_str("postgres://u:p@h"), "(redacted)");
309        assert_eq!(redact_str("plain"), "plain");
310    }
311
312    /// Live: needs `IMPROV_TEST_PG` (a libpq DSN to a scratch database).
313    fn live() -> Option<Backend> {
314        let Ok(dsn) = std::env::var("IMPROV_TEST_PG") else {
315            println!("skipped: IMPROV_TEST_PG not set");
316            return None;
317        };
318        Some(connect_postgres(&dsn).expect("connect IMPROV_TEST_PG"))
319    }
320
321    #[test]
322    fn pg_maps_every_column_type() {
323        let Some(mut be) = live() else { return };
324        let t = be
325            .query(
326                "SELECT true AS b, 2::int2 AS s, 3::int4 AS i, 4::int8 AS l, \
327                 1.5::float4 AS f, 2.5::float8 AS d, 3.25::numeric AS n, \
328                 'x'::text AS t, NULL::int4 AS z, NULL::text AS zt",
329            )
330            .unwrap();
331        assert_eq!(
332            t.columns,
333            ["b", "s", "i", "l", "f", "d", "n", "t", "z", "zt"]
334        );
335        assert_eq!(
336            t.rows,
337            vec![vec![
338                Cell::Bool(true),
339                Cell::Int(2),
340                Cell::Int(3),
341                Cell::Int(4),
342                Cell::Float(1.5),
343                Cell::Float(2.5),
344                Cell::Float(3.25),
345                Cell::Text("x".into()),
346                Cell::Null,
347                Cell::Null,
348            ]]
349        );
350        // An empty result still reports its columns.
351        let e = be.query("SELECT 1 AS a, 'q' AS b WHERE false").unwrap();
352        assert_eq!(e.columns, ["a", "b"]);
353        assert!(e.rows.is_empty());
354        let n = be
355            .query("SELECT -12345.678::numeric AS a, 0::numeric AS b, 'NaN'::numeric AS c")
356            .unwrap();
357        assert_eq!(n.rows[0][..2], [Cell::Float(-12345.678), Cell::Float(0.0)]);
358        assert!(matches!(n.rows[0][2], Cell::Float(f) if f.is_nan()));
359        // An unmappable type is an error, not a panic.
360        assert!(be.query("SELECT now() AS ts").is_err());
361        assert!(be.query("SELECT * FROM no_such_table").is_err());
362    }
363
364    #[test]
365    fn pg_refresh_query_runs_read_only() {
366        let Some(mut be) = live() else { return };
367        be.execute("DROP TABLE IF EXISTS improv_pg_ro", &[])
368            .unwrap();
369        be.execute("CREATE TABLE improv_pg_ro (v int8)", &[])
370            .unwrap();
371        be.execute("INSERT INTO improv_pg_ro VALUES ($1)", &[Cell::Int(7)])
372            .unwrap();
373        // A data-modifying CTE passes the keyword check but the READ ONLY
374        // transaction refuses it server-side.
375        let cte = "WITH x AS (DELETE FROM improv_pg_ro RETURNING *) SELECT * FROM x";
376        assert!(crate::check_read_only_sql(cte).is_ok());
377        assert!(be.query(cte).is_err());
378        assert!(be.query("DELETE FROM improv_pg_ro").is_err());
379        // The connection is still usable and the row survived.
380        let t = be.query("SELECT v FROM improv_pg_ro").unwrap();
381        assert_eq!(t.rows, vec![vec![Cell::Int(7)]]);
382        be.execute("DROP TABLE improv_pg_ro", &[]).unwrap();
383    }
384
385    #[test]
386    fn numeric_decoder_rejects_truncated_input() {
387        use postgres::types::FromSql;
388        let err = PgNumeric::from_sql(&Type::NUMERIC, &[0, 1]).err().unwrap();
389        assert_eq!(err.to_string(), "truncated numeric");
390        // ndigits=1 but the digit word is missing.
391        let err = PgNumeric::from_sql(&Type::NUMERIC, &[0, 1, 0, 0, 0, 0, 0, 0])
392            .err()
393            .unwrap();
394        assert_eq!(err.to_string(), "truncated numeric");
395        assert!(PgNumeric::accepts(&Type::NUMERIC) && !PgNumeric::accepts(&Type::TEXT));
396    }
397
398    /// Nothing listens on loopback port 1: a non-database (I/O) error, redacted.
399    #[test]
400    fn refused_connection_is_redacted() {
401        let Err(err) = connect_postgres("host=127.0.0.1 port=1 user=u password=hunter2") else {
402            panic!("connected to port 1");
403        };
404        let msg = err.to_string();
405        assert!(msg.contains("connection error (redacted)"), "{msg}");
406        assert!(!msg.contains("hunter2"), "{msg}");
407    }
408
409    #[test]
410    fn pg_binds_null_and_bool_and_redacts_execute_errors() {
411        let Some(mut be) = live() else { return };
412        be.execute("DROP TABLE IF EXISTS improv_pg_params", &[])
413            .unwrap();
414        be.execute("CREATE TABLE improv_pg_params (a int8, b bool)", &[])
415            .unwrap();
416        be.execute(
417            "INSERT INTO improv_pg_params VALUES ($1, $2), ($3, $4)",
418            &[Cell::Null, Cell::Bool(true), Cell::Int(1), Cell::Null],
419        )
420        .unwrap();
421        let t = be
422            .query("SELECT a, b FROM improv_pg_params ORDER BY a NULLS FIRST")
423            .unwrap();
424        assert_eq!(
425            t.rows,
426            [[Cell::Null, Cell::Bool(true)], [Cell::Int(1), Cell::Null]]
427        );
428        let err = be
429            .execute("INSERT INTO no_such_table VALUES (1)", &[])
430            .unwrap_err();
431        assert!(
432            err.to_string().contains("database error (redacted)"),
433            "{err}"
434        );
435        be.execute("DROP TABLE improv_pg_params", &[]).unwrap();
436    }
437}
improv_storage_sql/pg.rs

improv_storage_sql/
pg.rs
