//! [`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`]). //! The blocking helpers take that runtime from the calling thread's context //! (via [`tokio::runtime::Handle::current`]), so every thread that drives a //! `Connection` must have the shared runtime *entered* first (see //! [`helpers::BlockingPostgres`]). [`connect`] takes the runtime's handle from //! the reactor and enters it while establishing 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 errors; mod helpers; mod libpq; mod value; /// Register the `postgres` submodule (the `Connection` / `Cursor` classes, the /// DBAPI2 exception hierarchy 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)?)?; errors::register_exceptions(py, &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(()) } /// 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 runtime is entered for the /// duration of this call so the blocking connect resolves it via /// [`tokio::runtime::Handle::current`]; subsequent calls on the returned /// [`Connection`] rely on their own thread having the runtime entered. #[pyfunction] fn connect<'py>( py: Python<'py>, reactor: &Bound<'py, PyAny>, dsn: &str, ) -> PyResult> { let handle = runtime_handle(reactor)?; // The blocking helpers drive their futures on the runtime entered on the // calling thread (see [`helpers::BlockingPostgres`]), so enter the shared // runtime for the duration of establishing the connection. let _guard = handle.enter(); 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)?; // 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); 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) }