From 65ba851f2cfcd86f901d81e32cff17b37eeb2d48 Mon Sep 17 00:00:00 2001 From: Eric Eastwood Date: Wed, 17 Jun 2026 18:25:07 -0500 Subject: [PATCH] Concrete `Row` type --- rust/src/storage/db/mod.rs | 60 ++++++--------- rust/src/storage/db/python_db_pool.rs | 45 ++--------- rust/src/storage/db/rust_db_pool.rs | 107 +++++++------------------- rust/src/storage/store.rs | 5 +- 4 files changed, 65 insertions(+), 152 deletions(-) diff --git a/rust/src/storage/db/mod.rs b/rust/src/storage/db/mod.rs index aa18c93a79..9729bc3953 100644 --- a/rust/src/storage/db/mod.rs +++ b/rust/src/storage/db/mod.rs @@ -52,8 +52,7 @@ pub trait DatabasePool: Send + Sync { /// interact with the database #[async_trait::async_trait] pub trait Transaction: Send { - async fn query(&mut self, sql: &str, args: &[&str]) - -> Result>, anyhow::Error>; + async fn query(&mut self, sql: &str, args: &[&str]) -> Result, anyhow::Error>; } /// A single backend-agnostic value within a [`Row`]. @@ -68,39 +67,37 @@ pub enum Value { Int(i64), Float(f64), Text(String), - Bytes(Vec), } /// A row of data returned from the database by a query. /// -/// Modelled after [`tokio_postgres::Row`]: values are pulled out by their numeric -/// index with [`get`](Self::get) / [`try_get`](Self::try_get). Each database pool -/// implements this trait for its own native row type — the Python pool inspects -/// the Python type of each cell, while the `tokio-postgres` pool reuses that -/// crate's own `FromSql` machinery — converting cells into the engine-agnostic -/// [`Value`] returned by [`Row::value`]. -pub trait Row: std::fmt::Debug + Send { - /// Returns the number of values in the row. - fn len(&self) -> usize; +/// Each pool converts the cells its database driver hands back into the +/// engine-agnostic [`Value`] representation, so a row is simply a list of them. +/// Values are pulled out by their numeric index with [`RowExt::try_get`]. +pub type Row = Vec; - /// Returns whether the row contains no values. - fn is_empty(&self) -> bool { - self.len() == 0 - } - - /// Returns the backend-agnostic [`Value`] at `index`, converting from the - /// pool's native representation. - /// - /// Errors if `index` is out of bounds. - fn value(&self, index: usize) -> Result; -} - -impl dyn Row { +/// Extension methods for reading typed values out of a [`Row`]. +/// +/// Modelled after [`tokio_postgres::Row`]'s `try_get`: [`try_get`](Self::try_get) +/// converts the [`Value`] at a given index into the requested type via +/// [`FromValue`] (our analogue of `tokio-postgres`'s `FromSql`). +pub trait RowExt { /// Deserializes a value from the row, specified by its numeric index, /// returning an error if the index is out of bounds or the value cannot be /// converted into `T`. - pub fn try_get(&self, index: usize) -> Result { - T::from_value(self.value(index)?) + fn try_get(&self, index: usize) -> Result; +} + +impl RowExt for Row { + fn try_get(&self, index: usize) -> Result { + let value = self.get(index).cloned().ok_or_else(|| { + anyhow::anyhow!( + "tried to get column {index} but the row only has {} column(s)", + self.len() + ) + })?; + + T::from_value(value) } } @@ -150,15 +147,6 @@ impl FromValue for String { } } -impl FromValue for Vec { - fn from_value(value: Value) -> Result { - match value { - Value::Bytes(b) => Ok(b), - other => anyhow::bail!("cannot read {other:?} as bytes"), - } - } -} - impl FromValue for Option { fn from_value(value: Value) -> Result { match value { diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index f97694e248..07fef298b2 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/rust/src/storage/db/python_db_pool.rs @@ -17,7 +17,6 @@ //! - Database pool [`PythonDatabasePoolWrapper`] which allows you to start a... //! - transaction [`LoggingTransactionWrapper`] and query the database -use std::any::Any; use std::sync::{Arc, Mutex}; use futures::future::BoxFuture; @@ -25,7 +24,7 @@ use pyo3::{ exceptions::{PyRuntimeError, PyTypeError}, intern, prelude::*, - types::{PyBool, PyBytes, PyCFunction, PyFloat, PyInt, PyList, PyString}, + types::{PyBool, PyCFunction, PyFloat, PyInt, PyList, PyString}, }; use crate::http_client::await_deferred; @@ -236,12 +235,8 @@ impl LoggingTransactionWrapper { #[async_trait::async_trait] impl Transaction for LoggingTransactionWrapper { - async fn query( - &mut self, - sql: &str, - args: &[&str], - ) -> Result>, anyhow::Error> { - Python::attach(|py| -> PyResult>> { + async fn query(&mut self, sql: &str, args: &[&str]) -> Result, anyhow::Error> { + Python::attach(|py| -> PyResult> { // Convert the Rust `&[&str]` of SQL parameters into a Python sequence // so it can be passed through to the Python-side `execute`. Note that // `LoggingTransaction.execute` converts `?` placeholders into the @@ -257,14 +252,14 @@ impl Transaction for LoggingTransactionWrapper { .bind(py) .call_method0(intern!(py, "fetchall"))?; - let mut rows: Vec> = Vec::new(); + let mut rows: Vec = Vec::new(); for row_py in rows_py.try_iter()? { let row_py = row_py?; - let mut values: Vec = Vec::new(); + let mut row: Row = Vec::new(); for cell in row_py.try_iter()? { - values.push(py_cell_to_value(&cell?)?); + row.push(py_cell_to_value(&cell?)?); } - rows.push(Box::new(PythonRow { values })); + rows.push(row); } Ok(rows) @@ -281,7 +276,7 @@ fn py_cell_to_value(cell: &Bound<'_, PyAny>) -> PyResult { return Ok(Value::Null); } - // A `bool` *is* an `int` in Python, so ensure we try `bool` first. + // A `bool` *is* an `int` in SQLite, so ensure we try `bool` first. if let Ok(b) = cell.cast::() { Ok(Value::Bool(b.extract()?)) } else if let Ok(i) = cell.cast::() { @@ -290,8 +285,6 @@ fn py_cell_to_value(cell: &Bound<'_, PyAny>) -> PyResult { Ok(Value::Float(f.extract()?)) } else if let Ok(s) = cell.cast::() { Ok(Value::Text(s.to_string())) - } else if let Ok(bytes) = cell.cast::() { - Ok(Value::Bytes(bytes.as_bytes().to_vec())) } else { Err(PyTypeError::new_err(format!( "unsupported column type {} returned from the database", @@ -299,25 +292,3 @@ fn py_cell_to_value(cell: &Bound<'_, PyAny>) -> PyResult { ))) } } - -/// A [`Row`] backed by values pulled out of a Python `LoggingTransaction`. -#[derive(Debug)] -struct PythonRow { - /// Each cell, already converted from its Python type into a [`Value`]. - values: Vec, -} - -impl Row for PythonRow { - fn len(&self) -> usize { - self.values.len() - } - - fn value(&self, index: usize) -> Result { - self.values.get(index).cloned().ok_or_else(|| { - anyhow::anyhow!( - "tried to get column {index} but the row only has {} column(s)", - self.values.len() - ) - }) - } -} diff --git a/rust/src/storage/db/rust_db_pool.rs b/rust/src/storage/db/rust_db_pool.rs index 252bfa5215..be280c1f9b 100644 --- a/rust/src/storage/db/rust_db_pool.rs +++ b/rust/src/storage/db/rust_db_pool.rs @@ -17,11 +17,7 @@ // interfaces are compatible with `tokio-postgres`. use anyhow::Context; -use bb8_postgres::tokio_postgres::{ - self, - types::{ToSql, Type}, - IsolationLevel, -}; +use bb8_postgres::tokio_postgres::{self, types::ToSql, IsolationLevel}; use bb8_postgres::PostgresConnectionManager; use futures::future::BoxFuture; use postgres_native_tls::MakeTlsConnector; @@ -100,11 +96,7 @@ struct TokioPostgresTransaction<'a> { #[async_trait::async_trait] impl Transaction for TokioPostgresTransaction<'_> { - async fn query( - &mut self, - sql: &str, - args: &[&str], - ) -> Result>, anyhow::Error> { + async fn query(&mut self, sql: &str, args: &[&str]) -> Result, anyhow::Error> { // Synapse SQL uses `?` placeholders; `tokio-postgres` uses `$1`, `$2`, ... let sql = convert_param_style(sql); @@ -116,79 +108,38 @@ impl Transaction for TokioPostgresTransaction<'_> { .await .context("Failed to run query")?; - let out = rows - .into_iter() - .map(|row| Box::new(TokioPostgresRow { row }) as Box) - .collect(); - - Ok(out) + rows.iter().map(tokio_row_to_row).collect() } } -/// A [`Row`] backed by a native [`tokio_postgres::Row`]. -#[derive(Debug)] -struct TokioPostgresRow { - row: tokio_postgres::Row, +/// Convert a native [`tokio_postgres::Row`] into our engine-agnostic [`Row`]. +fn tokio_row_to_row(row: &tokio_postgres::Row) -> Result { + (0..row.len()) + .map(|index| tokio_cell_to_value(row, index)) + .collect() } -impl Row for TokioPostgresRow { - fn len(&self) -> usize { - self.row.len() - } - - fn value(&self, index: usize) -> Result { - let column = self.row.columns().get(index).ok_or_else(|| { - anyhow::anyhow!( - "tried to get column {index} but the row only has {} column(s)", - self.row.len() - ) - })?; - - // Dispatch on the column's Postgres type, leaning on `tokio-postgres`'s - // own `FromSql` impls to read the cell. Everything is read as - // `Option<_>` so a SQL `NULL` becomes `Value::Null`. - let ty = column.type_(); - let value = if *ty == Type::BOOL { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, Value::Bool) - } else if *ty == Type::INT2 { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, |v| Value::Int(v.into())) - } else if *ty == Type::INT4 { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, |v| Value::Int(v.into())) - } else if *ty == Type::INT8 { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, Value::Int) - } else if *ty == Type::FLOAT4 { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, |v| Value::Float(v.into())) - } else if *ty == Type::FLOAT8 { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, Value::Float) - } else if *ty == Type::TEXT - || *ty == Type::VARCHAR - || *ty == Type::BPCHAR - || *ty == Type::NAME - { - self.row - .try_get::<_, Option>(index)? - .map_or(Value::Null, Value::Text) - } else if *ty == Type::BYTEA { - self.row - .try_get::<_, Option>>(index)? - .map_or(Value::Null, Value::Bytes) - } else { - anyhow::bail!("unsupported Postgres column type `{ty}` at column {index}"); - }; - - Ok(value) +/// Convert a single cell of a [`tokio_postgres::Row`] into a [`Value`]. +/// +/// Dispatch on the column's Postgres type, leaning on `tokio-postgres`'s own +/// `FromSql` impls to read the cell. Everything is read as `Option<_>` so a SQL +/// `NULL` becomes `Value::Null`. +fn tokio_cell_to_value(row: &tokio_postgres::Row, index: usize) -> Result { + if let Ok(value) = row.try_get::<_, Option>(index) { + Ok(value.map_or(Value::Null, Value::Bool)) + } else if let Ok(value) = row.try_get::<_, Option>(index) { + Ok(value.map_or(Value::Null, Value::Int)) + } else if let Ok(value) = row.try_get::<_, Option>(index) { + Ok(value.map_or(Value::Null, Value::Float)) + } else if let Ok(value) = row.try_get::<_, Option>(index) { + Ok(value.map_or(Value::Null, Value::Text)) + } else { + let ty = row.columns()[index].type_(); + anyhow::bail!( + "Unsupported `tokio-postgres` type {} encountered when trying to convert it \ + to our generic database `Value` type. You probably just need to implement it in `tokio_cell_to_value(...)`.", + ty + ) } } diff --git a/rust/src/storage/store.rs b/rust/src/storage/store.rs index 2f1927fb73..f57f94e5b0 100644 --- a/rust/src/storage/store.rs +++ b/rust/src/storage/store.rs @@ -16,7 +16,10 @@ use futures::FutureExt; use serde::Serialize; -use crate::{config::SynapseConfig, storage::db::DatabasePool}; +use crate::{ + config::SynapseConfig, + storage::db::{DatabasePool, RowExt}, +}; /// Currently supported per-user features #[derive(Serialize, Debug)]