WIP

rust/src/database/mod.rs

@@ -15,6 +15,8 @@ use pyo3::types::PyModule;
 pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     let child = PyModule::new(py, "database")?;
 
+    postgres::register_module(py, &child)?;
+
     m.add_submodule(&child)?;
 
     // Mirror the convention used by other rust submodules so

rust/src/database/postgres/connection.rs

@@ -148,7 +148,7 @@ impl Cursor {
             return Ok(None);
         };
 
-        let py_row = pg_row_to_py(py, &row)?;
+        let py_row = pg_row_to_py(&row)?;
         Ok(Some(py_row))
     }
 
@@ -156,7 +156,7 @@ impl Cursor {
         let mut inner = self.py_lock()?;
         let rows = inner.query_state.fetch_all(py)?;
 
-        rows.into_iter().map(|row| pg_row_to_py(py, &row)).collect()
+        rows.into_iter().map(|row| pg_row_to_py(&row)).collect()
     }
 
     fn rowcount<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyInt>> {

rust/src/database/postgres/mod.rs

@@ -5,14 +5,54 @@
 //! methods are kept `pub` so that future Rust callers can drive them
 //! directly without going through the PyO3 wrappers.
 
+use log::warn;
 use pyo3::exceptions::PyRuntimeError;
 use pyo3::prelude::*;
+use pyo3::types::PyModule;
+
+use crate::database::postgres::helpers::BlockingPostgresResult;
+use crate::database::runtime::runtime;
 
 mod connection;
 mod cursor_state;
 mod helpers;
 mod value;
 
+pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    let child = PyModule::new(py, "postgres")?;
+
+    child.add_class::<connection::Connection>()?;
+    child.add_class::<connection::Cursor>()?;
+    child.add_function(wrap_pyfunction!(connect, &child)?)?;
+
+    m.add_submodule(&child)?;
+
+    // We need to manually add the module to sys.modules to make `from
+    // synapse.synapse_rust.database import postgres` work.
+    py.import("sys")?
+        .getattr("modules")?
+        .set_item("synapse.synapse_rust.database.postgres", child)?;
+
+    Ok(())
+}
+
 fn pg_err_to_py(e: tokio_postgres::Error) -> PyErr {
     PyRuntimeError::new_err(format!("postgres error: {e}"))
 }
+
+#[pyfunction]
+fn connect<'py>(py: Python<'py>, dsn: &str) -> PyResult<Bound<'py, connection::Connection>> {
+    let (client, connection) =
+        tokio_postgres::connect(dsn, tokio_postgres::NoTls).block_on_result(py)?;
+
+    // Spawn the connection task on the runtime.
+    runtime().spawn(async move {
+        if let Err(e) = connection.await {
+            warn!("postgres connection error: {e}");
+        }
+    });
+
+    let conn = connection::Connection::new(client);
+
+    Bound::new(py, conn)
+}

rust/src/database/postgres/value.rs

@@ -112,11 +112,11 @@ impl ToSql for PgValue {
                 float8_to_sql(v, buf);
                 Ok(IsNull::No)
             }
-            (&PgValue::Text(ref v), &Type::TEXT | &Type::VARCHAR | &Type::NAME | &Type::BPCHAR) => {
+            (PgValue::Text(v), &Type::TEXT | &Type::VARCHAR | &Type::NAME | &Type::BPCHAR) => {
                 text_to_sql(v, buf);
                 Ok(IsNull::No)
             }
-            (&PgValue::Bytea(ref v), &Type::BYTEA) => {
+            (PgValue::Bytea(v), &Type::BYTEA) => {
                 bytea_to_sql(v, buf);
                 Ok(IsNull::No)
             }
@@ -150,27 +150,16 @@ impl ToSql for PgValue {
     to_sql_checked!();
 }
 
-/// Convert one column of a `tokio_postgres::Row` into a Python object.
-pub fn pg_column_to_py(
-    py: Python<'_>,
-    row: &tokio_postgres::Row,
-    idx: usize,
-) -> PyResult<Option<Py<PyAny>>> {
-    let obj: PythonPgFromSql = row.try_get(idx).map_err(|e| {
-        PyValueError::new_err(format!(
-            "failed to decode column {idx} (type {}): {e}",
-            row.columns()[idx].type_()
-        ))
-    })?;
-
-    Ok(obj.0)
-}
-
-pub fn pg_row_to_py(py: Python<'_>, row: &tokio_postgres::Row) -> PyResult<Vec<Option<Py<PyAny>>>> {
+pub fn pg_row_to_py(row: &tokio_postgres::Row) -> PyResult<Vec<Option<Py<PyAny>>>> {
     let mut result = Vec::with_capacity(row.len());
     for idx in 0..row.len() {
-        let obj = pg_column_to_py(py, row, idx)?;
-        result.push(obj);
+        let obj: PythonPgFromSql = row.try_get(idx).map_err(|e| {
+            PyValueError::new_err(format!(
+                "failed to decode column {idx} (type {}): {e}",
+                row.columns()[idx].type_()
+            ))
+        })?;
+        result.push(obj.0);
     }
     Ok(result)
 }
@@ -233,7 +222,7 @@ impl PythonPgFromSql {
             }
             Type::FLOAT8 => {
                 let f = postgres_protocol::types::float8_from_sql(raw)?;
-                PyFloat::new(py, f.into()).into_any().unbind()
+                PyFloat::new(py, f).into_any().unbind()
             }
             Type::TEXT | Type::VARCHAR | Type::NAME | Type::BPCHAR => {
                 PyString::from_bytes(py, raw)?.into_any().unbind()