From e165762bb7c9f1eda2a0532baea1ddb6feba7313 Mon Sep 17 00:00:00 2001 From: Hans Yu Date: Tue, 14 Jul 2026 07:45:40 +0200 Subject: [PATCH] refactor: begin transaction to automatically commit (#41979) Co-authored-by: Evan Rusackas --- pytest.ini | 2 + superset/extensions/metadb.py | 33 ++--- .../unit_tests/db_engine_specs/test_sqlite.py | 9 +- .../unit_tests/extensions/test_sqlalchemy.py | 113 ++++++++++-------- .../test_add_deleted_at_to_dashboards.py | 14 +-- .../test_add_deleted_at_to_slices.py | 10 +- .../test_add_deleted_at_to_tables.py | 10 +- 7 files changed, 102 insertions(+), 89 deletions(-) diff --git a/pytest.ini b/pytest.ini index d2dba1e2707..64de5b177c1 100644 --- a/pytest.ini +++ b/pytest.ini @@ -36,6 +36,8 @@ filterwarnings = # error:"TableColumn" object is being merged into a Session:sqlalchemy.exc.RemovedIn20Warning # error:"TaggedObject" object is being merged into a Session:sqlalchemy.exc.RemovedIn20Warning # error:The autoload parameter is deprecated:sqlalchemy.exc.RemovedIn20Warning +# error:The connection.execute\(\) method:sqlalchemy.exc.RemovedIn20Warning + error:The current statement is being autocommitted using implicit autocommit:sqlalchemy.exc.RemovedIn20Warning error:The connection.execute\(\) method:sqlalchemy.exc.RemovedIn20Warning # error:The current statement is being autocommitted using implicit autocommit:sqlalchemy.exc.RemovedIn20Warning # error:The ``declarative_base\(\)`` function is now available:sqlalchemy.exc.RemovedIn20Warning diff --git a/superset/extensions/metadb.py b/superset/extensions/metadb.py index 99a3e82a6fb..9fb807a88c1 100644 --- a/superset/extensions/metadb.py +++ b/superset/extensions/metadb.py @@ -417,12 +417,12 @@ class SupersetShillelaghAdapter(Adapter): query = self._build_sql(bounds, order, limit, offset) with self.engine_context() as engine: - connection = engine.connect() - rows = connection.execute(query) - for i, row in enumerate(rows): - data = dict(zip(self.columns, row, strict=False)) - data["rowid"] = data[self._rowid] if self._rowid else i - yield data + with engine.connect() as connection: + rows = connection.execute(query) + for i, row in enumerate(rows): + data = dict(zip(self.columns, row, strict=False)) + data["rowid"] = data[self._rowid] if self._rowid else i + yield data @check_dml def insert_row(self, row: Row) -> int: @@ -445,15 +445,16 @@ class SupersetShillelaghAdapter(Adapter): query = self._table.insert().values(**row) with self.engine_context() as engine: - connection = engine.connect() - result = connection.execute(query) + with engine.begin() as connection: + result = connection.execute(query) - # return rowid - if self._rowid: - return result.inserted_primary_key[0] + # return rowid + if self._rowid: + return result.inserted_primary_key[0] query = select(func.count()).select_from(self._table) - return connection.execute(query).scalar() + with engine.connect() as connection: + return connection.execute(query).scalar() @check_dml @has_rowid @@ -464,8 +465,8 @@ class SupersetShillelaghAdapter(Adapter): query = self._table.delete().where(self._table.c[self._rowid] == row_id) with self.engine_context() as engine: - connection = engine.connect() - connection.execute(query) + with engine.begin() as connection: + connection.execute(query) @check_dml @has_rowid @@ -488,5 +489,5 @@ class SupersetShillelaghAdapter(Adapter): ) with self.engine_context() as engine: - connection = engine.connect() - connection.execute(query) + with engine.begin() as connection: + connection.execute(query) diff --git a/tests/unit_tests/db_engine_specs/test_sqlite.py b/tests/unit_tests/db_engine_specs/test_sqlite.py index 7efb84e81b9..677308a7701 100644 --- a/tests/unit_tests/db_engine_specs/test_sqlite.py +++ b/tests/unit_tests/db_engine_specs/test_sqlite.py @@ -121,12 +121,13 @@ def test_time_grain_expressions(dttm: str, grain: str, expected: str) -> None: from superset.db_engine_specs.sqlite import SqliteEngineSpec engine = create_engine("sqlite://") - connection = engine.connect() - connection.execute(text("CREATE TABLE t (dttm DATETIME)")) - connection.execute(text("INSERT INTO t VALUES (:dttm)"), {"dttm": dttm}) + with engine.begin() as connection: + connection.execute(text("CREATE TABLE t (dttm DATETIME)")) + connection.execute(text("INSERT INTO t VALUES (:dttm)"), {"dttm": dttm}) # pylint: disable=protected-access expression = SqliteEngineSpec._time_grain_expressions[grain].format(col="dttm") sql = f"SELECT {expression} FROM t" # noqa: S608 - result = connection.execute(text(sql)).scalar() + with engine.connect() as connection: + result = connection.execute(text(sql)).scalar() assert result == expected diff --git a/tests/unit_tests/extensions/test_sqlalchemy.py b/tests/unit_tests/extensions/test_sqlalchemy.py index 4778813f98d..1cfb539bc2a 100644 --- a/tests/unit_tests/extensions/test_sqlalchemy.py +++ b/tests/unit_tests/extensions/test_sqlalchemy.py @@ -63,16 +63,17 @@ def database1(session: Session) -> Iterator["Database"]: @pytest.fixture def table1(session: Session, database1: "Database") -> Iterator[None]: with database1.get_sqla_engine() as engine: - conn = engine.connect() - conn.execute( - text("CREATE TABLE table1 (a INTEGER NOT NULL PRIMARY KEY, b INTEGER)") - ) - conn.execute(text("INSERT INTO table1 (a, b) VALUES (1, 10), (2, 20)")) + with engine.begin() as conn: + conn.execute( + text("CREATE TABLE table1 (a INTEGER NOT NULL PRIMARY KEY, b INTEGER)") + ) + conn.execute(text("INSERT INTO table1 (a, b) VALUES (1, 10), (2, 20)")) db.session.commit() yield - conn.execute(text("DROP TABLE table1")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE table1")) db.session.commit() @@ -99,16 +100,19 @@ def database2(session: Session) -> Iterator["Database"]: @pytest.fixture def table2(session: Session, database2: "Database") -> Iterator[None]: with database2.get_sqla_engine() as engine: - conn = engine.connect() - conn.execute( - text("CREATE TABLE table2 (a INTEGER NOT NULL PRIMARY KEY, b TEXT)") - ) - conn.execute(text("INSERT INTO table2 (a, b) VALUES (1, 'ten'), (2, 'twenty')")) + with engine.begin() as conn: + conn.execute( + text("CREATE TABLE table2 (a INTEGER NOT NULL PRIMARY KEY, b TEXT)") + ) + conn.execute( + text("INSERT INTO table2 (a, b) VALUES (1, 'ten'), (2, 'twenty')") + ) db.session.commit() yield - conn.execute(text("DROP TABLE table2")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE table2")) db.session.commit() @@ -137,9 +141,9 @@ def test_superset(mocker: MockerFixture, app_context: None, table1: None) -> Non # Skip test if superset:// dialect can't be loaded (common in Docker) pytest.skip(f"Superset dialect not available: {e}") - conn = engine.connect() - results = conn.execute(text('SELECT * FROM "database1.table1"')) - assert list(results) == [(1, 10), (2, 20)] + with engine.connect() as conn: + results = conn.execute(text('SELECT * FROM "database1.table1"')) + assert list(results) == [(1, 10), (2, 20)] @with_config( @@ -178,9 +182,9 @@ def test_superset_limit(mocker: MockerFixture, app_context: None, table1: None) # Skip test if superset:// dialect can't be loaded (common in Docker) pytest.skip(f"Superset dialect not available: {e}") - conn = engine.connect() - results = conn.execute(text('SELECT * FROM "database1.table1"')) - assert list(results) == [(1, 10)] + with engine.connect() as conn: + results = conn.execute(text('SELECT * FROM "database1.table1"')) + assert list(results) == [(1, 10)] @with_feature_flags(ENABLE_SUPERSET_META_DB=True) @@ -213,16 +217,16 @@ def test_superset_joins( # Skip test if superset:// dialect can't be loaded (common in Docker) pytest.skip(f"Superset dialect not available: {e}") - conn = engine.connect() - results = conn.execute( - text(""" - SELECT t1.b, t2.b - FROM "database1.table1" AS t1 - JOIN "database2.table2" AS t2 - ON t1.a = t2.a - """) - ) - assert list(results) == [(10, "ten"), (20, "twenty")] + with engine.connect() as conn: + results = conn.execute( + text(""" + SELECT t1.b, t2.b + FROM "database1.table1" AS t1 + JOIN "database2.table2" AS t2 + ON t1.a = t2.a + """) + ) + assert list(results) == [(10, "ten"), (20, "twenty")] @with_feature_flags(ENABLE_SUPERSET_META_DB=True) @@ -257,22 +261,27 @@ def test_dml( # Skip test if superset:// dialect can't be loaded (common in Docker) pytest.skip(f"Superset dialect not available: {e}") - conn = engine.connect() + with engine.begin() as conn: + conn.execute(text('INSERT INTO "database1.table1" (a, b) VALUES (3, 30)')) + with engine.connect() as conn: + results = conn.execute(text('SELECT * FROM "database1.table1"')) + assert list(results) == [(1, 10), (2, 20), (3, 30)] + with engine.begin() as conn: + conn.execute(text('UPDATE "database1.table1" SET b=35 WHERE a=3')) + with engine.connect() as conn: + results = conn.execute(text('SELECT * FROM "database1.table1"')) + assert list(results) == [(1, 10), (2, 20), (3, 35)] + with engine.begin() as conn: + conn.execute(text('DELETE FROM "database1.table1" WHERE b>20')) + with engine.connect() as conn: + results = conn.execute(text('SELECT * FROM "database1.table1"')) + assert list(results) == [(1, 10), (2, 20)] - conn.execute(text('INSERT INTO "database1.table1" (a, b) VALUES (3, 30)')) - results = conn.execute(text('SELECT * FROM "database1.table1"')) - assert list(results) == [(1, 10), (2, 20), (3, 30)] - conn.execute(text('UPDATE "database1.table1" SET b=35 WHERE a=3')) - results = conn.execute(text('SELECT * FROM "database1.table1"')) - assert list(results) == [(1, 10), (2, 20), (3, 35)] - conn.execute(text('DELETE FROM "database1.table1" WHERE b>20')) - results = conn.execute(text('SELECT * FROM "database1.table1"')) - assert list(results) == [(1, 10), (2, 20)] - - with pytest.raises(ProgrammingError) as excinfo: - conn.execute( - text("""INSERT INTO "database2.table2" (a, b) VALUES (3, 'thirty')""") - ) + with engine.begin() as conn: + with pytest.raises(ProgrammingError) as excinfo: + conn.execute( + text("""INSERT INTO "database2.table2" (a, b) VALUES (3, 'thirty')""") + ) assert str(excinfo.value).strip() == ( "(shillelagh.exceptions.ProgrammingError) DML not enabled in database " '"database2"\n[SQL: INSERT INTO "database2.table2" (a, b) ' @@ -324,9 +333,9 @@ def test_security_manager( # Skip test if superset:// dialect can't be loaded (common in Docker) pytest.skip(f"Superset dialect not available: {e}") - conn = engine.connect() - with pytest.raises(SupersetSecurityException) as excinfo: - conn.execute(text('SELECT * FROM "database1.table1"')) + with engine.connect() as conn: + with pytest.raises(SupersetSecurityException) as excinfo: + conn.execute(text('SELECT * FROM "database1.table1"')) assert str(excinfo.value) == ( "You need access to the following tables: `table1`,\n " "`all_database_access` or `all_datasource_access` permission" @@ -358,13 +367,13 @@ def test_allowed_dbs(mocker: MockerFixture, app_context: None, table1: None) -> # Skip test if superset:// dialect can't be loaded (common in Docker) pytest.skip(f"Superset dialect not available: {e}") - conn = engine.connect() + with engine.connect() as conn: + results = conn.execute(text('SELECT * FROM "database1.table1"')) + assert list(results) == [(1, 10), (2, 20)] - results = conn.execute(text('SELECT * FROM "database1.table1"')) - assert list(results) == [(1, 10), (2, 20)] - - with pytest.raises(ProgrammingError) as excinfo: - conn.execute(text('SELECT * FROM "database2.table2"')) + with engine.connect() as conn: + with pytest.raises(ProgrammingError) as excinfo: + conn.execute(text('SELECT * FROM "database2.table2"')) assert str(excinfo.value) == ( """ (shillelagh.exceptions.ProgrammingError) Unsupported table: database2.table2 diff --git a/tests/unit_tests/migrations/test_add_deleted_at_to_dashboards.py b/tests/unit_tests/migrations/test_add_deleted_at_to_dashboards.py index 6bcc60f7956..66dae169177 100644 --- a/tests/unit_tests/migrations/test_add_deleted_at_to_dashboards.py +++ b/tests/unit_tests/migrations/test_add_deleted_at_to_dashboards.py @@ -103,7 +103,7 @@ def _indexes(engine: Engine) -> set[str]: def test_upgrade_adds_deleted_at_column_and_index(engine: Engine) -> None: - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -115,7 +115,7 @@ def test_upgrade_adds_deleted_at_column_and_index(engine: Engine) -> None: def test_downgrade_drops_deleted_at_column_and_index(engine: Engine) -> None: - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -138,7 +138,7 @@ def test_downgrade_preserves_soft_deleted_row_data(engine: Engine) -> None: hard-delete, restore, or rename BEFORE downgrading. See the "Rollback note" in ``UPDATING.md``. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -198,7 +198,7 @@ def test_sqlite_slug_constraint_swap_is_a_no_op(engine: Engine) -> None: "pre-condition: legacy slug constraint must exist before the migration runs" ) - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -207,7 +207,7 @@ def test_sqlite_slug_constraint_swap_is_a_no_op(engine: Engine) -> None: "upgrade() must keep the legacy slug constraint on SQLite (no-op branch)" ) - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.downgrade() @@ -222,7 +222,7 @@ def test_upgrade_is_idempotent(engine: Engine) -> None: idempotent skip-if-exists; running ``upgrade()`` twice must not raise. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -233,7 +233,7 @@ def test_downgrade_is_idempotent(engine: Engine) -> None: """``drop_columns`` / ``drop_index`` are skip-if-not-exists; running ``downgrade()`` twice must not raise. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() diff --git a/tests/unit_tests/migrations/test_add_deleted_at_to_slices.py b/tests/unit_tests/migrations/test_add_deleted_at_to_slices.py index 7a46c4da4cc..3510525dc6e 100644 --- a/tests/unit_tests/migrations/test_add_deleted_at_to_slices.py +++ b/tests/unit_tests/migrations/test_add_deleted_at_to_slices.py @@ -88,7 +88,7 @@ def _indexes(engine: Engine) -> set[str]: def test_upgrade_adds_deleted_at_column_and_index(engine: Engine) -> None: """upgrade() adds the nullable ``deleted_at`` column and its index.""" - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -101,7 +101,7 @@ def test_upgrade_adds_deleted_at_column_and_index(engine: Engine) -> None: def test_downgrade_drops_deleted_at_column_and_index(engine: Engine) -> None: """downgrade() removes both schema artifacts (column and index).""" - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -124,7 +124,7 @@ def test_downgrade_preserves_soft_deleted_row_data(engine: Engine) -> None: hard-delete, restore, or rename BEFORE downgrading. See the "Rollback note" in ``UPDATING.md``. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -171,7 +171,7 @@ def test_upgrade_is_idempotent(engine: Engine) -> None: idempotent skip-if-exists; running ``upgrade()`` twice must not raise. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -182,7 +182,7 @@ def test_downgrade_is_idempotent(engine: Engine) -> None: """``drop_columns`` / ``drop_index`` are skip-if-not-exists; running ``downgrade()`` twice must not raise. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() diff --git a/tests/unit_tests/migrations/test_add_deleted_at_to_tables.py b/tests/unit_tests/migrations/test_add_deleted_at_to_tables.py index 3ddeba215f7..61719ea2976 100644 --- a/tests/unit_tests/migrations/test_add_deleted_at_to_tables.py +++ b/tests/unit_tests/migrations/test_add_deleted_at_to_tables.py @@ -85,7 +85,7 @@ def _indexes(engine: Engine) -> set[str]: def test_upgrade_adds_deleted_at_column_and_index(engine: Engine) -> None: - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -97,7 +97,7 @@ def test_upgrade_adds_deleted_at_column_and_index(engine: Engine) -> None: def test_downgrade_drops_deleted_at_column_and_index(engine: Engine) -> None: - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -120,7 +120,7 @@ def test_downgrade_preserves_soft_deleted_row_data(engine: Engine) -> None: hard-delete, restore, or rename BEFORE downgrading. See the "Rollback note" in ``UPDATING.md``. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -170,7 +170,7 @@ def test_upgrade_is_idempotent(engine: Engine) -> None: idempotent skip-if-exists; running ``upgrade()`` twice must not raise. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade() @@ -181,7 +181,7 @@ def test_downgrade_is_idempotent(engine: Engine) -> None: """``drop_columns`` / ``drop_index`` are skip-if-not-exists; running ``downgrade()`` twice must not raise. """ - with engine.connect() as conn: + with engine.begin() as conn: ctx = MigrationContext.configure(conn) with Operations.context(ctx): migration.upgrade()