From 2a4ad76bda85efd2615529f041548d0516dd6a2f Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Thu, 2 Jul 2026 14:26:38 +0000 Subject: [PATCH] Add backend-agnostic Rust-native query helpers over the pool MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce a small enum-based layer so simple queries can be written once regardless of backend: DbPool/DbConn dispatch over each backend's own pool, and DbConn::query(sql, params) binds positional `?` placeholders and returns rows as backend-agnostic DbValue cells (read out with DbRowExt::try_get). No trait/async_trait machinery — a closed enum keeps the pool.get().await? / conn.query(...).await? surface simple. The Postgres arm rewrites `?` to `$1, $2, ...` (matching convert_param_style), reuses PgValue for parameter binding, and adds a Rust-native DbValueFromSql decoder (the non-Python counterpart of PythonPgFromSql). The shared column-type list is factored into accepts_column_type so the three mappings can't drift. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/src/database/mod.rs | 142 ++++++++++++++++++- rust/src/database/postgres/mod.rs | 1 + rust/src/database/postgres/pool.rs | 6 +- rust/src/database/postgres/query.rs | 82 +++++++++++ rust/src/database/postgres/value.rs | 132 +++++++++++++----- rust/src/database/value.rs | 208 ++++++++++++++++++++++++++++ 6 files changed, 536 insertions(+), 35 deletions(-) create mode 100644 rust/src/database/postgres/query.rs create mode 100644 rust/src/database/value.rs diff --git a/rust/src/database/mod.rs b/rust/src/database/mod.rs index 2d1c63310a..bb660e4044 100644 --- a/rust/src/database/mod.rs +++ b/rust/src/database/mod.rs @@ -1,14 +1,68 @@ -//! DBAPI2-shaped Connection / Cursor types implemented in Rust. +//! Rust database access. //! -//! Currently this provides a single Postgres backend ([`tokio_postgres`]); see -//! the [`postgres`] submodule. +//! Two layers live here: +//! - the DBAPI2-shaped `Connection`/`Cursor` shim (see the [`postgres`] +//! submodule), which lets existing Python transaction functions run against a +//! Rust connection unchanged, and +//! - the backend-agnostic Rust-native query helpers ([`DbPool`]/[`DbConn`] plus +//! [`value::DbValue`]), so simple queries can be written once and run against +//! any backend. +//! +//! Currently the only backend is Postgres ([`tokio_postgres`]); SQLite is +//! expected to follow as a sibling variant of [`DbPool`]/[`DbConn`]. pub mod postgres; pub mod runtime; +pub mod value; use pyo3::prelude::*; use pyo3::types::PyModule; +use crate::database::value::{DbRow, DbValue}; + +/// A backend-agnostic connection pool. +/// +/// A thin enum over each backend's own pool rather than a trait: with a small, +/// closed set of backends this keeps the ergonomic `pool.get().await?` / +/// `conn.query(...).await?` surface without `dyn`/`async_trait` machinery. +pub enum DbPool { + Postgres(postgres::pool::ConnectionPool), +} + +impl DbPool { + /// Check a connection out of the pool. + pub async fn get(&self) -> Result { + match self { + DbPool::Postgres(pool) => Ok(DbConn::Postgres(pool.get().await?)), + } + } +} + +/// A backend-agnostic connection checked out of a [`DbPool`]. +/// +/// Returned to the pool when dropped. +pub enum DbConn { + Postgres(postgres::pool::PooledConnection), +} + +impl DbConn { + /// Run a query and return all resulting rows. + /// + /// Parameters are bound positionally to `?` placeholders (Synapse's existing + /// convention); each backend rewrites them to its own placeholder syntax. + /// Result cells come back as [`DbValue`]s, read out with + /// [`value::DbRowExt::try_get`]. + pub async fn query( + &mut self, + sql: &str, + params: &[DbValue], + ) -> Result, anyhow::Error> { + match self { + DbConn::Postgres(conn) => postgres::query::query(conn, sql, params).await, + } + } +} + /// Register the `database` submodule (and its per-backend children) on the /// top-level `synapse_rust` module. pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { @@ -26,3 +80,85 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> Ok(()) } + +#[cfg(test)] +mod tests { + //! These exercise the end-to-end query helper against a live Postgres, so + //! they only run when `SYNAPSE_TEST_POSTGRES_DSN` is set (e.g. to + //! `host=postgres user=postgres password=postgres dbname=postgres`); + //! otherwise they no-op. The placeholder rewriting and value mapping have + //! their own server-free unit tests. + + use super::*; + use crate::database::runtime::runtime; + use crate::database::value::DbRowExt; + + fn test_dsn() -> Option { + std::env::var("SYNAPSE_TEST_POSTGRES_DSN").ok() + } + + fn test_pool(dsn: &str) -> DbPool { + DbPool::Postgres(postgres::pool::create_pool(dsn, 2).unwrap()) + } + + #[test] + fn query_binds_placeholders_and_decodes_scalar_rows() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + let pool = test_pool(&dsn); + + runtime().block_on(async { + let mut conn = pool.get().await.unwrap(); + let rows = conn + .query( + "SELECT ?::int AS a, ?::text AS b, ?::bool AS c", + &[ + DbValue::Int(7), + DbValue::Text("hi".into()), + DbValue::Bool(true), + ], + ) + .await + .unwrap(); + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].try_get::(0).unwrap(), 7); + assert_eq!(rows[0].try_get::(1).unwrap(), "hi"); + assert!(rows[0].try_get::(2).unwrap()); + }); + } + + #[test] + fn query_decodes_nulls_bytes_and_multiple_rows() { + let Some(dsn) = test_dsn() else { + eprintln!("skipping: set SYNAPSE_TEST_POSTGRES_DSN to run"); + return; + }; + let pool = test_pool(&dsn); + + runtime().block_on(async { + let mut conn = pool.get().await.unwrap(); + + // A NULL cell reads back as `None`; bytea round-trips byte-for-byte. + let rows = conn + .query( + "SELECT NULL::text AS a, ?::bytea AS b", + &[DbValue::Bytes(vec![0, 255])], + ) + .await + .unwrap(); + assert_eq!(rows[0].try_get::>(0).unwrap(), None); + assert_eq!(rows[0].try_get::>(1).unwrap(), vec![0u8, 255]); + + // Several rows come back in order. + let rows = conn + .query("SELECT g FROM generate_series(1, 3) AS g ORDER BY g", &[]) + .await + .unwrap(); + let got: Vec = rows.iter().map(|r| r.try_get::(0).unwrap()).collect(); + assert_eq!(got, vec![1, 2, 3]); + }); + } +} diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index 5a210b0b60..fa115d99c5 100644 --- a/rust/src/database/postgres/mod.rs +++ b/rust/src/database/postgres/mod.rs @@ -18,6 +18,7 @@ mod errors; mod helpers; mod libpq; pub mod pool; +pub(crate) mod query; mod value; /// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the diff --git a/rust/src/database/postgres/pool.rs b/rust/src/database/postgres/pool.rs index 08da9064fa..f29d4802d1 100644 --- a/rust/src/database/postgres/pool.rs +++ b/rust/src/database/postgres/pool.rs @@ -13,7 +13,7 @@ //! *same* pool, so both share a single set of connections rather than running //! two pools that could exhaust the server's connection limit between them. -use deadpool::managed::{Manager, Metrics, Pool, RecycleError, RecycleResult}; +use deadpool::managed::{Manager, Metrics, Object, Pool, RecycleError, RecycleResult}; use log::warn; use tokio_postgres::{Client, Config, NoTls}; @@ -71,6 +71,10 @@ impl Manager for ConnectionManager { /// A pool of [`tokio_postgres`] connections. pub type ConnectionPool = Pool; +/// A connection checked out of a [`ConnectionPool`]. Dereferences to the +/// underlying [`tokio_postgres::Client`]; returned to the pool when dropped. +pub type PooledConnection = Object; + /// Build a [`ConnectionPool`] from a libpq-style DSN, capped at `max_size` /// connections. pub fn create_pool(dsn: &str, max_size: usize) -> Result { diff --git a/rust/src/database/postgres/query.rs b/rust/src/database/postgres/query.rs new file mode 100644 index 0000000000..818c6dff67 --- /dev/null +++ b/rust/src/database/postgres/query.rs @@ -0,0 +1,82 @@ +//! The Rust-native `query` helper for the Postgres backend. +//! +//! This is the Postgres arm of [`crate::database::DbConn::query`]: it binds +//! backend-agnostic [`DbValue`] parameters and decodes rows back into +//! [`DbRow`]s, so simple queries can be written once regardless of backend. + +use tokio_postgres::types::ToSql; +use tokio_postgres::Client; + +use crate::database::postgres::value::{pg_row_to_db_row, PgValue}; +use crate::database::value::{DbRow, DbValue}; + +/// Run `sql` (with `?` placeholders bound to `params`) and return all rows. +/// +/// Parameters are bound positionally, in order, to the `?` placeholders — which +/// are first rewritten to Postgres' `$1, $2, ...` (see [`convert_placeholders`]). +pub(crate) async fn query( + client: &Client, + sql: &str, + params: &[DbValue], +) -> Result, anyhow::Error> { + let sql = convert_placeholders(sql); + + // Bind each agnostic `DbValue` as a Postgres parameter. `PgValue` is the + // `ToSql` type the rest of the backend already uses. + let pg_params: Vec = params.iter().map(PgValue::from_db_value).collect(); + let params_dyn: Vec<&(dyn ToSql + Sync)> = + pg_params.iter().map(|v| v as &(dyn ToSql + Sync)).collect(); + + // Passing the SQL as a `&str` lets `tokio_postgres` prepare it (inferring + // the parameter types our `ToSql` impl needs) and run it in one call. + let rows = client.query(sql.as_str(), ¶ms_dyn).await?; + + rows.iter().map(pg_row_to_db_row).collect() +} + +/// Rewrite `?` placeholders to Postgres' positional `$1, $2, ...`. +/// +/// A naive left-to-right substitution, matching Synapse's existing +/// `PostgresEngine.convert_param_style` (`?` -> `%s`): it does not skip a `?` +/// inside a string literal, so callers should parameterise rather than embed a +/// literal `?` in the SQL. +fn convert_placeholders(sql: &str) -> String { + let mut out = String::with_capacity(sql.len() + 8); + let mut n = 0u32; + for ch in sql.chars() { + if ch == '?' { + n += 1; + out.push('$'); + out.push_str(&n.to_string()); + } else { + out.push(ch); + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn convert_placeholders_numbers_each_question_mark() { + assert_eq!( + convert_placeholders("SELECT * FROM t WHERE a = ? AND b = ?"), + "SELECT * FROM t WHERE a = $1 AND b = $2" + ); + } + + #[test] + fn convert_placeholders_leaves_sql_without_params_untouched() { + assert_eq!(convert_placeholders("SELECT 1"), "SELECT 1"); + } + + #[test] + fn convert_placeholders_counts_past_nine() { + // The index is decimal, not a single digit. + let sql = "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; + let converted = convert_placeholders(sql); + assert!(converted.ends_with("$9, $10, $11)"), "{converted}"); + } +} diff --git a/rust/src/database/postgres/value.rs b/rust/src/database/postgres/value.rs index 5bb2c74eb4..b725107a7e 100644 --- a/rust/src/database/postgres/value.rs +++ b/rust/src/database/postgres/value.rs @@ -28,13 +28,16 @@ use std::error::Error; use bytes::BytesMut; use postgres_protocol::types::{ - bool_to_sql, bytea_to_sql, float4_to_sql, float8_to_sql, int2_to_sql, int4_to_sql, int8_to_sql, - text_to_sql, + bool_from_sql, bool_to_sql, bytea_to_sql, float4_from_sql, float4_to_sql, float8_from_sql, + float8_to_sql, int2_from_sql, int2_to_sql, int4_from_sql, int4_to_sql, int8_from_sql, + int8_to_sql, text_to_sql, }; use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::types::{PyBool, PyBytes, PyFloat, PyInt, PyString, PyTuple}; use pyo3::{prelude::*, BoundObject}; -use tokio_postgres::types::{to_sql_checked, IsNull, ToSql, Type, WrongType}; +use tokio_postgres::types::{to_sql_checked, FromSql, IsNull, ToSql, Type, WrongType}; + +use crate::database::value::{DbRow, DbValue}; /// Owned representation of a Python value that we can hand to [`tokio_postgres`] /// as a [`ToSql`] parameter. @@ -86,6 +89,42 @@ impl PgValue { obj.get_type().name()?, ))) } + + /// Bind a backend-agnostic [`DbValue`] (from the Rust-native query helpers) + /// as a Postgres parameter. This is the non-Python counterpart of + /// [`PgValue::from_py`]. + pub(crate) fn from_db_value(value: &DbValue) -> Self { + match value { + DbValue::Null => PgValue::Null, + DbValue::Bool(b) => PgValue::Bool(*b), + DbValue::Int(i) => PgValue::Int(*i), + DbValue::Float(f) => PgValue::Float(*f), + DbValue::Text(s) => PgValue::Text(s.as_str().into()), + DbValue::Bytes(b) => PgValue::Bytea(b.as_slice().into()), + } + } +} + +/// The Postgres column types the value mapping supports, in both directions. +/// +/// Shared by every `accepts` implementation here ([`PgValue`]'s `ToSql`, +/// [`PythonPgFromSql`] and [`DbValueFromSql`]) and kept in sync with their +/// `match` arms, so the supported set can't silently drift between them. +pub(crate) fn accepts_column_type(ty: &Type) -> bool { + matches!( + *ty, + Type::BOOL + | Type::INT2 + | Type::INT4 + | Type::INT8 + | Type::FLOAT4 + | Type::FLOAT8 + | Type::TEXT + | Type::VARCHAR + | Type::NAME + | Type::BPCHAR + | Type::BYTEA + ) } // Lets PyO3 extract a `PgValue` directly from a Python argument, e.g. when a @@ -168,20 +207,7 @@ impl ToSql for PgValue { } fn accepts(ty: &Type) -> bool { - matches!( - *ty, - Type::BOOL - | Type::INT2 - | Type::INT4 - | Type::INT8 - | Type::FLOAT4 - | Type::FLOAT8 - | Type::TEXT - | Type::VARCHAR - | Type::NAME - | Type::BPCHAR - | Type::BYTEA - ) + accepts_column_type(ty) } to_sql_checked!(); @@ -229,20 +255,7 @@ impl<'a> tokio_postgres::types::FromSql<'a> for PythonPgFromSql { } fn accepts(ty: &Type) -> bool { - matches!( - *ty, - Type::BOOL - | Type::INT2 - | Type::INT4 - | Type::INT8 - | Type::FLOAT4 - | Type::FLOAT8 - | Type::TEXT - | Type::VARCHAR - | Type::NAME - | Type::BPCHAR - | Type::BYTEA - ) + accepts_column_type(ty) } } @@ -294,6 +307,63 @@ impl PythonPgFromSql { } } +/// Convert a Postgres row into a backend-agnostic [`DbRow`] for the Rust-native +/// query helpers, one [`DbValue`] per column. +/// +/// The non-Python counterpart of [`pg_row_to_py`]; decodes each column via +/// [`DbValueFromSql`]. Errors (with the column index and type) if a column can't +/// be decoded. +pub(crate) fn pg_row_to_db_row(row: &tokio_postgres::Row) -> Result { + let mut out = Vec::with_capacity(row.len()); + for idx in 0..row.len() { + let value: DbValueFromSql = row.try_get(idx).map_err(|e| { + anyhow::anyhow!( + "failed to decode column {idx} (type {}): {e}", + row.columns()[idx].type_() + ) + })?; + out.push(value.0); + } + Ok(out) +} + +/// A column value decoded into a backend-agnostic [`DbValue`]. +/// +/// The non-Python counterpart of [`PythonPgFromSql`]: the same supported types +/// (see [`accepts_column_type`]), but with no GIL and no Python objects — just +/// the plain Rust value. +pub(crate) struct DbValueFromSql(pub DbValue); + +impl<'a> FromSql<'a> for DbValueFromSql { + fn from_sql(ty: &Type, raw: &'a [u8]) -> Result> { + let value = match *ty { + Type::BOOL => DbValue::Bool(bool_from_sql(raw)?), + Type::INT2 => DbValue::Int(int2_from_sql(raw)?.into()), + Type::INT4 => DbValue::Int(int4_from_sql(raw)?.into()), + Type::INT8 => DbValue::Int(int8_from_sql(raw)?), + Type::FLOAT4 => DbValue::Float(float4_from_sql(raw)?.into()), + Type::FLOAT8 => DbValue::Float(float8_from_sql(raw)?), + Type::TEXT | Type::VARCHAR | Type::NAME | Type::BPCHAR => { + DbValue::Text(std::str::from_utf8(raw)?.to_owned()) + } + Type::BYTEA => DbValue::Bytes(raw.to_vec()), + _ => { + // Unreachable unless `accepts` drifts out of sync with this match. + return Err(format!("unsupported column type for postgres: {ty}").into()); + } + }; + Ok(DbValueFromSql(value)) + } + + fn from_sql_null(_ty: &Type) -> Result> { + Ok(DbValueFromSql(DbValue::Null)) + } + + fn accepts(ty: &Type) -> bool { + accepts_column_type(ty) + } +} + #[cfg(test)] mod tests { //! These tests exercise the value mapping in isolation — no live Postgres diff --git a/rust/src/database/value.rs b/rust/src/database/value.rs new file mode 100644 index 0000000000..d37c633e35 --- /dev/null +++ b/rust/src/database/value.rs @@ -0,0 +1,208 @@ +//! Backend-agnostic value and row types for the Rust-native query helpers. +//! +//! [`DbValue`] is the common currency for both query *parameters* and result +//! *cells*, so simple queries can be written once and run against any backend +//! (currently Postgres; SQLite to follow). Each backend maps its driver's own +//! value types to and from [`DbValue`] at its edge — see +//! `postgres::value` for the Postgres mapping. +//! +//! The set is deliberately small (the common scalar types); richer types can be +//! added as the query helpers grow to need them. + +/// A single backend-agnostic value, used both as a query parameter and as a +/// cell in a returned [`DbRow`]. +#[derive(Debug, Clone, PartialEq)] +pub enum DbValue { + /// SQL `NULL`. + Null, + Bool(bool), + Int(i64), + Float(f64), + Text(String), + Bytes(Vec), +} + +// Ergonomic conversions so callers can write `&[1i64.into(), "x".into()]` (or +// pass typed values straight through) rather than spelling out `DbValue`. +impl From for DbValue { + fn from(v: bool) -> Self { + DbValue::Bool(v) + } +} +impl From for DbValue { + fn from(v: i64) -> Self { + DbValue::Int(v) + } +} +impl From for DbValue { + fn from(v: i32) -> Self { + DbValue::Int(v as i64) + } +} +impl From for DbValue { + fn from(v: f64) -> Self { + DbValue::Float(v) + } +} +impl From for DbValue { + fn from(v: String) -> Self { + DbValue::Text(v) + } +} +impl From<&str> for DbValue { + fn from(v: &str) -> Self { + DbValue::Text(v.to_owned()) + } +} +impl From> for DbValue { + fn from(v: Vec) -> Self { + DbValue::Bytes(v) + } +} +impl From<&[u8]> for DbValue { + fn from(v: &[u8]) -> Self { + DbValue::Bytes(v.to_vec()) + } +} +impl> From> for DbValue { + fn from(v: Option) -> Self { + match v { + Some(v) => v.into(), + None => DbValue::Null, + } + } +} + +/// A row of data returned by a query: one [`DbValue`] per column, read out by +/// numeric index with [`DbRowExt::try_get`]. +pub type DbRow = Vec; + +/// Reads a typed value out of a [`DbValue`], analogous to `tokio_postgres`'s +/// `FromSql`. `Option` reads a nullable column (SQL `NULL` -> `None`). +pub trait FromDbValue: Sized { + fn from_value(value: DbValue) -> Result; +} + +impl FromDbValue for bool { + fn from_value(value: DbValue) -> Result { + match value { + DbValue::Bool(b) => Ok(b), + // SQLite has no native boolean type and stores them as integers. + DbValue::Int(i) => Ok(i != 0), + other => anyhow::bail!("cannot read {other:?} as bool"), + } + } +} + +impl FromDbValue for i64 { + fn from_value(value: DbValue) -> Result { + match value { + DbValue::Int(i) => Ok(i), + DbValue::Bool(b) => Ok(b as i64), + other => anyhow::bail!("cannot read {other:?} as i64"), + } + } +} + +impl FromDbValue for f64 { + fn from_value(value: DbValue) -> Result { + match value { + DbValue::Float(f) => Ok(f), + DbValue::Int(i) => Ok(i as f64), + other => anyhow::bail!("cannot read {other:?} as f64"), + } + } +} + +impl FromDbValue for String { + fn from_value(value: DbValue) -> Result { + match value { + DbValue::Text(s) => Ok(s), + other => anyhow::bail!("cannot read {other:?} as String"), + } + } +} + +impl FromDbValue for Vec { + fn from_value(value: DbValue) -> Result { + match value { + DbValue::Bytes(b) => Ok(b), + other => anyhow::bail!("cannot read {other:?} as bytes"), + } + } +} + +impl FromDbValue for Option { + fn from_value(value: DbValue) -> Result { + match value { + DbValue::Null => Ok(None), + other => Ok(Some(T::from_value(other)?)), + } + } +} + +/// Extension methods for reading typed values out of a [`DbRow`]. +pub trait DbRowExt { + /// Read the value at `index`, converting it into `T` via [`FromDbValue`]. + /// Errors if the index is out of bounds or the value can't become `T`. + fn try_get(&self, index: usize) -> Result; +} + +impl DbRowExt for DbRow { + fn try_get(&self, index: usize) -> Result { + let value = self.get(index).cloned().ok_or_else(|| { + anyhow::anyhow!( + "tried to read column {index} but the row only has {} column(s)", + self.len() + ) + })?; + T::from_value(value) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn from_impls_classify_scalars() { + assert_eq!(DbValue::from(true), DbValue::Bool(true)); + assert_eq!(DbValue::from(7i64), DbValue::Int(7)); + assert_eq!(DbValue::from(7i32), DbValue::Int(7)); + assert_eq!(DbValue::from(1.5f64), DbValue::Float(1.5)); + assert_eq!(DbValue::from("x"), DbValue::Text("x".to_owned())); + assert_eq!(DbValue::from(vec![1u8, 2]), DbValue::Bytes(vec![1, 2])); + // `Option` folds `None` to NULL and `Some` to the inner conversion. + assert_eq!(DbValue::from(None::), DbValue::Null); + assert_eq!(DbValue::from(Some(3i64)), DbValue::Int(3)); + } + + #[test] + fn try_get_reads_typed_values() { + let row: DbRow = vec![ + DbValue::Int(42), + DbValue::Text("hi".to_owned()), + DbValue::Null, + ]; + assert_eq!(row.try_get::(0).unwrap(), 42); + assert_eq!(row.try_get::(1).unwrap(), "hi"); + // A NULL read as `Option` is `None`... + assert_eq!(row.try_get::>(2).unwrap(), None); + // ...but read as a non-optional type it's an error, not a silent default. + assert!(row.try_get::(2).is_err()); + } + + #[test] + fn try_get_out_of_bounds_errors() { + let row: DbRow = vec![DbValue::Int(1)]; + let err = row.try_get::(5).unwrap_err(); + assert!(err.to_string().contains("only has 1 column"), "{err}"); + } + + #[test] + fn bool_reads_from_sqlite_style_integer() { + // SQLite returns booleans as integers; `FromDbValue for bool` bridges it. + assert!(bool::from_value(DbValue::Int(1)).unwrap()); + assert!(!bool::from_value(DbValue::Int(0)).unwrap()); + } +}