From 66737780dd14c5ade81e03a8a0d1a1138eb8979c Mon Sep 17 00:00:00 2001 From: Eric Eastwood Date: Thu, 4 Jun 2026 18:44:04 -0500 Subject: [PATCH] Slow going --- Cargo.lock | 1 + rust/Cargo.toml | 1 + rust/src/storage/db/mod.rs | 17 ++++++-- rust/src/storage/db/python_db_pool.rs | 62 +++++++++++++++++++++------ rust/src/storage/db/rust_db_pool.rs | 25 ++++++----- 5 files changed, 79 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5759240eb9..5faf827e1d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1831,6 +1831,7 @@ name = "synapse" version = "0.1.0" dependencies = [ "anyhow", + "async-trait", "base64", "bb8", "bb8-postgres", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index dfd6d87b61..1e16b8769c 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -22,6 +22,7 @@ crate-type = ["lib", "cdylib"] name = "synapse.synapse_rust" [dependencies] +async-trait = "0.1.89" anyhow = "1.0.63" base64 = "0.22.1" bytes = "1.6.0" diff --git a/rust/src/storage/db/mod.rs b/rust/src/storage/db/mod.rs index 7c6c4fbb9a..14f4a2692e 100644 --- a/rust/src/storage/db/mod.rs +++ b/rust/src/storage/db/mod.rs @@ -16,13 +16,24 @@ pub mod python_db_pool; pub mod rust_db_pool; +#[async_trait::async_trait] pub trait DatabasePool { - async fn get_transaction(&self, description: &str) -> dyn Transaction; + /// TODO + /// + /// Arguments: + /// description of the transaction, for logging and metrics + async fn get_transaction( + &self, + description: &str, + ) -> Result, anyhow::Error>; } /// A [`tokio_postgres::Transaction`] looking thing that we can use on the Rust side to /// interact with the database +#[async_trait::async_trait] pub trait Transaction { - async fn query(&self, sql: &str, args: &[&str]) -> (); - async fn commit(&self) -> (); + async fn query(&self, sql: &str, args: &[&str]) -> Vec; + async fn commit(self) -> Result<(), anyhow::Error>; } + +pub type Row = Vec; diff --git a/rust/src/storage/db/python_db_pool.rs b/rust/src/storage/db/python_db_pool.rs index 979b8069aa..a2636716cf 100644 --- a/rust/src/storage/db/python_db_pool.rs +++ b/rust/src/storage/db/python_db_pool.rs @@ -15,7 +15,7 @@ use pyo3::{intern, prelude::*}; -use crate::storage::db::{DatabasePool, Transaction}; +use crate::storage::db::{DatabasePool, Row, Transaction}; /// The database engines we support in the Python side of Synapse #[derive(Copy, Clone, Debug)] @@ -34,14 +34,40 @@ impl DatabaseEngine { } } -pub struct PythonDatabasePool {} +/// Wrapper for a `DatabasePool` from the Python side of Synapse. +pub struct PythonDatabasePool { + /// The underlying `DatabasePool` + database_pool_py: Bound<'py, PyAny>, +} +#[async_trait::async_trait] impl DatabasePool for PythonDatabasePool { - pub fn get_transaction(&self, description: &str) -> dyn Transaction { - todo!("..."); - // let execute_fn = self.raw.getattr(intern!(self.raw.py(), "runInteraction"))?; - // execute_fn.call1((sql, args))?; - // Ok(()) + async fn get_transaction( + &self, + description: &str, + ) -> Result, anyhow::Error> { + // Synapse has built-in retry functionality and can call this function multiple + // times under certain failure modes. Normally, everything in the transaction + // happens in the callback but since we have a little bit of a different API + // surface, we instead extract the transaction for us to use outside. + // + // Re-using `runInteraction`, means we get all of the logging, metrics, etc for + // free. + let callback_func = + PyCFunction::new_closure(py, None, None, move |args, _| -> PyResult> { + // TODO: Error if already called + + let py = args.py(); + let txn_py = args.get_item(0)?; + txn + }); + + let execute_fn = self + .database_pool_py + .getattr(intern!(self.database_pool_py.py(), "runInteraction"))?; + execute_fn.call1((description, callback_func))?; + + Ok(Box::new(txn)) } } @@ -71,7 +97,7 @@ fn detect_engine(txn_py: &Bound<'_, PyAny>) -> PyResult { /// Use [`execute`](Self::execute) (or other methods) while holding the GIL. pub struct LoggingTransactionWrapper { /// The underlying `LoggingTransaction` - raw: Py, + logging_transaction_py: Py, /// Disambiguate which underlying database engine we're working with pub database_engine: DatabaseEngine, @@ -83,10 +109,10 @@ impl<'a, 'py> FromPyObject<'a, 'py> for LoggingTransactionWrapper { /// Extract from a Python `LoggingTransaction` passed as an argument. /// /// The resulting wrapper has `done_tx = None`; Python owns the transaction lifetime. - fn extract(obj: Borrowed<'a, 'py, PyAny>) -> PyResult { - let database_engine = detect_engine(&obj.to_owned())?; + fn extract(logging_transaction_py: Borrowed<'a, 'py, PyAny>) -> PyResult { + let database_engine = detect_engine(&logging_transaction_py.to_owned())?; Ok(Self { - raw: obj.to_owned().unbind(), + logging_transaction_py: logging_transaction_py.to_owned().unbind(), database_engine, }) } @@ -99,14 +125,22 @@ impl LoggingTransactionWrapper { sql: &str, args: &Bound<'py, PyAny>, ) -> PyResult<()> { - let execute_fn = self.raw.bind(py).getattr(intern!(py, "execute"))?; + let execute_fn = self + .logging_transaction_py + .bind(py) + .getattr(intern!(py, "execute"))?; execute_fn.call1((sql, args))?; Ok(()) } } +#[async_trait::async_trait] impl Transaction for LoggingTransactionWrapper { - fn query(&self, sql: &str, args: &[&str]) -> () { - self.execute(sql, args); + async fn query(&self, sql: &str, args: &[&str]) -> Vec { + self.execute(sql, args).await; + } + + async fn commit(&self) -> Result<(), anyhow::Error> { + // In Synapse, `commit` is part of `LoggingDatabaseConnection` } } diff --git a/rust/src/storage/db/rust_db_pool.rs b/rust/src/storage/db/rust_db_pool.rs index 83206bb5fc..c12cd81199 100644 --- a/rust/src/storage/db/rust_db_pool.rs +++ b/rust/src/storage/db/rust_db_pool.rs @@ -13,23 +13,26 @@ * */ - // TODO: remove. This is just here to make sure our `DatabasePool`/`Transaction` // interfaces are compatible with `tokio-postgres`. use anyhow::Context; -use bb8_postgres::{tokio_postgres::Row, PostgresConnectionManager}; +use bb8_postgres::PostgresConnectionManager; use postgres_native_tls::MakeTlsConnector; -use crate::storage::db::{DatabasePool, Transaction}; +use crate::storage::db::{DatabasePool, Row, Transaction}; /// Native Rust database access backed by `tokio-postgres` (for use in synapse-rust-apps) pub struct RustDatabasePool { db_pool: bb8::Pool>, } +#[async_trait::async_trait] impl DatabasePool for RustDatabasePool { - async fn get_transaction(&self, description: &str) -> dyn Transaction { + async fn get_transaction( + &self, + _description: &str, + ) -> Result, anyhow::Error> { let mut conn = self .db_pool .get() @@ -37,14 +40,15 @@ impl DatabasePool for RustDatabasePool { .await .context("Failed to acquire database connection")?; + + // TODO: Set repeatable-read isolation level (like Synapse) let txn = conn .transaction() // .instrument(tracing::info_span!("start transaction")) .await .context("Failed to start transaction")?; - // TODO: Set isolation level - txn + Ok(Box::new(TokioPostgresTransaction { txn })) } } @@ -52,20 +56,21 @@ struct TokioPostgresTransaction<'a> { txn: bb8_postgres::tokio_postgres::Transaction<'a>, } +#[async_trait::async_trait] impl Transaction for TokioPostgresTransaction<'_> { - async fn query(&self, sql: &str, args: &[&str]) -> () { - todo!("TODO"); + async fn query(&self, sql: &str, args: &[&str]) -> Vec { + // TODO: Convert `?` SQL param style to `tokio-postgres` compatible let rows = self.txn.query(sql, args).await?; rows } - async fn commit(&self) -> () { + async fn commit(self) -> Result<(), anyhow::Error> { self.txn .commit() // .instrument(tracing::info_span!("commit transaction")) .await - .context("Failed to commit transaction")?; + .context("Failed to commit transaction") } }