diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 1048c65fb3..07eddb6f44 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -387,6 +387,15 @@ impl Connection { Ok(self.lock()?.in_txn) } + /// The current autocommit mode. + /// + /// Mirrors psycopg2's readable `connection.autocommit`; `prepare_database` + /// reads it to decide whether it needs to open a transaction explicitly. + #[getter] + fn autocommit(&self) -> PyResult { + Ok(self.lock()?.autocommit) + } + /// Switch autocommit mode on or off. /// /// In autocommit mode no implicit `BEGIN` is issued, so each statement runs diff --git a/rust/src/database/postgres/value.rs b/rust/src/database/postgres/value.rs index b725107a7e..7f27194a68 100644 --- a/rust/src/database/postgres/value.rs +++ b/rust/src/database/postgres/value.rs @@ -4,9 +4,9 @@ //! Kept in its own module so the cursor code stays focused on the DBAPI shape //! rather than the type-mapping table. //! -//! First cut: int / float / bool / str / bytes / None. Lists (for -//! `ANY($1)`-style queries) and richer types — json, decimal, timestamps — -//! are deferred to a follow-up. +//! First cut: int / float / bool / str / bytes / None, plus lists of those (for +//! `column = ANY($1)` / `!= ALL($1)` queries, which Synapse uses on Postgres). +//! Richer types — json, decimal, timestamps — are deferred to a follow-up. //! //! The mapping is column-type-driven on the way *out* (a single Python `int` //! becomes `INT2`/`INT4`/`INT8` depending on the column it is bound to) and @@ -20,22 +20,25 @@ //! | `Float` | `float` | `FLOAT4`, `FLOAT8` | //! | `Text` | `str` | `TEXT`, `VARCHAR`, `NAME`, `BPCHAR` | //! | `Bytea` | `bytes` | `BYTEA` | +//! | `Array` | `list` | any array of the above (e.g. `INT8[]`) | //! -//! Both directions share these type lists via the two `accepts` methods, which -//! must stay in sync with the `match` arms below. +//! Decoding (the way *in*) doesn't produce arrays — Synapse only binds them as +//! parameters — so only the `ToSql` side handles the `Array` variant. The scalar +//! type lists are shared via [`accepts_column_type`], kept in sync with the +//! `match` arms below. use std::error::Error; use bytes::BytesMut; use postgres_protocol::types::{ - bool_from_sql, bool_to_sql, bytea_to_sql, float4_from_sql, float4_to_sql, float8_from_sql, - float8_to_sql, int2_from_sql, int2_to_sql, int4_from_sql, int4_to_sql, int8_from_sql, - int8_to_sql, text_to_sql, + array_to_sql, bool_from_sql, bool_to_sql, bytea_to_sql, float4_from_sql, float4_to_sql, + float8_from_sql, float8_to_sql, int2_from_sql, int2_to_sql, int4_from_sql, int4_to_sql, + int8_from_sql, int8_to_sql, text_to_sql, ArrayDimension, }; use pyo3::exceptions::{PyTypeError, PyValueError}; -use pyo3::types::{PyBool, PyBytes, PyFloat, PyInt, PyString, PyTuple}; +use pyo3::types::{PyBool, PyBytes, PyFloat, PyInt, PyList, PyString, PyTuple}; use pyo3::{prelude::*, BoundObject}; -use tokio_postgres::types::{to_sql_checked, FromSql, IsNull, ToSql, Type, WrongType}; +use tokio_postgres::types::{to_sql_checked, FromSql, IsNull, Kind, ToSql, Type, WrongType}; use crate::database::value::{DbRow, DbValue}; @@ -49,6 +52,9 @@ pub enum PgValue { Float(f64), Text(Box), Bytea(Box<[u8]>), + /// A Python `list`, bound to a Postgres array column (e.g. for + /// `column = ANY($1)`). Each element is itself a [`PgValue`]. + Array(Vec), } impl PgValue { @@ -84,6 +90,15 @@ impl PgValue { if let Ok(b) = obj.cast::() { return Ok(PgValue::Bytea(b.as_bytes().into())); } + // A list is bound as a Postgres array (for `= ANY($1)` / `!= ALL($1)`). + // Each element is classified recursively. + if let Ok(list) = obj.cast::() { + let elements = list + .iter() + .map(|item| PgValue::from_py(&item)) + .collect::>>()?; + return Ok(PgValue::Array(elements)); + } Err(PyTypeError::new_err(format!( "unsupported parameter type for postgres: {}", obj.get_type().name()?, @@ -196,6 +211,43 @@ impl ToSql for PgValue { bytea_to_sql(v, buf); Ok(IsNull::No) } + (PgValue::Array(elements), ty) => { + // The column type is the array; its element type drives how each + // element is encoded. + let element_ty = match ty.kind() { + Kind::Array(element) => element, + _ => { + return Err( + format!("array parameter can't be bound to column type {ty}").into(), + ) + } + }; + + // An empty array has zero dimensions on the wire; a non-empty one + // is a single dimension with the SQL-standard lower bound of 1. + let dimensions = if elements.is_empty() { + Vec::new() + } else { + vec![ArrayDimension { + len: elements.len() as i32, + lower_bound: 1, + }] + }; + + array_to_sql( + dimensions, + element_ty.oid(), + elements.iter(), + // `array_to_sql`'s serializer uses `postgres_protocol`'s own + // `IsNull`, distinct from `tokio_postgres`'s; map between them. + |element, buf| match element.to_sql(element_ty, buf)? { + IsNull::No => Ok(postgres_protocol::IsNull::No), + IsNull::Yes => Ok(postgres_protocol::IsNull::Yes), + }, + buf, + )?; + Ok(IsNull::No) + } // If we get here then the caller has passed a value that doesn't // match the type of the column. (&PgValue::Bool(_), _) => Err(Box::new(WrongType::new::(ty.clone()))), @@ -207,7 +259,9 @@ impl ToSql for PgValue { } fn accepts(ty: &Type) -> bool { + // Scalars, plus arrays of a supported scalar element type. accepts_column_type(ty) + || matches!(ty.kind(), Kind::Array(element) if accepts_column_type(element)) } to_sql_checked!(); @@ -442,15 +496,70 @@ mod tests { fn from_py_rejects_unsupported_type() { Python::initialize(); Python::attach(|py| { - // A list is not a scalar we know how to bind. - let list = pyo3::types::PyList::new(py, [1, 2, 3]).unwrap(); - let err = PgValue::from_py(&list.into_any()).unwrap_err(); + // A dict is not something we know how to bind. + let dict = pyo3::types::PyDict::new(py); + let err = PgValue::from_py(&dict.into_any()).unwrap_err(); assert!(err.is_instance_of::(py)); // The message names the offending type, which is the useful part. - assert!(err.to_string().contains("list"), "got: {err}"); + assert!(err.to_string().contains("dict"), "got: {err}"); }); } + #[test] + fn from_py_classifies_list_as_array() { + Python::initialize(); + Python::attach(|py| { + let list = PyList::new(py, [1i64, 2, 3]).unwrap(); + match PgValue::from_py(&list.into_any()).unwrap() { + PgValue::Array(elements) => assert!(matches!( + elements.as_slice(), + [PgValue::Int(1), PgValue::Int(2), PgValue::Int(3)] + )), + other => panic!("expected Array, got {other:?}"), + } + + // Element types are classified individually (here, strings). + let list = PyList::new(py, ["a", "b"]).unwrap(); + match PgValue::from_py(&list.into_any()).unwrap() { + PgValue::Array(elements) => assert_eq!(elements.len(), 2), + other => panic!("expected Array, got {other:?}"), + } + }); + } + + #[test] + fn to_sql_encodes_arrays() { + // An `INT8[]` array encodes without error and produces a non-empty + // buffer; the element type comes from the array column's element type. + let array = PgValue::Array(vec![PgValue::Int(1), PgValue::Int(2)]); + let (bytes, is_null) = encode(&array, &Type::INT8_ARRAY); + assert!(!is_null); + assert!(!bytes.is_empty()); + + // A `TEXT[]` array likewise. + let array = PgValue::Array(vec![PgValue::Text("x".into())]); + assert!(encode_result(&array, &Type::TEXT_ARRAY).is_ok()); + + // An empty array is valid (zero dimensions). + assert!(encode_result(&PgValue::Array(vec![]), &Type::INT8_ARRAY).is_ok()); + + // Binding an array to a non-array column is an error. + assert!(encode_result(&PgValue::Array(vec![PgValue::Int(1)]), &Type::INT8).is_err()); + + // An element whose type doesn't match the array's element type errors. + let bad = PgValue::Array(vec![PgValue::Text("x".into())]); + assert!(encode_result(&bad, &Type::INT8_ARRAY).is_err()); + } + + #[test] + fn accepts_arrays_of_supported_elements() { + assert!(::accepts(&Type::INT8_ARRAY)); + assert!(::accepts(&Type::TEXT_ARRAY)); + assert!(::accepts(&Type::BOOL_ARRAY)); + // An array of an unsupported element type is rejected. + assert!(!::accepts(&Type::JSON_ARRAY)); + } + #[test] fn to_sql_encodes_each_type_for_its_column() { // NULL is encoded as "no bytes, IsNull::Yes" regardless of column type. diff --git a/synapse/storage/database.py b/synapse/storage/database.py index 3ad3a612b3..1dee889067 100644 --- a/synapse/storage/database.py +++ b/synapse/storage/database.py @@ -560,7 +560,13 @@ class LoggingTransaction: self._do_execute(self.txn.executemany, sql, *args) def executescript(self, sql: str) -> None: - if isinstance(self.database_engine, Sqlite3Engine): + # Both the sqlite driver and the Rust Postgres shim expose a + # multi-statement `executescript` on the cursor; psycopg2 does not (its + # engine runs scripts via a plain `execute`). + if isinstance(self.database_engine, Sqlite3Engine) or ( + isinstance(self.database_engine, PostgresEngine) + and not self.database_engine.uses_psycopg2_extras + ): self._do_execute(self.txn.executescript, sql) # type: ignore[attr-defined] else: raise NotImplementedError( diff --git a/synapse/storage/rust_dbapi.py b/synapse/storage/rust_dbapi.py index 78486cea39..2766c41863 100644 --- a/synapse/storage/rust_dbapi.py +++ b/synapse/storage/rust_dbapi.py @@ -209,6 +209,12 @@ class Cursor: self._exhausted = False self._cursor.executemany(sql, [list(p) for p in seq_of_parameters]) + def executescript(self, script: str) -> None: + # The shim runs a multi-statement script on the simple-query protocol + # (no parameters, no fetchable rows). + self._exhausted = False + self._cursor.executescript(script) + def _next(self) -> Any: """Fetch one row, or `None` once the result set is exhausted.""" if self._exhausted: @@ -294,6 +300,10 @@ class Connection: def set_autocommit(self, autocommit: bool) -> None: self._conn.set_autocommit(autocommit) + @property + def autocommit(self) -> bool: + return bool(self._conn.autocommit) + def is_closed(self) -> bool: return bool(self._conn.is_closed()) diff --git a/tests/storage/test_rust_dbapi.py b/tests/storage/test_rust_dbapi.py index fc72455aa5..ccf564fe1b 100644 --- a/tests/storage/test_rust_dbapi.py +++ b/tests/storage/test_rust_dbapi.py @@ -253,3 +253,29 @@ class RustDBAPIAdapterTestCase(unittest.TestCase): self.assertEqual(txn.description[0][0], "only") db_conn.commit() + + def test_array_parameter(self) -> None: + # A list parameter binds as a Postgres array, for `= ANY($1)` queries. + cursor = self.conn.cursor() + cursor.execute("SELECT 5 = ANY($1::int[])", ([1, 5, 9],)) + self.assertEqual(cursor.fetchone(), (True,)) + cursor.execute("SELECT 7 = ANY($1::int[])", ([1, 5, 9],)) + self.assertEqual(cursor.fetchone(), (False,)) + self.conn.commit() + + def test_executescript(self) -> None: + # A multi-statement script runs via the shim's simple-query path. + cursor = self.conn.cursor() + cursor.executescript( + "CREATE TEMP TABLE s (a int); INSERT INTO s VALUES (1), (2);" + ) + cursor.execute("SELECT count(*) FROM s") + self.assertEqual(cursor.fetchone(), (2,)) + self.conn.commit() + + def test_autocommit_property(self) -> None: + self.assertFalse(self.conn.autocommit) + self.conn.set_autocommit(True) + self.assertTrue(self.conn.autocommit) + self.conn.set_autocommit(False) + self.assertFalse(self.conn.autocommit)