synapse/rust/src/database/postgres/mod.rs

//! [`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::<connection::Connection>()?;
    child.add_class::<connection::Cursor>()?;
    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<Bound<'py, connection::Connection>> {
    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<tokio_postgres::Config, Error> {
    let mut config = dsn.parse::<tokio_postgres::Config>()?;

    // `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)
}