mirror of
https://github.com/element-hq/synapse.git
synced 2026-08-28 20:08:16 +00:00
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:
co-authored by
Claude Opus 4.8
parent
883a43a416
commit
848100e9ef
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user