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)"