Add Postgres Connection/Cursor types and connect()

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-06 09:44:51 +00:00
co-authored by Claude Opus 4.8
parent 883a43a416
commit 848100e9ef
5 changed files with 672 additions and 29 deletions
+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();