diff --git a/rust/src/database/mod.rs b/rust/src/database/mod.rs index 578ac7a168..706ef80736 100644 --- a/rust/src/database/mod.rs +++ b/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 diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 86dbf74bac..0db2b899ed 100644 --- a/rust/src/database/postgres/connection.rs +++ b/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> { diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index ce7f17bd7e..5be089b990 100644 --- a/rust/src/database/postgres/mod.rs +++ b/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::()?; + child.add_class::()?; + 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> { + 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) +} diff --git a/rust/src/database/postgres/value.rs b/rust/src/database/postgres/value.rs index 8a4db44d1f..0b8bc4476a 100644 --- a/rust/src/database/postgres/value.rs +++ b/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>> { - 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>>> { +pub fn pg_row_to_py(row: &tokio_postgres::Row) -> PyResult>>> { 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()