Slow going

Cargo.lock

@@ -1831,6 +1831,7 @@ name = "synapse"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "async-trait",
  "base64",
  "bb8",
  "bb8-postgres",

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"

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<Box<dyn Transaction>, 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<Row>;
+    async fn commit(self) -> Result<(), anyhow::Error>;
 }
+
+pub type Row = Vec<String>;

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<Box<dyn Transaction>, 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<Py<PyAny>> {
+                // 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<DatabaseEngine> {
 /// Use [`execute`](Self::execute) (or other methods) while holding the GIL.
 pub struct LoggingTransactionWrapper {
     /// The underlying `LoggingTransaction`
-    raw: Py<PyAny>,
+    logging_transaction_py: Py<PyAny>,
 
     /// 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<Self> {
-        let database_engine = detect_engine(&obj.to_owned())?;
+    fn extract(logging_transaction_py: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
+        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<Row> {
+        self.execute(sql, args).await;
+    }
+
+    async fn commit(&self) -> Result<(), anyhow::Error> {
+        // In Synapse, `commit` is part of `LoggingDatabaseConnection`
     }
 }

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<PostgresConnectionManager<MakeTlsConnector>>,
 }
 
+#[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<Box<dyn Transaction>, 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<Row> {
+        // 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")
     }
 }