//! [`tokio_postgres`]-backed `Connection` / `Cursor` types exposed to Python. //! //! The driver itself is async; we drive it from sync Python methods via the //! extension's shared multi-thread tokio runtime (see [`crate::tokio_runtime`]). //! [`connect`] takes the runtime's handle from the reactor once and hands it to //! the [`Connection`], which carries it for the life of the connection. use anyhow::Error; use log::warn; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyModule; use crate::database::postgres::helpers::BlockingPostgresResult; use crate::tokio_runtime::runtime_handle; mod connection; mod cursor_state; mod helpers; mod libpq; mod value; /// Register the `postgres` submodule (the `Connection` / `Cursor` classes and /// the `connect` factory) under the parent `database` module. 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(()) } /// Map a [`tokio_postgres`] error into a Python `RuntimeError`. fn pg_err_to_py(e: tokio_postgres::Error) -> PyErr { PyRuntimeError::new_err(format!("postgres error: {e}")) } /// Open a new Postgres connection from a libpq-style DSN. /// /// Blocks until the connection is established, then spawns the long-lived /// connection task (which drives the socket) onto the shared runtime and /// hands back a `Connection` wrapping the client. /// /// `reactor` is the Twisted reactor the extension's shared runtime is stored /// on; the runtime is started on demand if the reactor hasn't run yet (so this /// works during schema setup and in tests). The resulting handle is stored on /// the returned [`Connection`] and used for every subsequent call on it. #[pyfunction] fn connect<'py>( py: Python<'py>, reactor: &Bound<'py, PyAny>, dsn: &str, ) -> PyResult> { let handle = runtime_handle(reactor)?; let config = fixup_default_host(dsn) .map_err(|e| PyRuntimeError::new_err(format!("Failed to prepare DSN: {e}")))?; // TLS is not yet supported: unlike libpq (whose default is // `sslmode=prefer`), we never negotiate TLS regardless of the DSN's // sslmode. Supporting it is left to a follow-up. let (client, connection) = config .connect(tokio_postgres::NoTls) .block_on_result(py, &handle)?; // Spawn the connection task on the shared runtime. handle.spawn(async move { if let Err(e) = connection.await { warn!("postgres connection error: {e}"); } }); let conn = connection::Connection::new(client, handle); Bound::new(py, conn) } /// Fix up a DSN to ensure it has a host, using libpq's default host if /// necessary. /// /// [`tokio_postgres`] has a different default host than libpq, which is what /// Synapse previously used (and is what e.g. `psql` uses). libpq's default host /// is configurable, so when the DSN omits a host we ask libpq what its default /// would be and use that instead (see [`libpq::default_host`]). fn fixup_default_host(dsn: &str) -> Result { let mut config = dsn.parse::()?; // `tokio_postgres` parses only the DSN string (it does not consult `PGHOST` // or the compiled-in default), so an empty host list means the DSN really // omitted the host. A DSN that gives a `hostaddr` instead of a `host` is // still connectable as-is, so leave it alone too — injecting a default host // there would just confuse TLS/SNI. if !config.get_hosts().is_empty() || !config.get_hostaddrs().is_empty() { return Ok(config); } // Resolve libpq's default host without connecting (see `libpq::default_host`). let host = libpq::default_host()?; config.host(&host); Ok(config) }