From 2e6fdbbfeecec9d8cbeef26de40ed6da03e9e73a Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Mon, 29 Jun 2026 19:53:56 +0000 Subject: [PATCH] Add Postgres Connection/Cursor types and connect() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the Python-facing `Connection` / `Cursor` pair and the `connect` factory, implementing enough of the PEP-249 (DBAPI2) shape for Synapse's needs. A `Connection` owns a single `tokio_postgres::Client`. The client *moves* between the connection and an in-flight cursor rather than being shared: it lives in the connection between interactions, is taken out for the duration of a cursor, and is handed back when the cursor finishes. That single-owner baton (an `Option` slot on each side) makes it structurally impossible to use the connection mid-transaction or to run two overlapping transactions on one socket — both become a clean "already closed" error. The transaction lifecycle (`cursor` opens with `BEGIN`; `finish` COMMIT/ROLLBACKs and hands the client back; `CursorGuard` rolls back an abandoned transaction on drop) is included here. On any transaction-control error the client is dropped rather than returned, closing the socket — safer than handing a possibly-broken connection back to what will become a connection pool. `connect()` parses a libpq-style DSN, blocks until connected, and spawns the long-lived connection task onto the shared runtime (the libpq default-host fixup is a follow-up). The cursor query methods (`execute` / `fetch_one` / `fetch_all` / `fetch_next_batch` / `rowcount`) delegate to the `CursorQueryState` machine. `run_interaction` — the high-level glue that opens a cursor, runs the callback, and commits/rolls back — follows in the next change, so `cursor`/`finish`/`CursorGuard` carry a transitional `#[allow(dead_code)]` until then. Now that the value/helpers/cursor_state modules are consumed internally, their visibility is tightened from `pub` back to private. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/src/database/postgres/connection.rs | 433 +++++++++++++++++++++ rust/src/database/postgres/cursor_state.rs | 93 ++++- rust/src/database/postgres/mod.rs | 65 +++- synapse/synapse_rust/database/__init__.pyi | 15 + synapse/synapse_rust/database/postgres.pyi | 89 +++++ 5 files changed, 666 insertions(+), 29 deletions(-) create mode 100644 rust/src/database/postgres/connection.rs create mode 100644 synapse/synapse_rust/database/__init__.pyi create mode 100644 synapse/synapse_rust/database/postgres.pyi diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs new file mode 100644 index 0000000000..6ad158de4b --- /dev/null +++ b/rust/src/database/postgres/connection.rs @@ -0,0 +1,433 @@ +//! The Python-facing [`Connection`] / [`Cursor`] pair for the Postgres backend. +//! +//! These implement enough of the PEP-249 (DBAPI2) shape for Synapse's needs, +//! arranged so that the existing Python transaction driver +//! (`synapse.storage.database.new_transaction`) can drive them *unchanged*: it +//! opens a [`Cursor`] with [`Connection::cursor`], runs the interaction +//! function against it, and then commits or rolls back at the **connection** +//! level with [`Connection::commit`] / [`Connection::rollback`]. +//! +//! The driver underneath is async but the Python API is sync, so every database +//! call is driven to completion on the shared tokio runtime via the `block_on` +//! helpers (see [`super::helpers::BlockingPostgresResult`]). +//! +//! ## Connection owns the transaction; the cursor is a thin view +//! +//! The single [`tokio_postgres::Client`] lives in the [`Connection`] for the +//! whole life of the connection. Transaction control (`BEGIN` / `COMMIT` / +//! `ROLLBACK`) is issued *on the connection*, because that is where Synapse's +//! driver issues it — `conn.commit()` runs while the cursor that produced the +//! rows is still open (the driver only closes the cursor afterwards). +//! +//! A [`Cursor`] is therefore cheap: it holds a clone of the owning +//! [`Connection`] (an `Arc`) plus its own result-set state +//! ([`CursorQueryState`]). It borrows the client only for the brief moment it +//! takes to *start* a query (`prepare` + `query_raw`); the resulting row stream +//! is self-contained (`'static`), so once a query has been issued the cursor +//! reads rows from its own state without touching the connection again. Many +//! cursors can share one connection this way, though in practice Synapse's +//! driver uses one at a time. +//! +//! ## Implicit transactions (matching psycopg2) +//! +//! psycopg2 is transactional by default: the first statement after `connect` +//! (or after a `commit`/`rollback`) implicitly opens a transaction. +//! `tokio_postgres`, by contrast, is autocommit by default. To behave like +//! psycopg2 we track whether a transaction is open (`in_txn`) and lazily issue a +//! `BEGIN` before the first statement of each transaction — unless the +//! connection has been put into autocommit mode (see [`Connection::set_autocommit`]). +//! `commit`/`rollback` end the transaction and clear the flag; with no +//! transaction open they are no-ops, just like psycopg2. +//! +//! ## Dropping the `Client` on error +//! +//! A *query* error (bad SQL, a constraint violation, an integer out of range, +//! …) leaves the connection open with its transaction in the aborted state, +//! exactly as psycopg2 does: the error propagates to Python and the driver is +//! expected to `rollback()`. We do **not** throw the connection away for these. +//! +//! The transaction-control statements are different. If `COMMIT` or `ROLLBACK` +//! itself fails we no longer know what state the server-side session is in, so +//! we drop the `Client` (closing the socket) rather than hand a possibly-broken +//! connection back for reuse. Likewise, [`Connection::close`] drops the client; +//! the server rolls back any transaction left open when the socket closes. + +use std::sync::{Arc, Mutex, MutexGuard, TryLockError}; + +use pyo3::{ + exceptions::PyRuntimeError, + prelude::*, + types::{PyInt, PyTuple}, +}; +use tokio_postgres::Client; + +use crate::database::postgres::{ + cursor_state::CursorQueryState, helpers::BlockingPostgresResult, value::PgValue, +}; + +/// `try_lock` a mutex that is single-threaded by contract, mapping its two +/// failure modes to Python errors. +/// +/// A poisoned mutex means a panic happened while it was held, so the guarded +/// state can't be trusted: `on_poison` resets it (closing a connection / +/// discarding a result set) before we error. `WouldBlock` means the object is +/// being used from two threads at once, which is a caller bug rather than +/// something to block on. `noun` names the object in both messages (e.g. +/// `"connection"`, `"cursor"`). +fn try_lock_or_reset<'a, T>( + mutex: &'a Mutex, + noun: &str, + on_poison: impl FnOnce(&mut T), +) -> PyResult> { + match mutex.try_lock() { + Ok(guard) => Ok(guard), + Err(TryLockError::Poisoned(poisoned)) => { + on_poison(&mut poisoned.into_inner()); + Err(PyRuntimeError::new_err(format!("{noun} mutex poisoned"))) + } + Err(TryLockError::WouldBlock) => Err(PyRuntimeError::new_err(format!( + "{noun} is being used in another thread and cannot be used concurrently" + ))), + } +} + +/// A single Postgres connection exposed to Python. +/// +/// Owns the [`tokio_postgres::Client`] for its whole life and is the authority +/// on transaction state. The `Arc>` lets cursors hold a cheap clone +/// (so they can reach the client to start a query) while keeping all access to +/// the client serialised. +#[pyclass(frozen, skip_from_py_object)] +#[derive(Clone)] +pub struct Connection { + inner: Arc>, +} + +/// The mutable guts of a [`Connection`], behind its mutex. +struct ConnInner { + /// The driver client. `None` once the connection has been closed (or thrown + /// away after a transaction-control error); any further use is an error. + client: Option, + /// Whether a transaction is currently open (a `BEGIN` has been issued and + /// not yet matched by a `COMMIT`/`ROLLBACK`). Drives the lazy `BEGIN`. + in_txn: bool, + /// In autocommit mode we never issue an implicit `BEGIN`, so each statement + /// runs in its own implicit transaction. Defaults to `false`, matching + /// psycopg2's transactional default. + autocommit: bool, +} + +impl Connection { + /// Wrap a freshly-established `Client` in a `Connection`. + pub fn new(client: Client) -> Self { + Self { + inner: Arc::new(Mutex::new(ConnInner { + client: Some(client), + in_txn: false, + autocommit: false, + })), + } + } + + /// Lock the inner state. + /// + /// Uses `try_lock` rather than `lock`: a connection is used from a single + /// thread at a time by contract (Synapse hands one connection to one + /// worker thread), so contention means it's being used from two threads at + /// once, which we surface as an error instead of blocking. A poisoned mutex + /// (a panic happened mid-operation) closes the connection — we no longer + /// know the session state — and errors. + fn lock(&self) -> PyResult> { + try_lock_or_reset(&self.inner, "connection", |inner| { + // On poison we no longer know the session state, so close the + // connection: drop the client and clear the transaction flag. + inner.client = None; + inner.in_txn = false; + }) + } + + /// Borrow the client just long enough to run `f` (typically starting a + /// query), opening an implicit transaction first if one isn't already open. + /// + /// The borrow ends as soon as `f` returns; `f` is expected to hand back an + /// owned, self-contained value (e.g. a `'static` row stream) rather than + /// anything tied to the client. Errors if the connection is closed. + fn with_client( + &self, + py: Python<'_>, + f: impl FnOnce(&Client) -> PyResult, + ) -> PyResult { + let mut guard = self.lock()?; + + // Lazily open a transaction so statements are transactional by default, + // matching psycopg2. We set `in_txn` *after* a successful `BEGIN` but + // before running `f`, so that if `f` (the user's statement) fails the + // open-but-aborted transaction is still tracked and `rollback()` knows + // to clean it up. + if !guard.autocommit && !guard.in_txn { + { + let client = client_ref(&guard)?; + client.execute("BEGIN", &[]).block_on_result(py)?; + } + guard.in_txn = true; + } + + let client = client_ref(&guard)?; + f(client) + } + + /// Issue a transaction-control statement (`COMMIT`/`ROLLBACK`) if a + /// transaction is open; a no-op otherwise. + /// + /// On success the transaction flag is cleared. On failure the client is + /// dropped (closing the socket): after a failed commit/rollback the session + /// state is unknown, so the connection is thrown away rather than reused. + fn end_txn(&self, py: Python<'_>, stmt: &'static str) -> PyResult<()> { + let mut guard = self.lock()?; + + if !guard.in_txn { + return Ok(()); + } + + // If the client is already gone the server has rolled the transaction + // back for us; just clear the flag. + if guard.client.is_none() { + guard.in_txn = false; + return Ok(()); + } + + let result = { + let client = client_ref(&guard)?; + client.execute(stmt, &[]).block_on_result(py) + }; + + match result { + Ok(_) => { + guard.in_txn = false; + Ok(()) + } + Err(err) => { + // Unknown session state: drop the connection rather than reuse it. + guard.client = None; + guard.in_txn = false; + Err(err) + } + } + } +} + +/// Borrow the live client out of a locked inner state, or error if the +/// connection has been closed. +fn client_ref(guard: &ConnInner) -> PyResult<&Client> { + guard + .client + .as_ref() + .ok_or_else(|| PyRuntimeError::new_err("connection already closed")) +} + +#[pymethods] +impl Connection { + /// Open a new cursor over this connection. + /// + /// Cheap: no I/O and no `BEGIN` happens here (the transaction is opened + /// lazily on the first `execute`). The returned cursor shares this + /// connection's client. + fn cursor(&self) -> Cursor { + Cursor::new(self.clone()) + } + + /// Commit the current transaction, if one is open. A no-op otherwise. + fn commit(&self, py: Python<'_>) -> PyResult<()> { + self.end_txn(py, "COMMIT") + } + + /// Roll back the current transaction, if one is open. A no-op otherwise. + fn rollback(&self, py: Python<'_>) -> PyResult<()> { + self.end_txn(py, "ROLLBACK") + } + + /// Close the connection, dropping the underlying client. + /// + /// Dropping the client closes the socket; the server rolls back any + /// transaction that was still open. Idempotent: closing an + /// already-closed connection is fine. + fn close(&self) -> PyResult<()> { + let mut guard = self.lock()?; + guard.client = None; + guard.in_txn = false; + Ok(()) + } + + /// Switch autocommit mode on or off. + /// + /// In autocommit mode no implicit `BEGIN` is issued, so each statement runs + /// in its own transaction. Mirrors psycopg2's `set_session(autocommit=...)`, + /// including its rule that the mode can't be changed while a transaction is + /// in progress. + fn set_autocommit(&self, autocommit: bool) -> PyResult<()> { + let mut guard = self.lock()?; + if guard.in_txn { + return Err(PyRuntimeError::new_err( + "cannot change autocommit mode while a transaction is in progress", + )); + } + guard.autocommit = autocommit; + Ok(()) + } + + /// Context-manager entry: returns the connection itself. + fn __enter__(slf: Py) -> Py { + slf + } + + /// Context-manager exit: commit if the block completed normally, roll back + /// if it raised. Like psycopg2, this does *not* close the connection, and + /// it does not suppress the exception (returns `False`). + fn __exit__( + &self, + py: Python<'_>, + exc_type: Option>, + _exc_value: Option>, + _traceback: Option>, + ) -> PyResult { + if exc_type.is_some() { + self.rollback(py)?; + } else { + self.commit(py)?; + } + Ok(false) + } +} + +/// A PEP-249-style cursor over a connection's current transaction. +/// +/// Cheap to create: it holds a clone of the owning [`Connection`] plus its own +/// result-set state ([`CursorQueryState`]). Transaction control lives on the +/// connection, not here. All interior state is behind a mutex so the cursor can +/// be `frozen` (shared via `Arc`) yet still mutated by its methods. +#[pyclass(frozen, skip_from_py_object)] +#[derive(Clone)] +pub struct Cursor { + /// The owning connection, used to reach the client when starting a query. + connection: Connection, + /// State of the most recent `execute` (live row stream, rowcount, etc.). + state: Arc>, +} + +impl Cursor { + /// Build a cursor over `connection`. + fn new(connection: Connection) -> Self { + Self { + connection, + state: Arc::new(Mutex::new(CursorQueryState::new())), + } + } + + /// Lock the cursor's query state. + /// + /// Like [`Connection::lock`], uses `try_lock`: a cursor is single-threaded + /// by contract, so contention means concurrent use, which we surface as an + /// error. On a poisoned mutex we drop the result set (resetting to `Idle`) + /// and error. + fn lock_state(&self) -> PyResult> { + // On poison we drop the (untrusted) result set, resetting to `Idle`. + try_lock_or_reset(&self.state, "cursor", |state| { + *state = CursorQueryState::new() + }) + } +} + +#[pymethods] +impl Cursor { + /// Execute `query`, optionally with positional `params` bound to `$1`, + /// `$2`, ... placeholders. + /// + /// Any previous result set is discarded. After this returns, rows (if any) + /// can be read with `fetch_one`/`fetch_all`/`fetch_next_batch`. + #[pyo3(signature = (query, params = None))] + fn execute(&self, py: Python<'_>, query: &str, params: Option>) -> PyResult<()> { + // Drop any previous result set before starting the new query. + self.lock_state()?.new_query(); + + // Borrow the connection's client only for as long as it takes to start + // the query. `query_raw` returns a `'static` `RowStream`, so the borrow + // ends here and the cursor owns the stream from now on. + let (stream, description) = self.connection.with_client(py, |client| { + let statement = client.prepare(query).block_on_result(py)?; + + let stream = client + .query_raw(&statement, params.unwrap_or_default()) + .block_on_result(py)?; + + // The column names back the (future) PEP-249 `Cursor.description`; + // pull them out of the prepared statement here so `cursor_state` + // stays decoupled from `tokio_postgres::Column`. + let description = statement + .columns() + .iter() + .map(|c| c.name().to_string()) + .collect::>(); + + Ok((stream, description)) + })?; + + // A statement with no result columns (INSERT/UPDATE/DELETE/DDL without + // RETURNING) produces no rows for anyone to fetch, so nothing would + // otherwise poll its stream — and a `query_raw` stream only reports the + // server's response, including any error (e.g. a constraint violation) + // and the affected-row count, once polled. Drive it to completion now so + // such errors surface here at `execute` time (as psycopg2 does) rather + // than being lost when the result is never fetched — most visibly under + // autocommit, where there is no later `commit` to surface them. + let returns_rows = !description.is_empty(); + let mut state = self.lock_state()?; + state.on_query_start(stream, description); + if !returns_rows { + state.finish_no_rows(py)?; + } + + Ok(()) + } + + /// Return the next row of the current result set, or `None` if exhausted. + fn fetch_one<'py>(&self, py: Python<'py>) -> PyResult>> { + self.lock_state()?.fetch_one(py) + } + + /// Drain and return all remaining rows of the current result set. + fn fetch_all<'py>(&self, py: Python<'py>) -> PyResult>> { + self.lock_state()?.fetch_all(py) + } + + /// Fetch the next batch of rows from the current result set. + /// + /// Blocks for the first row, then returns any further rows that are already + /// available without blocking. Returns an empty list only once the result + /// set is exhausted. `capacity` is a hint for the size of the returned + /// buffer, not a limit on the number of rows returned. + #[pyo3(signature = (capacity = 100))] + fn fetch_next_batch<'py>( + &self, + py: Python<'py>, + capacity: usize, + ) -> PyResult>> { + self.lock_state()?.fetch_next_batch(py, capacity) + } + + /// Return the PEP-249 `rowcount` for the last statement. + /// + /// This is the number of rows affected by a DML statement; for queries + /// where it isn't (yet) known it follows PEP-249 and returns `-1`. + fn rowcount<'py>(&self, py: Python<'py>) -> PyResult> { + self.lock_state()?.rowcount(py) + } + + /// Close the cursor, discarding any in-flight result set. + /// + /// This does *not* touch the transaction — that's the connection's job. + /// Idempotent. + fn close(&self) -> PyResult<()> { + *self.lock_state()? = CursorQueryState::new(); + Ok(()) + } +} diff --git a/rust/src/database/postgres/cursor_state.rs b/rust/src/database/postgres/cursor_state.rs index 653c78f4d5..f8f8ea2898 100644 --- a/rust/src/database/postgres/cursor_state.rs +++ b/rust/src/database/postgres/cursor_state.rs @@ -284,16 +284,19 @@ impl CursorQueryState { Ok(rows) } - /// Return the affected-row count, draining the stream first if needed. + /// Drive a non-row-returning statement's stream to completion, surfacing + /// any execution error and capturing the affected-row count. A no-op unless + /// the cursor is `Active`. /// - /// Unlike the `fetch_*` methods this is always valid: reading the rowcount - /// of an already-exhausted (`Closed`) cursor returns the captured count - /// rather than erroring, per PEP-249. - pub fn rowcount<'py>(&mut self, py: Python<'py>) -> PyResult> { - // `rows_affected()` is only valid after the stream is drained, so we - // drain it here. This is OK as in Python the rowcount should only be - // accessed for queries that DO NOT return rows, e.g. INSERT, UPDATE, - // DELETE. + /// `execute` calls this for statements with no result columns + /// (INSERT/UPDATE/DELETE/DDL without RETURNING). Their `query_raw` stream is + /// never fetched, but a `query_raw` stream only reports the server's + /// response — including an error such as a constraint violation, and the + /// affected-row count — once it is polled. Draining here makes such an error + /// surface at `execute` time, as psycopg2 does, instead of being silently + /// lost when the result is never fetched (most visibly under autocommit, + /// where there is no later `commit` to surface it). + pub fn finish_no_rows(&mut self, py: Python<'_>) -> PyResult<()> { if let Self::Active { stream, description, @@ -309,6 +312,19 @@ impl CursorQueryState { rowcount, }; } + Ok(()) + } + + /// Return the affected-row count, draining the stream first if needed. + /// + /// Unlike the `fetch_*` methods this is always valid: reading the rowcount + /// of an already-exhausted (`Closed`) cursor returns the captured count + /// rather than erroring, per PEP-249. + pub fn rowcount<'py>(&mut self, py: Python<'py>) -> PyResult> { + // `rows_affected()` is only valid after the stream is drained, so drain + // it here. This is OK as in Python the rowcount should only be accessed + // for queries that DO NOT return rows, e.g. INSERT, UPDATE, DELETE. + self.finish_no_rows(py)?; match self { Self::Closed { @@ -510,6 +526,65 @@ mod tests { }); } + #[test] + fn finish_no_rows_drains_and_captures_rowcount() { + Python::initialize(); + Python::attach(|py| { + // A non-row-returning statement: no rows to yield, a command-tag + // rowcount, and an empty column list. + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::new(), + rows_affected: Some(3), + }, + vec![], + ); + + state.finish_no_rows(py).unwrap(); + assert!(matches!( + state, + CursorQueryState::Closed { + rowcount: Some(3), + .. + } + )); + }); + } + + #[test] + fn finish_no_rows_surfaces_execution_error() { + Python::initialize(); + Python::attach(|py| { + // The statement fails server-side (e.g. a constraint violation), + // surfaced only when the stream is polled — which `finish_no_rows` + // does, so the error is raised rather than silently swallowed. + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::from([Err(FakeError("boom"))]), + rows_affected: None, + }, + vec![], + ); + + let err = state.finish_no_rows(py).unwrap_err(); + assert!(err.to_string().contains("boom"), "{err}"); + // A stream error resets the cursor to `Idle`. + assert!(matches!(state, CursorQueryState::Idle)); + }); + } + + #[test] + fn finish_no_rows_on_idle_is_a_noop() { + Python::initialize(); + Python::attach(|py| { + let mut state = CursorQueryState::::new(); + state.finish_no_rows(py).unwrap(); + assert!(matches!(state, CursorQueryState::Idle)); + }); + } + #[test] fn fetch_all_drains_and_closes() { Python::initialize(); diff --git a/rust/src/database/postgres/mod.rs b/rust/src/database/postgres/mod.rs index 1513743cb3..7de4af7ede 100644 --- a/rust/src/database/postgres/mod.rs +++ b/rust/src/database/postgres/mod.rs @@ -1,33 +1,30 @@ -//! [`tokio_postgres`]-backed Postgres backend for the Rust `database` module. +//! [`tokio_postgres`]-backed `Connection` / `Cursor` types exposed to Python. //! -//! This module will grow the Python-facing `Connection` / `Cursor` classes and -//! the `connect` factory; for now it hosts the value-mapping layer ([`value`]) -//! that converts between Python objects and Postgres' binary wire format. -//! -//! The driver itself is async; the eventual `Connection` / `Cursor` types will -//! drive it from sync Python methods via a shared multi-thread tokio runtime. +//! The driver itself is async; we drive it from sync Python methods via a +//! shared multi-thread tokio runtime (see `super::runtime`). +use log::warn; use pyo3::exceptions::PyRuntimeError; use pyo3::prelude::*; use pyo3::types::PyModule; -// `pub` (rather than private) so the not-yet-consumed public items in these -// submodules are reachable from the crate root as public API. This is what -// stops clippy's `dead_code` lint from firing on them before the -// cursor/connection code (added in later changes) wires them up; the visibility -// is tightened back to private once that happens. -pub mod cursor_state; -pub mod helpers; -pub mod value; +use crate::database::postgres::helpers::BlockingPostgresResult; +use crate::database::runtime::runtime; -/// Register the `postgres` submodule under the parent `database` module. -/// -/// The `Connection` / `Cursor` classes and the `connect` factory are added in -/// later changes; for now this just creates the (otherwise empty) submodule so -/// the module tree — and the `value` mapping layer hanging off it — exists. +mod connection; +mod cursor_state; +mod helpers; +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 @@ -43,3 +40,31 @@ pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> 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. +#[pyfunction] +fn connect<'py>(py: Python<'py>, dsn: &str) -> PyResult> { + let config = dsn + .parse::() + .map_err(|e| PyRuntimeError::new_err(format!("Failed to parse 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 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/synapse/synapse_rust/database/__init__.pyi b/synapse/synapse_rust/database/__init__.pyi new file mode 100644 index 0000000000..87fc834c61 --- /dev/null +++ b/synapse/synapse_rust/database/__init__.pyi @@ -0,0 +1,15 @@ +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . + +# The `database` submodule of the Rust extension. Its `postgres` submodule +# (stubbed in `postgres.pyi`) holds the native connection pool + DBAPI2-ish +# connection/cursor shim and the PEP-249 exception hierarchy. diff --git a/synapse/synapse_rust/database/postgres.pyi b/synapse/synapse_rust/database/postgres.pyi new file mode 100644 index 0000000000..24303b3416 --- /dev/null +++ b/synapse/synapse_rust/database/postgres.pyi @@ -0,0 +1,89 @@ +# This file is licensed under the Affero General Public License (AGPL) version 3. +# +# Copyright (C) 2026 Element Creations Ltd +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as +# published by the Free Software Foundation, either version 3 of the +# License, or (at your option) any later version. +# +# See the GNU Affero General Public License for more details: +# . + +from types import TracebackType +from typing import Any, Optional, Sequence + +class ConnectionPool: + """A pool of native `tokio_postgres` connections. + + Built once from a libpq-style DSN; connections are opened lazily. See + :mod:`synapse.storage.rust_dbapi` / :mod:`synapse.storage.rust_pool` for the + DBAPI2 / Twisted adapters that wrap it. + """ + + def __init__( + self, + dsn: str, + max_size: int = 10, + *, + synchronous_commit: bool = True, + statement_timeout_ms: Optional[int] = None, + ) -> None: ... + def connect(self) -> Connection: + """Check a connection out of the pool, blocking until one is available. + + Raises an ``OperationalError`` if the configured ``checkout_timeout_ms`` + elapses first. + """ + + def close(self) -> None: ... + +class Connection: + """A connection checked out of a :class:`ConnectionPool`.""" + + @property + def autocommit(self) -> bool: ... + def cursor(self) -> Cursor: ... + def commit(self) -> None: ... + def rollback(self) -> None: ... + def close(self) -> None: ... + def is_closed(self) -> bool: ... + def in_transaction(self) -> bool: ... + def set_autocommit(self, autocommit: bool) -> None: ... + def __enter__(self) -> Connection: ... + def __exit__( + self, + exc_type: Optional[type[BaseException]], + exc: Optional[BaseException], + tb: Optional[TracebackType], + ) -> bool: ... + +class Cursor: + """A PEP-249-style cursor over a connection's current transaction. + + Result-fetching uses ``fetch_one`` / ``fetch_all`` / ``fetch_next_batch`` and + ``rowcount`` / ``description`` as *methods* (not DBAPI2 properties); the + :class:`synapse.storage.rust_dbapi.Cursor` adapter presents the DBAPI2 shape. + """ + + def execute(self, query: str, params: Optional[Sequence[Any]] = None) -> None: ... + def executemany(self, query: str, params_seq: Sequence[Sequence[Any]]) -> None: ... + def executescript(self, script: str) -> None: ... + def fetch_one(self) -> Optional[tuple[Any, ...]]: ... + def fetch_all(self) -> list[tuple[Any, ...]]: ... + def fetch_next_batch(self, capacity: int = 100) -> list[tuple[Any, ...]]: ... + def rowcount(self) -> int: ... + def description(self) -> Optional[list[tuple[Any, ...]]]: ... + def close(self) -> None: ... + +# The PEP-249 exception hierarchy raised by the backend. `DatabaseError` and its +# subclasses carry the SQLSTATE as `pgcode` (a 5-character string, or `None`), +# matching psycopg2 so engine code such as `is_deadlock` can read `error.pgcode`. +class Error(Exception): ... + +class DatabaseError(Error): + pgcode: Optional[str] + +class OperationalError(DatabaseError): ... +class IntegrityError(DatabaseError): ... +class ProgrammingError(DatabaseError): ...