From d3af0c4e181126d922c7f222b6d02bc4a536821e Mon Sep 17 00:00:00 2001 From: Eric Eastwood Date: Wed, 17 Jun 2026 17:48:35 -0500 Subject: [PATCH] LLM further `Row` --- rust/src/storage/db/mod.rs | 113 ++++++++++++++++++++++++-- rust/src/storage/db/python_db_pool.rs | 76 +++++++++++++---- rust/src/storage/db/rust_db_pool.rs | 97 ++++++++++++++++++---- rust/src/storage/store.rs | 2 +- 4 files changed, 249 insertions(+), 39 deletions(-) diff --git a/rust/src/storage/db/mod.rs b/rust/src/storage/db/mod.rs index 79001f5b8c..aa18c93a79 100644 --- a/rust/src/storage/db/mod.rs +++ b/rust/src/storage/db/mod.rs @@ -14,7 +14,6 @@ */ use std::future::Future; -use std::str::FromStr; use futures::future::BoxFuture; @@ -53,22 +52,118 @@ 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`]. +/// +/// Each pool maps the values its database driver hands back into this common +/// set, so callers can work with one representation regardless of engine. +#[derive(Debug, Clone, PartialEq)] +pub enum Value { + /// A SQL `NULL`. + Null, + Bool(bool), + Int(i64), + Float(f64), + Text(String), + Bytes(Vec), } /// A row of data returned from the database by a query. -pub trait Row { +/// +/// 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; - /// Deserializes a value from the row. + /// 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. /// - /// The value can be specified by its numeric index in the row. - fn try_get(&self, index: usize) -> Result; + /// Errors if `index` is out of bounds. + fn value(&self, index: usize) -> Result; } -impl std::fmt::Debug for dyn Row { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - // TODO +impl dyn Row { + /// 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)?) + } +} + +/// Converts a backend-agnostic [`Value`] into a concrete Rust type, analogous to +/// `tokio-postgres`'s `FromSql`. +pub trait FromValue: Sized { + fn from_value(value: Value) -> Result; +} + +impl FromValue for bool { + fn from_value(value: Value) -> Result { + match value { + Value::Bool(b) => Ok(b), + // SQLite has no native boolean type and stores them as integers. + Value::Int(i) => Ok(i != 0), + other => anyhow::bail!("cannot read {other:?} as bool"), + } + } +} + +impl FromValue for i64 { + fn from_value(value: Value) -> Result { + match value { + Value::Int(i) => Ok(i), + Value::Bool(b) => Ok(b as i64), + other => anyhow::bail!("cannot read {other:?} as i64"), + } + } +} + +impl FromValue for f64 { + fn from_value(value: Value) -> Result { + match value { + Value::Float(f) => Ok(f), + Value::Int(i) => Ok(i as f64), + other => anyhow::bail!("cannot read {other:?} as f64"), + } + } +} + +impl FromValue for String { + fn from_value(value: Value) -> Result { + match value { + Value::Text(s) => Ok(s), + other => anyhow::bail!("cannot read {other:?} as 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 { + Value::Null => Ok(None), + other => Ok(Some(T::from_value(other)?)), + } } } diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index 8005247c26..f97694e248 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/rust/src/storage/db/python_db_pool.rs @@ -21,10 +21,15 @@ use std::any::Any; use std::sync::{Arc, Mutex}; use futures::future::BoxFuture; -use pyo3::{exceptions::PyRuntimeError, intern, prelude::*, types::PyCFunction, types::PyList}; +use pyo3::{ + exceptions::{PyRuntimeError, PyTypeError}, + intern, + prelude::*, + types::{PyBool, PyBytes, PyCFunction, PyFloat, PyInt, PyList, PyString}, +}; use crate::http_client::await_deferred; -use crate::storage::db::{DatabasePool, Row, Transaction}; +use crate::storage::db::{DatabasePool, Row, Transaction, Value}; /// The database engines we support in the Python side of Synapse #[derive(Copy, Clone, Debug)] @@ -231,8 +236,12 @@ 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 @@ -241,23 +250,21 @@ impl Transaction for LoggingTransactionWrapper { let args = PyList::new(py, args)?; self.execute(py, sql, args.as_any())?; - // Pull the rows back out. Each cell is converted to its textual - // representation so we have a single engine-agnostic `Row` type; - // callers parse the strings into richer types as needed. + // Pull the rows back out, converting each cell from its Python type + // into the engine-agnostic `Value` representation as we go. let rows_py = self .logging_transaction_py .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 row: Row = Vec::new(); + let mut values: Vec = Vec::new(); for cell in row_py.try_iter()? { - let cell = cell?; - row.push(cell.str()?.to_string_lossy().into_owned()); + values.push(py_cell_to_value(&cell?)?); } - rows.push(row); + rows.push(Box::new(PythonRow { values })); } Ok(rows) @@ -266,10 +273,51 @@ impl Transaction for LoggingTransactionWrapper { } } +/// Convert a single cell from a Python row into a backend-agnostic [`Value`] by +/// inspecting its Python type (the pyo3 equivalent of `isinstance` checks). +fn py_cell_to_value(cell: &Bound<'_, PyAny>) -> PyResult { + // `None` maps to SQL `NULL`. + if cell.is_none() { + return Ok(Value::Null); + } + + // A `bool` *is* an `int` in Python, so ensure we try `bool` first. + if let Ok(b) = cell.cast::() { + Ok(Value::Bool(b.extract()?)) + } else if let Ok(i) = cell.cast::() { + Ok(Value::Int(i.extract()?)) + } else if let Ok(f) = cell.cast::() { + 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", + cell.get_type().name()? + ))) + } +} + +/// A [`Row`] backed by values pulled out of a Python `LoggingTransaction`. +#[derive(Debug)] struct PythonRow { - // TODO + /// Each cell, already converted from its Python type into a [`Value`]. + values: Vec, } impl Row for PythonRow { - // TODO + 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 284bd18924..252bfa5215 100644 --- a/rust/src/storage/db/rust_db_pool.rs +++ b/rust/src/storage/db/rust_db_pool.rs @@ -17,12 +17,16 @@ // interfaces are compatible with `tokio-postgres`. use anyhow::Context; -use bb8_postgres::tokio_postgres::{self, types::ToSql, IsolationLevel}; +use bb8_postgres::tokio_postgres::{ + self, + types::{ToSql, Type}, + IsolationLevel, +}; use bb8_postgres::PostgresConnectionManager; use futures::future::BoxFuture; use postgres_native_tls::MakeTlsConnector; -use crate::storage::db::{DatabasePool, Row, Transaction}; +use crate::storage::db::{DatabasePool, Row, Transaction, Value}; /// Native Rust database access backed by `tokio-postgres` (for use in synapse-rust-apps) pub struct RustDatabasePool { @@ -96,7 +100,11 @@ 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); @@ -108,23 +116,82 @@ impl Transaction for TokioPostgresTransaction<'_> { .await .context("Failed to run query")?; - let mut out: Vec = Vec::with_capacity(rows.len()); - for row in rows { - let mut cells: Row = Vec::with_capacity(row.len()); - for i in 0..row.len() { - // Best-effort textual extraction to match the engine-agnostic - // `Row` type. A real implementation would map column types - // properly; this only exists to prove the interface fits. - let value: String = row.try_get(i).unwrap_or_default(); - cells.push(value); - } - out.push(cells); - } + let out = rows + .into_iter() + .map(|row| Box::new(TokioPostgresRow { row }) as Box) + .collect(); Ok(out) } } +/// A [`Row`] backed by a native [`tokio_postgres::Row`]. +#[derive(Debug)] +struct TokioPostgresRow { + row: tokio_postgres::Row, +} + +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 `?`-style placeholders into `tokio-postgres`'s `$1`, `$2`, ... style. fn convert_param_style(sql: &str) -> String { let mut out = String::with_capacity(sql.len()); diff --git a/rust/src/storage/store.rs b/rust/src/storage/store.rs index 02191bf519..2f1927fb73 100644 --- a/rust/src/storage/store.rs +++ b/rust/src/storage/store.rs @@ -97,7 +97,7 @@ impl Store

{ // If there is no row, default to disabled [] => false, // Found an entry for the user - [row] => row.try_get(0), + [row] => row.try_get(0)?, _ => { panic!("Programming error") }