LLM further Row

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<Vec<dyn Row>, anyhow::Error>;
+    async fn query(&mut self, sql: &str, args: &[&str])
+        -> Result<Vec<Box<dyn Row>>, 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<u8>),
 }
 
 /// 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<T>(&self, index: usize) -> Result<T, anyhow::Error>;
+    /// Errors if `index` is out of bounds.
+    fn value(&self, index: usize) -> Result<Value, anyhow::Error>;
 }
 
-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<T: FromValue>(&self, index: usize) -> Result<T, anyhow::Error> {
+        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<Self, anyhow::Error>;
+}
+
+impl FromValue for bool {
+    fn from_value(value: Value) -> Result<Self, anyhow::Error> {
+        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<Self, anyhow::Error> {
+        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<Self, anyhow::Error> {
+        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<Self, anyhow::Error> {
+        match value {
+            Value::Text(s) => Ok(s),
+            other => anyhow::bail!("cannot read {other:?} as String"),
+        }
+    }
+}
+
+impl FromValue for Vec<u8> {
+    fn from_value(value: Value) -> Result<Self, anyhow::Error> {
+        match value {
+            Value::Bytes(b) => Ok(b),
+            other => anyhow::bail!("cannot read {other:?} as bytes"),
+        }
+    }
+}
+
+impl<T: FromValue> FromValue for Option<T> {
+    fn from_value(value: Value) -> Result<Self, anyhow::Error> {
+        match value {
+            Value::Null => Ok(None),
+            other => Ok(Some(T::from_value(other)?)),
+        }
     }
 }

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<Vec<Row>, anyhow::Error> {
-        Python::attach(|py| -> PyResult<Vec<Row>> {
+    async fn query(
+        &mut self,
+        sql: &str,
+        args: &[&str],
+    ) -> Result<Vec<Box<dyn Row>>, anyhow::Error> {
+        Python::attach(|py| -> PyResult<Vec<Box<dyn Row>>> {
             // 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<Row> = Vec::new();
+            let mut rows: Vec<Box<dyn Row>> = 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<Value> = 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<Value> {
+    // `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::<PyBool>() {
+        Ok(Value::Bool(b.extract()?))
+    } else if let Ok(i) = cell.cast::<PyInt>() {
+        Ok(Value::Int(i.extract()?))
+    } else if let Ok(f) = cell.cast::<PyFloat>() {
+        Ok(Value::Float(f.extract()?))
+    } else if let Ok(s) = cell.cast::<PyString>() {
+        Ok(Value::Text(s.to_string()))
+    } else if let Ok(bytes) = cell.cast::<PyBytes>() {
+        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<Value>,
 }
 
 impl Row for PythonRow {
-    // TODO
+    fn len(&self) -> usize {
+        self.values.len()
+    }
+
+    fn value(&self, index: usize) -> Result<Value, anyhow::Error> {
+        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()
+            )
+        })
+    }
 }

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<Vec<Row>, anyhow::Error> {
+    async fn query(
+        &mut self,
+        sql: &str,
+        args: &[&str],
+    ) -> Result<Vec<Box<dyn Row>>, 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<Row> = 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<dyn Row>)
+            .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<Value, anyhow::Error> {
+        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<bool>>(index)?
+                .map_or(Value::Null, Value::Bool)
+        } else if *ty == Type::INT2 {
+            self.row
+                .try_get::<_, Option<i16>>(index)?
+                .map_or(Value::Null, |v| Value::Int(v.into()))
+        } else if *ty == Type::INT4 {
+            self.row
+                .try_get::<_, Option<i32>>(index)?
+                .map_or(Value::Null, |v| Value::Int(v.into()))
+        } else if *ty == Type::INT8 {
+            self.row
+                .try_get::<_, Option<i64>>(index)?
+                .map_or(Value::Null, Value::Int)
+        } else if *ty == Type::FLOAT4 {
+            self.row
+                .try_get::<_, Option<f32>>(index)?
+                .map_or(Value::Null, |v| Value::Float(v.into()))
+        } else if *ty == Type::FLOAT8 {
+            self.row
+                .try_get::<_, Option<f64>>(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<String>>(index)?
+                .map_or(Value::Null, Value::Text)
+        } else if *ty == Type::BYTEA {
+            self.row
+                .try_get::<_, Option<Vec<u8>>>(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());

rust/src/storage/store.rs

@@ -97,7 +97,7 @@ impl<P: DatabasePool> Store<P> {
                         // 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")
                         }