dandri/synapse
1
0
Fork 0
mirror of https://github.com/element-hq/synapse.git synced 2026-08-14 13:40:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add Postgres Connection/Cursor types and connect()

Browse Source 
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<Client>` 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) <noreply@anthropic.com>
This commit is contained in:
Erik Johnston
2026-07-11 09:04:57 +00:00
co-authored by Claude Opus 4.8
parent 3d49cd0406
commit 2e6fdbbfee
5 changed files with 666 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/src/database/postgres/connection.rs
+433
View File
@@ -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<T>,
noun: &str,
on_poison: impl FnOnce(&mut T),
) -> PyResult<MutexGuard<'a, T>> {
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<Mutex<...>>` 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<Mutex<ConnInner>>,
}
/// 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<Client>,
/// 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<MutexGuard<'_, ConnInner>> {
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<R>(
&self,
py: Python<'_>,
f: impl FnOnce(&Client) -> PyResult<R>,
) -> PyResult<R> {
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<Self>) -> Py<Self> {
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<Bound<'_, PyAny>>,
_exc_value: Option<Bound<'_, PyAny>>,
_traceback: Option<Bound<'_, PyAny>>,
) -> PyResult<bool> {
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<Mutex<CursorQueryState>>,
}
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<MutexGuard<'_, CursorQueryState>> {
// 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<Vec<PgValue>>) -> 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::<Vec<_>>();
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<Option<Bound<'py, PyTuple>>> {
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<Vec<Bound<'py, PyTuple>>> {
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<Vec<Bound<'py, PyTuple>>> {
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<Bound<'py, PyInt>> {
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(())
}
}
rust/src/database/postgres/cursor_state.rs
+84 -9
View File
@@ -284,16 +284,19 @@ impl<S: CursorRowStream> CursorQueryState<S> {
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<Bound<'py, PyInt>> {
// `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<S: CursorRowStream> CursorQueryState<S> {
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<Bound<'py, PyInt>> {
// `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::<FakeStream>::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::<FakeStream>::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::<FakeStream>::new();
state.finish_no_rows(py).unwrap();
assert!(matches!(state, CursorQueryState::Idle));
});
}
#[test]
fn fetch_all_drains_and_closes() {
Python::initialize();
rust/src/database/postgres/mod.rs
+45 -20
View File
@@ -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::<connection::Connection>()?;
child.add_class::<connection::Cursor>()?;
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<Bound<'py, connection::Connection>> {
let config = dsn
.parse::<tokio_postgres::Config>()
.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)
}
synapse/synapse_rust/database/__init__.pyi
+15
View File
@@ -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:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
# 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.
synapse/synapse_rust/database/postgres.pyi
+89
View File
@@ -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:
# <https://www.gnu.org/licenses/agpl-3.0.html>.
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): ...