From f5768de6908af3dd980f0c6598e192fb1e2b299f Mon Sep 17 00:00:00 2001 From: Erik Johnston Date: Wed, 1 Jul 2026 15:19:37 +0000 Subject: [PATCH] Add Cursor.executescript for multi-statement SQL MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Synapse's schema setup (`prepare_database.py`) runs `;`-separated SQL scripts via the engine's `executescript`. `Cursor.execute` can't serve those: it `prepare`s the query, and Postgres rejects multiple commands in a prepared statement. Add `Cursor.executescript`, which runs the whole script on the simple-query protocol (`batch_execute`), which does allow multiple statements. It takes no parameters and produces no fetchable rows. The script runs inside the connection's *current* transaction (opened lazily like `execute`) and is left open for the caller to commit. It deliberately does NOT reproduce the commit-any-pending-transaction-first behaviour of `sqlite3.executescript` (which psycopg2's engine mirrors with a leading `COMMIT`). That forced commit actually undercuts the atomicity `prepare_database` sets out to get — it opens a transaction so "upgrades are either applied completely, or not at all", but the first script's implicit commit ends it. Running the script within the ongoing transaction instead lets successive scripts accumulate and be committed once, which is both simpler and more correct. The only engine-level piece left to layer on top is the auto-increment placeholder substitution. Co-Authored-By: Claude Opus 4.8 (1M context) --- rust/src/database/postgres/connection.rs | 26 ++++++++ tests/synapse_rust/test_database_postgres.py | 65 ++++++++++++++++++++ 2 files changed, 91 insertions(+) diff --git a/rust/src/database/postgres/connection.rs b/rust/src/database/postgres/connection.rs index 6ad158de4b..b21af633d6 100644 --- a/rust/src/database/postgres/connection.rs +++ b/rust/src/database/postgres/connection.rs @@ -389,6 +389,32 @@ impl Cursor { Ok(()) } + /// Execute a multi-statement SQL script (statements separated by `;`). + /// + /// Unlike [`Cursor::execute`], which `prepare`s a single statement, this + /// runs the whole script on the simple-query protocol (`batch_execute`), so + /// it may contain many `;`-separated statements — as Synapse's schema files + /// do. It takes no parameters and produces no fetchable rows. + /// + /// This is a thin primitive: the script runs inside the connection's current + /// transaction, opening one lazily like `execute` and leaving it open for + /// the caller to commit. The higher-level engine `executescript` — which + /// also substitutes the auto-increment placeholder — is layered on top. + /// Note it does *not* commit any prior transaction first: unlike the + /// psycopg2 engine (which still prefixes `COMMIT; BEGIN TRANSACTION;`, + /// committing script-by-script), the Rust engine deliberately keeps a whole + /// sequence of schema/delta scripts in one transaction so it is applied + /// either completely or not at all — see + /// `BaseDatabaseEngine.executescript`'s docstring for the contract. + fn executescript(&self, py: Python<'_>, script: &str) -> PyResult<()> { + // A script yields no fetchable rows, so drop any previous result set. + self.lock_state()?.new_query(); + + self.connection.with_client(py, |client| { + client.batch_execute(script).block_on_result(py) + }) + } + /// Return the next row of the current result set, or `None` if exhausted. fn fetch_one<'py>(&self, py: Python<'py>) -> PyResult>> { self.lock_state()?.fetch_one(py) diff --git a/tests/synapse_rust/test_database_postgres.py b/tests/synapse_rust/test_database_postgres.py index ce8daea07f..492e04e1d8 100644 --- a/tests/synapse_rust/test_database_postgres.py +++ b/tests/synapse_rust/test_database_postgres.py @@ -816,6 +816,71 @@ class PostgresConnectionDrivenTestCase(unittest.TestCase): self.conn.set_autocommit(True) self.conn.rollback() + # -- executescript (multi-statement) ------------------------------------ + + def test_executescript_runs_all_statements(self) -> None: + """A `;`-separated script runs every statement (unlike `execute`, which + prepares a single one).""" + table = "rust_pg_script" + try: + cursor = self.conn.cursor() + cursor.executescript( + f"CREATE TABLE {table} (id int); " + f"INSERT INTO {table} VALUES (1); " + f"INSERT INTO {table} VALUES (2), (3);" + ) + self.conn.commit() + + read = self.conn.cursor() + read.execute(f"SELECT id FROM {table} ORDER BY id") + self.assertEqual(read.fetch_all(), [(1,), (2,), (3,)]) + self.conn.commit() + finally: + self._exec_commit(f"DROP TABLE IF EXISTS {table}") + + def test_executescript_leaves_transaction_open_for_caller(self) -> None: + """The script runs in the connection's transaction, left open — so a + following `rollback()` undoes it (it was not autocommitted).""" + table = "rust_pg_script_open" + try: + self.conn.cursor().executescript(f"CREATE TABLE {table} (id int);") + self.conn.rollback() + self.assertFalse(self._table_exists(table)) + self.conn.commit() + finally: + self._exec_commit(f"DROP TABLE IF EXISTS {table}") + + def test_executescript_error_surfaces_as_database_error(self) -> None: + """A failing statement mid-script surfaces via the exception hierarchy; + the aborted transaction rolls back cleanly.""" + cursor = self.conn.cursor() + with self.assertRaises(postgres.DatabaseError): + cursor.executescript("CREATE TABLE rust_pg_script_bad (id int); NOT SQL;") + self.conn.rollback() + # The CREATE was in the same aborted, rolled-back transaction, so it + # left nothing behind. + self.assertFalse(self._table_exists("rust_pg_script_bad")) + self.conn.commit() + + def test_successive_scripts_share_one_transaction(self) -> None: + """Successive `executescript` calls accumulate in the same open + transaction -- there is no implicit commit between them (unlike + `sqlite3.executescript`) -- so a single rollback discards them all. + This is the atomicity `prepare_database` relies on.""" + try: + cursor = self.conn.cursor() + cursor.executescript("CREATE TABLE rust_pg_script_a (id int);") + cursor.executescript("CREATE TABLE rust_pg_script_b (id int);") + # Nothing was committed between the two calls, so one rollback + # undoes both. + self.conn.rollback() + self.assertFalse(self._table_exists("rust_pg_script_a")) + self.assertFalse(self._table_exists("rust_pg_script_b")) + self.conn.commit() + finally: + self._exec_commit("DROP TABLE IF EXISTS rust_pg_script_a") + self._exec_commit("DROP TABLE IF EXISTS rust_pg_script_b") + @unittest.skip_unless( bool(USE_POSTGRES_FOR_TESTS), "requires a Postgres server (set SYNAPSE_POSTGRES)"