diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index b21af633d6..166e1067fb 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -57,7 +57,7 @@ use std::sync::{Arc, Mutex, MutexGuard, TryLockError}; use pyo3::{ exceptions::PyRuntimeError, prelude::*, - types::{PyInt, PyTuple}, + types::{PyInt, PyList, PyTuple}, }; use tokio_postgres::Client; @@ -448,6 +448,45 @@ impl Cursor { self.lock_state()?.rowcount(py) } + /// Return the PEP-249 `description` for the current result set, or `None`. + /// + /// This is a list with one entry per column, each a 7-tuple + /// `(name, type_code, display_size, internal_size, precision, scale, + /// null_ok)` as PEP-249 specifies. Only the column name is populated; the + /// remaining six fields are always `None` — that is all Synapse needs, as + /// it only ever reads `column[0]`. + /// + /// It is `None` when there is no row-returning result set to describe: + /// before any query, after an error reset the cursor, or for a statement + /// that returns no rows (e.g. a bare `INSERT`), matching psycopg2. + fn description<'py>(&self, py: Python<'py>) -> PyResult>> { + let state = self.lock_state()?; + let Some(columns) = state.description() else { + return Ok(None); + }; + + let rows = columns + .iter() + .map(|name| { + // PEP-249's 7-tuple; only `name` carries a meaningful value. + PyTuple::new( + py, + [ + name.into_pyobject(py)?.into_any(), + py.None().into_bound(py), + py.None().into_bound(py), + py.None().into_bound(py), + py.None().into_bound(py), + py.None().into_bound(py), + py.None().into_bound(py), + ], + ) + }) + .collect::>>()?; + + Ok(Some(PyList::new(py, rows)?)) + } + /// Close the cursor, discarding any in-flight result set. /// /// This does *not* touch the transaction — that's the connection's job. diff --git a/rust/src/database/postgres/cursor_state.rs b/rust/src/database/postgres/cursor_state.rs index df2e287e37..7198ed957e 100644 --- a/rust/src/database/postgres/cursor_state.rs +++ b/rust/src/database/postgres/cursor_state.rs @@ -107,9 +107,7 @@ pub enum CursorQueryState { /// Live row stream for the current query. stream: FusedStream, /// Column names for the result set (empty for a DML statement). - /// - /// TODO: currently write-only; kept to back a future PEP-249 - /// `Cursor.description` accessor. + /// Exposed to Python via [`CursorQueryState::description`]. description: Vec, }, /// The result set has been fully consumed and exhaustion reported — a fetch @@ -117,13 +115,9 @@ pub enum CursorQueryState { /// `fetch_*` is a programming error; `rowcount` still returns the count. Closed { /// Column names for the result set, carried over from `Active` so they - /// survive once the rows are gone. - /// - /// TODO: currently write-only; like `Active::description` it is kept to - /// back a future PEP-249 `Cursor.description` accessor. `#[allow]`d - /// until that reader lands rather than dropped, so the column metadata - /// isn't silently lost at the `Active` -> `Closed` transition. - #[allow(dead_code)] + /// survive once the rows are gone (PEP-249 keeps `description` + /// available after the rows have been fetched). Exposed to Python via + /// [`CursorQueryState::description`]. description: Vec, /// PEP-249 `rowcount` from the command tag, if it was captured. rowcount: Option, @@ -344,6 +338,34 @@ impl CursorQueryState { } } + /// The column names of the current result set, or `None` when there is no + /// row-returning result set to describe. + /// + /// This backs the PEP-249 `Cursor.description`. It is `None` in the `Idle` + /// state (no query has run yet, or an error reset the cursor) and also for + /// a statement that returns no columns — e.g. an `INSERT`/`UPDATE`/`DELETE` + /// without `RETURNING` — matching psycopg2, which reports `description` as + /// `None` for such statements. For a row-returning statement the names stay + /// available after the rows have been fetched (the `Closed` state), as + /// PEP-249 requires. + /// + /// Only the column *names* are tracked (see `on_query_start`); it is the + /// caller's job to shape them into PEP-249's 7-tuples. + pub fn description(&self) -> Option<&[String]> { + let columns = match self { + Self::Active { description, .. } | Self::Closed { description, .. } => description, + Self::Idle => return None, + }; + + // A statement that returns no columns (DML) has an empty column list; + // PEP-249 / psycopg2 report `description` as `None` in that case. + if columns.is_empty() { + None + } else { + Some(columns) + } + } + /// The error to raise when a `fetch_*` method finds no rows available, /// distinguishing "no query was ever run" from "the result set has already /// been exhausted". Only meaningful for the non-`Active` states. @@ -771,6 +793,63 @@ mod tests { }); } + #[test] + fn description_is_none_before_any_query() { + let state = CursorQueryState::::new(); + assert!(state.description().is_none()); + } + + #[test] + fn description_reports_columns_while_active_and_after_close() { + Python::initialize(); + Python::attach(|py| { + // `active_with` builds an `Active` state whose columns are `["col"]`. + let mut state = active_with(vec![vec![1]], Some(0)); + assert_eq!(state.description(), Some(["col".to_string()].as_slice())); + + // Draining the rows moves the cursor to `Closed`, but the column + // names survive so `description` is still available (PEP-249). + let _ = state.fetch_all(py).unwrap(); + assert!(matches!(state, CursorQueryState::Closed { .. })); + assert_eq!(state.description(), Some(["col".to_string()].as_slice())); + }); + } + + #[test] + fn description_is_none_for_a_column_less_result() { + // A DML statement yields no columns; `description` should be `None` + // even while the (empty) result set is `Active`. + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::new(), + rows_affected: Some(3), + }, + vec![], + ); + assert!(state.description().is_none()); + } + + #[test] + fn description_is_none_after_an_error_resets_to_idle() { + Python::initialize(); + Python::attach(|py| { + let mut state = CursorQueryState::::new(); + state.on_query_start( + FakeStream { + items: VecDeque::from(vec![Err(FakeError("boom"))]), + rows_affected: None, + }, + vec!["col".to_string()], + ); + + // The error resets the cursor to `Idle`, dropping the columns. + let _ = state.fetch_one(py).unwrap_err(); + assert!(matches!(state, CursorQueryState::Idle)); + assert!(state.description().is_none()); + }); + } + #[test] fn new_query_resets_an_active_cursor() { Python::initialize(); diff --git a/tests/synapse_rust/test_database_postgres.py b/tests/synapse_rust/test_database_postgres.py index 492e04e1d8..b3578509b7 100644 --- a/tests/synapse_rust/test_database_postgres.py +++ b/tests/synapse_rust/test_database_postgres.py @@ -485,6 +485,55 @@ class PostgresConnectionTestCase(unittest.TestCase): self.assertEqual(run_interaction(self.conn, interaction), [3, 2, 3]) + # ------------------------------------------------------------------ + # description + # ------------------------------------------------------------------ + + def test_description_is_none_before_query(self) -> None: + """A cursor that has not run a query has no description.""" + + def interaction(cursor: Any) -> Any: + return cursor.description() + + self.assertIsNone(run_interaction(self.conn, interaction)) + + def test_description_reports_column_names(self) -> None: + """A row-returning statement describes its columns; only the name is + populated, in a PEP-249 7-tuple.""" + + def interaction(cursor: Any) -> Any: + cursor.execute("SELECT 1 AS a, 'x'::text AS b") + return cursor.description() + + description = run_interaction(self.conn, interaction) + self.assertEqual([col[0] for col in description], ["a", "b"]) + # Each entry is a PEP-249 7-tuple with only the name populated. + for col in description: + self.assertEqual(len(col), 7) + self.assertTrue(all(field is None for field in col[1:])) + + def test_description_available_after_fetch(self) -> None: + """The description survives after the rows have been fetched.""" + + def interaction(cursor: Any) -> Any: + cursor.execute("SELECT 1 AS a") + cursor.fetch_all() # exhausts the result set + return cursor.description() + + description = run_interaction(self.conn, interaction) + self.assertEqual([col[0] for col in description], ["a"]) + + def test_description_is_none_for_dml(self) -> None: + """A statement that returns no rows (a bare INSERT) has no + description, matching psycopg2.""" + + def interaction(cursor: Any) -> Any: + cursor.execute("CREATE TEMP TABLE d (id int)") + cursor.execute("INSERT INTO d VALUES (1)") + return cursor.description() + + self.assertIsNone(run_interaction(self.conn, interaction)) + # ------------------------------------------------------------------ # Transaction handling (COMMIT on success, ROLLBACK on error) # ------------------------------------------------------------------