diff --git a/pytest.ini b/pytest.ini index cc1cf882b39..573a81fa66b 100644 --- a/pytest.ini +++ b/pytest.ini @@ -41,7 +41,7 @@ filterwarnings = # error:The current statement is being autocommitted using implicit autocommit:sqlalchemy.exc.RemovedIn20Warning # error:The `database` package is deprecated:sqlalchemy.exc.RemovedIn20Warning # error:The ``declarative_base\(\)`` function is now available:sqlalchemy.exc.RemovedIn20Warning -# error:The Engine.execute\(\) method is considered legacy:sqlalchemy.exc.RemovedIn20Warning + error:The Engine.execute\(\) method is considered legacy:sqlalchemy.exc.RemovedIn20Warning error:The legacy calling style of select\(\) is deprecated:sqlalchemy.exc.RemovedIn20Warning error:The "whens" argument to case:sqlalchemy.exc.RemovedIn20Warning # error:"User" object is being merged into a Session:sqlalchemy.exc.RemovedIn20Warning diff --git a/scripts/benchmark_migration.py b/scripts/benchmark_migration.py index f2b9e258cdc..550d3fda418 100644 --- a/scripts/benchmark_migration.py +++ b/scripts/benchmark_migration.py @@ -153,10 +153,11 @@ def main( # noqa: C901 ) print(f"Migration goes from {down_revision} to {revision}") - current_revision = db.engine.execute( - text("SELECT version_num FROM alembic_version") - ).scalar() - print(f"Current version of the DB is {current_revision}") + with db.engine.connect() as conn: + current_revision = conn.execute( + text("SELECT version_num FROM alembic_version") + ).scalar() + print(f"Current version of the DB is {current_revision}") if current_revision != down_revision: if not force: diff --git a/superset/cli/test_db.py b/superset/cli/test_db.py index 3eaecb803c1..124efdb7eb3 100644 --- a/superset/cli/test_db.py +++ b/superset/cli/test_db.py @@ -120,14 +120,16 @@ def test_datetime(console: Console, engine: Engine) -> None: now = datetime.now() console.print("Inserting timestamp value...") - insert_stmt = insert(table).values(ts=now) - engine.execute(insert_stmt) + with engine.begin() as conn: + insert_stmt = insert(table).values(ts=now) + conn.execute(insert_stmt) console.print("Reading timestamp value...") - select_stmt = select(table) - row = engine.execute(select_stmt).fetchone() - assert row[0] == now - console.print(":thumbs_up: [green]Success!") + with engine.connect() as conn: + select_stmt = select(table) + row = conn.execute(select_stmt).fetchone() + assert row[0] == now + console.print(":thumbs_up: [green]Success!") @click.command() diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py index db3d382f2ad..d4b1f23444c 100644 --- a/superset/db_engine_specs/databricks.py +++ b/superset/db_engine_specs/databricks.py @@ -737,11 +737,14 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): return default_catalog with database.get_sqla_engine() as engine: - catalogs = {catalog for (catalog,) in engine.execute(text("SHOW CATALOGS"))} - if len(catalogs) == 1: - return catalogs.pop() + with engine.connect() as conn: + catalogs = { + catalog for (catalog,) in conn.execute(text("SHOW CATALOGS")) + } + if len(catalogs) == 1: + return catalogs.pop() - return engine.execute(text("SELECT current_catalog()")).scalar() + return conn.execute(text("SELECT current_catalog()")).scalar() @classmethod def get_prequeries( @@ -765,7 +768,8 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): database: Database, inspector: Inspector, ) -> set[str]: - return {catalog for (catalog,) in inspector.bind.execute(text("SHOW CATALOGS"))} + with inspector.engine.connect() as conn: + return {catalog for (catalog,) in conn.execute(text("SHOW CATALOGS"))} class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): @@ -939,7 +943,8 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): database: Database, inspector: Inspector, ) -> set[str]: - return {catalog for (catalog,) in inspector.bind.execute(text("SHOW CATALOGS"))} + with inspector.engine.connect() as conn: + return {catalog for (catalog,) in conn.execute(text("SHOW CATALOGS"))} @classmethod def adjust_engine_params( diff --git a/superset/db_engine_specs/doris.py b/superset/db_engine_specs/doris.py index 27a69bdf183..6aeb3d288a8 100644 --- a/superset/db_engine_specs/doris.py +++ b/superset/db_engine_specs/doris.py @@ -307,9 +307,10 @@ class DorisEngineSpec(MySQLEngineSpec): # if not, iterate over existing catalogs and find the current one with database.get_sqla_engine() as engine: - for catalog in engine.execute("SHOW CATALOGS"): - if catalog.IsCurrent: - return catalog.CatalogName + with engine.connect() as conn: + for catalog in conn.execute(text("SHOW CATALOGS")): + if catalog.IsCurrent: + return catalog.CatalogName # fallback to "internal" return DEFAULT_CATALOG @@ -326,8 +327,9 @@ class DorisEngineSpec(MySQLEngineSpec): CatalogId, CatalogName, Type, IsCurrent, CreateTime, LastUpdateTime, Comment We need to extract just the CatalogName column. """ - result = inspector.bind.execute(text("SHOW CATALOGS")) - return {row.CatalogName for row in result} + with inspector.engine.connect() as conn: + result = conn.execute(text("SHOW CATALOGS")) + return {row.CatalogName for row in result} @classmethod def get_schema_from_engine_params( diff --git a/superset/db_engine_specs/duckdb.py b/superset/db_engine_specs/duckdb.py index b9b2616c1c9..db2b21a43dd 100644 --- a/superset/db_engine_specs/duckdb.py +++ b/superset/db_engine_specs/duckdb.py @@ -470,9 +470,10 @@ class MotherDuckEngineSpec(DuckDBEngineSpec): database: Database, inspector: Inspector, ) -> set[str]: - return { - catalog - for (catalog,) in inspector.bind.execute( - text("SELECT alias FROM MD_ALL_DATABASES() WHERE is_attached;") - ) - } + with inspector.engine.connect() as conn: + return { + catalog + for (catalog,) in conn.execute( + text("SELECT alias FROM MD_ALL_DATABASES() WHERE is_attached;") + ) + } diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index 3d89e122314..d17d57ca42f 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -234,7 +234,8 @@ class HiveEngineSpec(PrestoEngineSpec): catalog=table.catalog, schema=table.schema, ) as engine: - engine.execute(text(f"DROP TABLE IF EXISTS {str(table)}")) + with engine.begin() as conn: + conn.execute(text(f"DROP TABLE IF EXISTS {str(table)}")) def _get_hive_type(dtype: np.dtype[Any]) -> str: hive_type_by_dtype = { @@ -260,22 +261,25 @@ class HiveEngineSpec(PrestoEngineSpec): catalog=table.catalog, schema=table.schema, ) as engine: - engine.execute( - text( - f""" - CREATE TABLE {str(table)} ({schema_definition}) - STORED AS PARQUET - LOCATION :location - """ - ), - location=upload_to_s3( - filename=file.name, - upload_prefix=app.config["CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC"]( - database, g.user, table.schema + with engine.begin() as conn: + conn.execute( + text( + f""" + CREATE TABLE {str(table)} ({schema_definition}) + STORED AS PARQUET + LOCATION :location + """ ), - table=table, - ), - ) + { + "location": upload_to_s3( + filename=file.name, + upload_prefix=app.config[ + "CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC" + ](database, g.user, table.schema), + table=table, + ), + }, + ) @classmethod def convert_dttm( diff --git a/superset/db_engine_specs/impala.py b/superset/db_engine_specs/impala.py index 086839a0b77..b73a169d260 100644 --- a/superset/db_engine_specs/impala.py +++ b/superset/db_engine_specs/impala.py @@ -96,11 +96,12 @@ class ImpalaEngineSpec(BaseEngineSpec): @classmethod def get_schema_names(cls, inspector: Inspector) -> set[str]: - return { - row[0] - for row in inspector.engine.execute(text("SHOW SCHEMAS")) - if not row[0].startswith("_") - } + with inspector.engine.connect() as conn: + return { + row[0] + for row in conn.execute(text("SHOW SCHEMAS")) + if not row[0].startswith("_") + } @classmethod def has_implicit_cancel(cls) -> bool: diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 6bb4ce5cd49..49489822d13 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -782,15 +782,16 @@ class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec): In Postgres, a catalog is called a "database". """ - return { - catalog - for (catalog,) in inspector.bind.execute( - text(""" + with inspector.engine.connect() as conn: + return { + catalog + for (catalog,) in conn.execute( + text(""" SELECT datname FROM pg_database WHERE datistemplate = false; - """) - ) - } + """) + ) + } @classmethod def get_table_names( diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index b9ca92c75e5..d48fbe7bf95 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -326,7 +326,8 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta): """ Get all catalogs. """ - return {catalog for (catalog,) in inspector.bind.execute(text("SHOW CATALOGS"))} + with inspector.engine.connect() as conn: + return {catalog for (catalog,) in conn.execute(text("SHOW CATALOGS"))} @classmethod def adjust_engine_params( @@ -712,9 +713,8 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta): :return: list of column objects """ full_table_name = cls.quote_table(table, inspector.engine.dialect) - return inspector.bind.execute( - text(f"SHOW COLUMNS FROM {full_table_name}") - ).fetchall() + with inspector.engine.connect() as conn: + return conn.execute(text(f"SHOW COLUMNS FROM {full_table_name}")).fetchall() @classmethod def _create_column_info( diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 1723d9b5963..89b120d7f61 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -265,12 +265,13 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec): In Snowflake, a catalog is called a "database". """ - return { - catalog - for (catalog,) in inspector.bind.execute( - text("SELECT DATABASE_NAME from information_schema.databases") - ) - } + with inspector.engine.connect() as conn: + return { + catalog + for (catalog,) in conn.execute( + text("SELECT DATABASE_NAME from information_schema.databases") + ) + } @classmethod def epoch_to_dttm(cls) -> str: diff --git a/superset/db_engine_specs/starrocks.py b/superset/db_engine_specs/starrocks.py index 35278a4688b..12175d0226c 100644 --- a/superset/db_engine_specs/starrocks.py +++ b/superset/db_engine_specs/starrocks.py @@ -344,24 +344,25 @@ class StarRocksEngineSpec(MySQLEngineSpec): The command returns columns: Catalog, Type, Comment """ try: - result = inspector.bind.execute(text("SHOW CATALOGS")) - catalogs = set() + with inspector.engine.connect() as conn: + result = conn.execute(text("SHOW CATALOGS")) + catalogs = set() - for row in result: - try: - if hasattr(row, "keys") and "Catalog" in row.keys(): - catalogs.add(row["Catalog"]) - elif hasattr(row, "Catalog"): - catalogs.add(row.Catalog) - else: - catalogs.add(row[0]) - except (AttributeError, TypeError, IndexError, KeyError) as ex: - logger.warning( - "Unable to extract catalog name from row: %s (%s)", row, ex - ) - continue + for row in result: + try: + if hasattr(row, "keys") and "Catalog" in row.keys(): + catalogs.add(row["Catalog"]) + elif hasattr(row, "Catalog"): + catalogs.add(row.Catalog) + else: + catalogs.add(row[0]) + except (AttributeError, TypeError, IndexError, KeyError) as ex: + logger.warning( + "Unable to extract catalog name from row: %s (%s)", row, ex + ) + continue - return catalogs + return catalogs except Exception as ex: # pylint: disable=broad-except logger.exception("Error fetching catalog names from SHOW CATALOGS: %s", ex) return set() @@ -375,8 +376,9 @@ class StarRocksEngineSpec(MySQLEngineSpec): (e.g., "catalog." sets the context to that catalog). """ try: - result = inspector.bind.execute(text("SHOW DATABASES")) - return {row[0] for row in result} + with inspector.engine.connect() as conn: + result = conn.execute(text("SHOW DATABASES")) + return {row[0] for row in result} except Exception as ex: # pylint: disable=broad-except logger.exception("Error fetching schema names from SHOW DATABASES: %s", ex) return set() diff --git a/superset/utils/mock_data.py b/superset/utils/mock_data.py index 557547a4bfd..9325e455edf 100644 --- a/superset/utils/mock_data.py +++ b/superset/utils/mock_data.py @@ -205,11 +205,12 @@ def add_data( table = DBTable(table_name, metadata, *column_objects) metadata.create_all(engine) - if not append: - engine.execute(table.delete()) + with engine.begin() as conn: + if not append: + conn.execute(table.delete()) - data = generate_data(columns, num_rows) - engine.execute(table.insert(), data) + data = generate_data(columns, num_rows) + conn.execute(table.insert(), data) def get_column_objects(columns: list[ColumnInfo]) -> list[Column]: @@ -287,9 +288,10 @@ def add_sample_rows(model: type[Model], count: int) -> Iterator[Model]: def get_valid_foreign_key(column: Column) -> Any: foreign_key = list(column.foreign_keys)[0] table_name, column_name = foreign_key.target_fullname.split(".", 1) - return db.engine.execute( - text(f"SELECT {column_name} FROM {table_name} LIMIT 1") # noqa: S608 - ).scalar() + with db.engine.connect() as conn: + return conn.execute( + text(f"SELECT {column_name} FROM {table_name} LIMIT 1") # noqa: S608 + ).scalar() def generate_value(column: Column) -> Any: diff --git a/tests/example_data/data_loading/pandas/pandas_data_loader.py b/tests/example_data/data_loading/pandas/pandas_data_loader.py index af3416b60d9..c7dfa09eddd 100644 --- a/tests/example_data/data_loading/pandas/pandas_data_loader.py +++ b/tests/example_data/data_loading/pandas/pandas_data_loader.py @@ -75,7 +75,8 @@ class PandasDataLoader(DataLoader): return None def remove_table(self, table_name: str) -> None: - self._db_engine.execute(text(f"DROP TABLE IF EXISTS {table_name}")) + with self._db_engine.begin() as conn: + conn.execute(text(f"DROP TABLE IF EXISTS {table_name}")) class TableToDfConvertor(ABC): diff --git a/tests/integration_tests/celery_tests.py b/tests/integration_tests/celery_tests.py index 15807dd9575..565b4f0ead2 100644 --- a/tests/integration_tests/celery_tests.py +++ b/tests/integration_tests/celery_tests.py @@ -121,7 +121,8 @@ def drop_table_if_exists(table_name: str, table_type: CTASMethod) -> None: sql = f"DROP {table_type.name} IF EXISTS {table_name}" database = get_example_database() with database.get_sqla_engine() as engine: - engine.execute(text(sql)) + with engine.begin() as conn: + conn.execute(text(sql)) def quote_f(value: Optional[str]): @@ -532,7 +533,8 @@ def test_teardown_with_app_context(): def delete_tmp_view_or_table(name: str, ctas_method: CTASMethod): - db.get_engine().execute(text(f"DROP {ctas_method.name} IF EXISTS {name}")) + with db.get_engine().begin() as conn: + conn.execute(text(f"DROP {ctas_method.name} IF EXISTS {name}")) def wait_for_success(result): diff --git a/tests/integration_tests/conftest.py b/tests/integration_tests/conftest.py index 2dde51433c5..c720a3a5b09 100644 --- a/tests/integration_tests/conftest.py +++ b/tests/integration_tests/conftest.py @@ -27,7 +27,7 @@ import pytest from flask.ctx import AppContext from flask_appbuilder.security.sqla import models as ab_models from sqlalchemy import text -from sqlalchemy.engine import Engine +from sqlalchemy.engine import Connection from superset import db, security_manager from superset.extensions import feature_flag_manager @@ -146,15 +146,15 @@ def setup_sample_data() -> Any: db.session.commit() -def drop_from_schema(engine: Engine, schema_name: str): - schemas = engine.execute(text("SHOW SCHEMAS")).fetchall() # noqa: F541 +def drop_from_schema(conn: Connection, schema_name: str): + schemas = conn.execute(text("SHOW SCHEMAS")).fetchall() # noqa: F541 if schema_name not in [s[0] for s in schemas]: # schema doesn't exist return - tables_or_views = engine.execute(text(f"SHOW TABLES in {schema_name}")).fetchall() + tables_or_views = conn.execute(text(f"SHOW TABLES in {schema_name}")).fetchall() for tv in tables_or_views: - engine.execute(text(f"DROP TABLE IF EXISTS {schema_name}.{tv[0]}")) - engine.execute(text(f"DROP VIEW IF EXISTS {schema_name}.{tv[0]}")) + conn.execute(text(f"DROP TABLE IF EXISTS {schema_name}.{tv[0]}")) + conn.execute(text(f"DROP VIEW IF EXISTS {schema_name}.{tv[0]}")) @pytest.fixture(scope="session") @@ -213,13 +213,14 @@ def setup_presto_if_needed(): if backend in {"presto", "hive"}: database = get_example_database() with database.get_sqla_engine() as engine: - drop_from_schema(engine, CTAS_SCHEMA_NAME) - engine.execute(text(f"DROP SCHEMA IF EXISTS {CTAS_SCHEMA_NAME}")) - engine.execute(text(f"CREATE SCHEMA {CTAS_SCHEMA_NAME}")) + with engine.begin() as conn: + drop_from_schema(conn, CTAS_SCHEMA_NAME) + conn.execute(text(f"DROP SCHEMA IF EXISTS {CTAS_SCHEMA_NAME}")) + conn.execute(text(f"CREATE SCHEMA {CTAS_SCHEMA_NAME}")) - drop_from_schema(engine, ADMIN_SCHEMA_NAME) - engine.execute(text(f"DROP SCHEMA IF EXISTS {ADMIN_SCHEMA_NAME}")) - engine.execute(text(f"CREATE SCHEMA {ADMIN_SCHEMA_NAME}")) + drop_from_schema(conn, ADMIN_SCHEMA_NAME) + conn.execute(text(f"DROP SCHEMA IF EXISTS {ADMIN_SCHEMA_NAME}")) + conn.execute(text(f"CREATE SCHEMA {ADMIN_SCHEMA_NAME}")) def with_feature_flags(**mock_feature_flags): @@ -356,34 +357,35 @@ def physical_dataset(): with example_database.get_sqla_engine() as engine: quoter = get_identifier_quoter(engine.name) # sqlite can only execute one statement at a time - engine.execute( - text(f""" - CREATE TABLE IF NOT EXISTS physical_dataset( - col1 INTEGER, - col2 VARCHAR(255), - col3 DECIMAL(4,2), - col4 VARCHAR(255), - col5 TIMESTAMP DEFAULT '1970-01-01 00:00:01', - col6 TIMESTAMP DEFAULT '1970-01-01 00:00:01', - {quoter("time column with spaces")} TIMESTAMP DEFAULT '1970-01-01 00:00:01' - ); - """) - ) - engine.execute( - text(""" - INSERT INTO physical_dataset values - (0, 'a', 1.0, NULL, '2000-01-01 00:00:00', '2002-01-03 00:00:00', '2002-01-03 00:00:00'), - (1, 'b', 1.1, NULL, '2000-01-02 00:00:00', '2002-02-04 00:00:00', '2002-02-04 00:00:00'), - (2, 'c', 1.2, NULL, '2000-01-03 00:00:00', '2002-03-07 00:00:00', '2002-03-07 00:00:00'), - (3, 'd', 1.3, NULL, '2000-01-04 00:00:00', '2002-04-12 00:00:00', '2002-04-12 00:00:00'), - (4, 'e', 1.4, NULL, '2000-01-05 00:00:00', '2002-05-11 00:00:00', '2002-05-11 00:00:00'), - (5, 'f', 1.5, NULL, '2000-01-06 00:00:00', '2002-06-13 00:00:00', '2002-06-13 00:00:00'), - (6, 'g', 1.6, NULL, '2000-01-07 00:00:00', '2002-07-15 00:00:00', '2002-07-15 00:00:00'), - (7, 'h', 1.7, NULL, '2000-01-08 00:00:00', '2002-08-18 00:00:00', '2002-08-18 00:00:00'), - (8, 'i', 1.8, NULL, '2000-01-09 00:00:00', '2002-09-20 00:00:00', '2002-09-20 00:00:00'), - (9, 'j', 1.9, NULL, '2000-01-10 00:00:00', '2002-10-22 00:00:00', '2002-10-22 00:00:00'); - """) # noqa: E501 - ) + with engine.begin() as conn: + conn.execute( + text(f""" + CREATE TABLE IF NOT EXISTS physical_dataset( + col1 INTEGER, + col2 VARCHAR(255), + col3 DECIMAL(4,2), + col4 VARCHAR(255), + col5 TIMESTAMP DEFAULT '1970-01-01 00:00:01', + col6 TIMESTAMP DEFAULT '1970-01-01 00:00:01', + {quoter("time column with spaces")} TIMESTAMP DEFAULT '1970-01-01 00:00:01' + ); + """) # noqa: E501 + ) + conn.execute( + text(""" + INSERT INTO physical_dataset values + (0, 'a', 1.0, NULL, '2000-01-01 00:00:00', '2002-01-03 00:00:00', '2002-01-03 00:00:00'), + (1, 'b', 1.1, NULL, '2000-01-02 00:00:00', '2002-02-04 00:00:00', '2002-02-04 00:00:00'), + (2, 'c', 1.2, NULL, '2000-01-03 00:00:00', '2002-03-07 00:00:00', '2002-03-07 00:00:00'), + (3, 'd', 1.3, NULL, '2000-01-04 00:00:00', '2002-04-12 00:00:00', '2002-04-12 00:00:00'), + (4, 'e', 1.4, NULL, '2000-01-05 00:00:00', '2002-05-11 00:00:00', '2002-05-11 00:00:00'), + (5, 'f', 1.5, NULL, '2000-01-06 00:00:00', '2002-06-13 00:00:00', '2002-06-13 00:00:00'), + (6, 'g', 1.6, NULL, '2000-01-07 00:00:00', '2002-07-15 00:00:00', '2002-07-15 00:00:00'), + (7, 'h', 1.7, NULL, '2000-01-08 00:00:00', '2002-08-18 00:00:00', '2002-08-18 00:00:00'), + (8, 'i', 1.8, NULL, '2000-01-09 00:00:00', '2002-09-20 00:00:00', '2002-09-20 00:00:00'), + (9, 'j', 1.9, NULL, '2000-01-10 00:00:00', '2002-10-22 00:00:00', '2002-10-22 00:00:00'); + """) # noqa: E501 + ) dataset = SqlaTable( table_name="physical_dataset", @@ -407,11 +409,9 @@ def physical_dataset(): yield dataset - engine.execute( - text(""" - DROP TABLE physical_dataset; - """) - ) + with example_database.get_sqla_engine() as engine: + with engine.begin() as conn: + conn.execute(text("DROP TABLE physical_dataset")) dataset = db.session.query(SqlaTable).filter_by(table_name="physical_dataset").all() for ds in dataset: db.session.delete(ds) diff --git a/tests/integration_tests/databases/commands/upload_test.py b/tests/integration_tests/databases/commands/upload_test.py index f3262763ec0..b4edc617979 100644 --- a/tests/integration_tests/databases/commands/upload_test.py +++ b/tests/integration_tests/databases/commands/upload_test.py @@ -80,8 +80,9 @@ def _setup_csv_upload(allowed_schemas: list[str] | None = None): upload_db = get_upload_db() with upload_db.get_sqla_engine() as engine: - engine.execute(text(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE}")) - engine.execute(text(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE_W_SCHEMA}")) + with engine.begin() as conn: + conn.execute(text(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE}")) + conn.execute(text(f"DROP TABLE IF EXISTS {CSV_UPLOAD_TABLE_W_SCHEMA}")) db.session.delete(upload_db) db.session.commit() @@ -114,12 +115,13 @@ def test_csv_upload_with_nulls(): CSVReader({"null_values": ["N/A", "None"]}), ).run() with upload_database.get_sqla_engine() as engine: - data = engine.execute(text(f"SELECT * from {CSV_UPLOAD_TABLE}")).fetchall() # noqa: S608 - assert data == [ - ("name1", None, "city1", "1-1-1980"), - ("name2", 29, None, "1-1-1981"), - ("name3", 28, "city3", "1-1-1982"), - ] + with engine.connect() as conn: + data = conn.execute(text(f"SELECT * from {CSV_UPLOAD_TABLE}")).fetchall() # noqa: S608 + assert data == [ + ("name1", None, "city1", "1-1-1980"), + ("name2", 29, None, "1-1-1981"), + ("name3", 28, "city3", "1-1-1982"), + ] @pytest.mark.usefixtures("setup_csv_upload_with_context") @@ -158,23 +160,26 @@ def test_csv_upload_with_index(): CSVReader({"dataframe_index": True, "index_label": "id"}), ).run() with upload_database.get_sqla_engine() as engine: - data = engine.execute(text(f"SELECT * from {CSV_UPLOAD_TABLE}")).fetchall() # noqa: S608 - assert data == [ - (0, "name1", 30, "city1", "1-1-1980"), - (1, "name2", 29, "city2", "1-1-1981"), - (2, "name3", 28, "city3", "1-1-1982"), - ] - # assert column names - assert [ # noqa: C416 - col - for col in engine.execute(text(f"SELECT * from {CSV_UPLOAD_TABLE}")).keys() # noqa: S608 - ] == [ - "id", - "Name", - "Age", - "City", - "Birth", - ] + with engine.connect() as conn: + data = conn.execute(text(f"SELECT * from {CSV_UPLOAD_TABLE}")).fetchall() # noqa: S608 + assert data == [ + (0, "name1", 30, "city1", "1-1-1980"), + (1, "name2", 29, "city2", "1-1-1981"), + (2, "name3", 28, "city3", "1-1-1982"), + ] + # assert column names + assert [ # noqa: C416 + col + for col in conn.execute( + text(f"SELECT * from {CSV_UPLOAD_TABLE}") # noqa: S608 + ).keys() # noqa: S608 + ] == [ + "id", + "Name", + "Age", + "City", + "Birth", + ] @only_postgresql diff --git a/tests/integration_tests/datasets/api_tests.py b/tests/integration_tests/datasets/api_tests.py index 7932c960f59..94d1fa6af9d 100644 --- a/tests/integration_tests/datasets/api_tests.py +++ b/tests/integration_tests/datasets/api_tests.py @@ -922,9 +922,15 @@ class TestDatasetApi(SupersetTestCase): example_db = get_example_database() with example_db.get_sqla_engine() as engine: - engine.execute( - text(f"CREATE TABLE {CTAS_SCHEMA_NAME}.birth_names AS SELECT 2 as two") - ) + with engine.begin() as conn: + conn.execute( + text( + f""" + CREATE TABLE {CTAS_SCHEMA_NAME}.birth_names AS + SELECT 2 as two + """ + ) + ) self.login(ADMIN_USERNAME) table_data = { @@ -943,7 +949,8 @@ class TestDatasetApi(SupersetTestCase): rv = self.client.delete(uri) assert rv.status_code == 200 with example_db.get_sqla_engine() as engine: - engine.execute(text(f"DROP TABLE {CTAS_SCHEMA_NAME}.birth_names")) + with engine.begin() as conn: + conn.execute(text(f"DROP TABLE {CTAS_SCHEMA_NAME}.birth_names")) def test_create_dataset_validate_database(self): """ @@ -3003,10 +3010,11 @@ class TestDatasetApi(SupersetTestCase): examples_db = get_example_database() with examples_db.get_sqla_engine() as engine: - engine.execute(text("DROP TABLE IF EXISTS test_create_sqla_table_api")) - engine.execute( - text("CREATE TABLE test_create_sqla_table_api AS SELECT 2 as col") - ) + with engine.begin() as conn: + conn.execute(text("DROP TABLE IF EXISTS test_create_sqla_table_api")) + conn.execute( + text("CREATE TABLE test_create_sqla_table_api AS SELECT 2 as col") + ) rv = self.client.post( "api/v1/dataset/get_or_create/", @@ -3030,7 +3038,8 @@ class TestDatasetApi(SupersetTestCase): self.items_to_delete = [table] with examples_db.get_sqla_engine() as engine: - engine.execute(text("DROP TABLE test_create_sqla_table_api")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE test_create_sqla_table_api")) def test_get_or_create_dataset_disambiguates_by_schema(self): """ diff --git a/tests/integration_tests/datasets/commands_tests.py b/tests/integration_tests/datasets/commands_tests.py index abfd1ca97ed..54b448142cf 100644 --- a/tests/integration_tests/datasets/commands_tests.py +++ b/tests/integration_tests/datasets/commands_tests.py @@ -279,8 +279,9 @@ class TestExportDatasetsCommand(SupersetTestCase): mock_g.user = security_manager.find_user("admin") examples_db = get_example_database() with examples_db.get_sqla_engine() as engine: - engine.execute("DROP TABLE IF EXISTS 中文") - engine.execute("CREATE TABLE 中文 AS SELECT 2 as col") + with engine.begin() as conn: + conn.execute(text("DROP TABLE IF EXISTS 中文")) + conn.execute(text("CREATE TABLE 中文 AS SELECT 2 as col")) # scope cleanup to the example database so datasets with the same name # on other databases are left untouched stale = db.session.query(SqlaTable).filter_by( @@ -307,7 +308,8 @@ class TestExportDatasetsCommand(SupersetTestCase): db.session.delete(example_dataset) db.session.commit() with examples_db.get_sqla_engine() as engine: - engine.execute("DROP TABLE 中文") + with engine.begin() as conn: + conn.execute(text("DROP TABLE 中文")) db.session.commit() @@ -618,10 +620,11 @@ class TestCreateDatasetCommand(SupersetTestCase): def test_create_dataset_command(self): examples_db = get_example_database() with examples_db.get_sqla_engine() as engine: - engine.execute(text("DROP TABLE IF EXISTS test_create_dataset_command")) - engine.execute( - text("CREATE TABLE test_create_dataset_command AS SELECT 2 as col") - ) + with engine.begin() as conn: + conn.execute(text("DROP TABLE IF EXISTS test_create_dataset_command")) + conn.execute( + text("CREATE TABLE test_create_dataset_command AS SELECT 2 as col") + ) with override_user(security_manager.find_user("admin")): table = CreateDatasetCommand( @@ -642,7 +645,8 @@ class TestCreateDatasetCommand(SupersetTestCase): db.session.delete(table) db.session.commit() with examples_db.get_sqla_engine() as engine: - engine.execute(text("DROP TABLE test_create_dataset_command")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE test_create_dataset_command")) db.session.commit() def test_create_dataset_command_not_allowed(self): diff --git a/tests/integration_tests/datasource_tests.py b/tests/integration_tests/datasource_tests.py index b5d56d45e8e..a79e80e4d29 100644 --- a/tests/integration_tests/datasource_tests.py +++ b/tests/integration_tests/datasource_tests.py @@ -63,27 +63,29 @@ def create_test_table_context(database: Database): full_table_name = f"{schema}.test_table" if schema else "test_table" with database.get_sqla_engine() as engine: - engine.execute( - text(f""" - CREATE TABLE IF NOT EXISTS {full_table_name} AS - SELECT 1 as first, 2 as second - """) - ) - engine.execute( - text(f""" - INSERT INTO {full_table_name} (first, second) VALUES (1, 2) - """) # noqa: S608 - ) - engine.execute( - text(f""" - INSERT INTO {full_table_name} (first, second) VALUES (3, 4) - """) # noqa: S608 - ) + with engine.begin() as conn: + conn.execute( + text(f""" + CREATE TABLE IF NOT EXISTS {full_table_name} AS + SELECT 1 as first, 2 as second + """) + ) + conn.execute( + text(f""" + INSERT INTO {full_table_name} (first, second) VALUES (1, 2) + """) # noqa: S608 + ) + conn.execute( + text(f""" + INSERT INTO {full_table_name} (first, second) VALUES (3, 4) + """) # noqa: S608 + ) yield db.session with database.get_sqla_engine() as engine: - engine.execute(text(f"DROP TABLE {full_table_name}")) + with engine.begin() as conn: + conn.execute(text(f"DROP TABLE {full_table_name}")) @contextmanager diff --git a/tests/integration_tests/db_engine_specs/hive_tests.py b/tests/integration_tests/db_engine_specs/hive_tests.py index 1ab1738ebf7..df6db6723af 100644 --- a/tests/integration_tests/db_engine_specs/hive_tests.py +++ b/tests/integration_tests/db_engine_specs/hive_tests.py @@ -183,12 +183,14 @@ def test_df_to_sql_if_exists_replace(mock_upload_to_s3, mock_g): app.config["CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC"]: lambda *args: "" # noqa: F722 mock_upload_to_s3.return_value = "mock-location" mock_g.user = True + mock_execute = mock.MagicMock(return_value=True) + mock_connection = mock.MagicMock() + mock_connection.execute = mock_execute + mock_engine = mock.MagicMock() + mock_engine.begin().__enter__.return_value = mock_connection mock_database = mock.MagicMock() mock_database.get_df.return_value.empty = False - mock_execute = mock.MagicMock(return_value=True) - mock_database.get_sqla_engine.return_value.__enter__.return_value.execute = ( - mock_execute - ) + mock_database.get_sqla_engine().__enter__.return_value = mock_engine table_name = "foobar" with app.app_context(): @@ -213,12 +215,14 @@ def test_df_to_sql_if_exists_replace_with_schema(mock_upload_to_s3, mock_g): app.config["CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC"]: lambda *args: "" # noqa: F722 mock_upload_to_s3.return_value = "mock-location" mock_g.user = True + mock_execute = mock.MagicMock(return_value=True) + mock_connection = mock.MagicMock() + mock_connection.execute = mock_execute + mock_engine = mock.MagicMock() + mock_engine.begin().__enter__.return_value = mock_connection mock_database = mock.MagicMock() mock_database.get_df.return_value.empty = False - mock_execute = mock.MagicMock(return_value=True) - mock_database.get_sqla_engine.return_value.__enter__.return_value.execute = ( - mock_execute - ) + mock_database.get_sqla_engine().__enter__.return_value = mock_engine table_name = "foobar" schema = "schema" diff --git a/tests/integration_tests/db_engine_specs/presto_tests.py b/tests/integration_tests/db_engine_specs/presto_tests.py index 10d9067f75b..3f8d8b51c5a 100644 --- a/tests/integration_tests/db_engine_specs/presto_tests.py +++ b/tests/integration_tests/db_engine_specs/presto_tests.py @@ -80,14 +80,16 @@ class TestPrestoDbEngineSpec(SupersetTestCase): assert result == {"a", "d"} def verify_presto_column(self, column, expected_results): - inspector = mock.Mock() + inspector = mock.MagicMock() preparer = inspector.engine.dialect.identifier_preparer preparer.quote_identifier = preparer.quote = preparer.quote_schema = ( lambda x: f'"{x}"' ) row = mock.Mock() row.Column, row.Type, row.Null = column - inspector.bind.execute.return_value.fetchall = mock.Mock(return_value=[row]) + inspector.engine.connect().__enter__().execute().fetchall = mock.Mock( + return_value=[row] + ) results = PrestoEngineSpec.get_columns(inspector, Table("", "")) assert len(expected_results) == len(results) for expected_result, result in zip(expected_results, results, strict=False): @@ -831,16 +833,17 @@ class TestPrestoDbEngineSpec(SupersetTestCase): preparer.quote_identifier = preparer.quote = preparer.quote_schema = ( lambda x: f'"{x}"' ) - inspector.bind.execute.return_value.fetchall = mock.MagicMock( - return_value=["a", "b"] + inspector.engine.connect().__enter__().execute.return_value.fetchall = ( + mock.MagicMock(return_value=["a", "b"]) ) table_name = "table_name" result = PrestoEngineSpec._show_columns(inspector, Table(table_name)) assert result == ["a", "b"] - assert_called_once_with_text( - inspector.bind.execute, - f'SHOW COLUMNS FROM "{table_name}"', - ) + with inspector.engine.connect() as conn: + assert_called_once_with_text( + conn.execute, + f'SHOW COLUMNS FROM "{table_name}"', + ) def test_show_columns_with_schema(self): inspector = mock.MagicMock() @@ -848,16 +851,17 @@ class TestPrestoDbEngineSpec(SupersetTestCase): preparer.quote_identifier = preparer.quote = preparer.quote_schema = ( lambda x: f'"{x}"' ) - inspector.bind.execute.return_value.fetchall = mock.MagicMock( - return_value=["a", "b"] + inspector.engine.connect().__enter__().execute.return_value.fetchall = ( + mock.MagicMock(return_value=["a", "b"]) ) table_name = "table_name" schema = "schema" result = PrestoEngineSpec._show_columns(inspector, Table(table_name, schema)) assert result == ["a", "b"] - assert_called_once_with_text( - inspector.bind.execute, f'SHOW COLUMNS FROM "{schema}"."{table_name}"' - ) + with inspector.engine.connect() as conn: + assert_called_once_with_text( + conn.execute, f'SHOW COLUMNS FROM "{schema}"."{table_name}"' + ) def test_is_column_name_quoted(self): column_name = "mock" diff --git a/tests/integration_tests/fixtures/energy_dashboard.py b/tests/integration_tests/fixtures/energy_dashboard.py index 82f75d9805a..78425e2c8a9 100644 --- a/tests/integration_tests/fixtures/energy_dashboard.py +++ b/tests/integration_tests/fixtures/energy_dashboard.py @@ -53,7 +53,8 @@ def load_energy_table_data(): yield with app.app_context(): with get_example_database().get_sqla_engine() as engine: - engine.execute(text("DROP TABLE IF EXISTS energy_usage")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE IF EXISTS energy_usage")) @pytest.fixture diff --git a/tests/integration_tests/fixtures/unicode_dashboard.py b/tests/integration_tests/fixtures/unicode_dashboard.py index 885ff500f80..d54c5ea4084 100644 --- a/tests/integration_tests/fixtures/unicode_dashboard.py +++ b/tests/integration_tests/fixtures/unicode_dashboard.py @@ -52,7 +52,8 @@ def load_unicode_data(): yield with app.app_context(): with get_example_database().get_sqla_engine() as engine: - engine.execute(text("DROP TABLE IF EXISTS unicode_test")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE IF EXISTS unicode_test")) @pytest.fixture diff --git a/tests/integration_tests/fixtures/world_bank_dashboard.py b/tests/integration_tests/fixtures/world_bank_dashboard.py index eed27d995fd..e412c6b1cd0 100644 --- a/tests/integration_tests/fixtures/world_bank_dashboard.py +++ b/tests/integration_tests/fixtures/world_bank_dashboard.py @@ -67,7 +67,8 @@ def load_world_bank_data(): yield with app.app_context(): with get_example_database().get_sqla_engine() as engine: - engine.execute(text("DROP TABLE IF EXISTS wb_health_population")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE IF EXISTS wb_health_population")) @pytest.fixture diff --git a/tests/integration_tests/reports/commands_tests.py b/tests/integration_tests/reports/commands_tests.py index e7762a796e1..5ebb0f2d232 100644 --- a/tests/integration_tests/reports/commands_tests.py +++ b/tests/integration_tests/reports/commands_tests.py @@ -178,18 +178,20 @@ def assert_log(state: str, error_message: Optional[str] = None): @contextmanager def create_test_table_context(database: Database): with database.get_sqla_engine() as engine: - engine.execute( - text(""" - CREATE TABLE IF NOT EXISTS test_table AS - SELECT 1 as first, 2 as second - """) - ) - engine.execute(text("INSERT INTO test_table (first, second) VALUES (1, 2)")) - engine.execute(text("INSERT INTO test_table (first, second) VALUES (3, 4)")) + with engine.begin() as conn: + conn.execute( + text(""" + CREATE TABLE IF NOT EXISTS test_table AS + SELECT 1 as first, 2 as second + """) + ) + conn.execute(text("INSERT INTO test_table (first, second) VALUES (1, 2)")) + conn.execute(text("INSERT INTO test_table (first, second) VALUES (3, 4)")) yield db.session with database.get_sqla_engine() as engine: - engine.execute(text("DROP TABLE test_table")) + with engine.begin() as conn: + conn.execute(text("DROP TABLE test_table")) @pytest.fixture diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py index 0f01abfe6ae..21bc572a5fa 100644 --- a/tests/integration_tests/sqllab_tests.py +++ b/tests/integration_tests/sqllab_tests.py @@ -184,20 +184,22 @@ class TestSqlLab(SupersetTestCase): db.session.commit() examples_db = get_example_database() with examples_db.get_sqla_engine() as engine: - data = engine.execute( - text(f"SELECT * FROM admin_database.{tmp_table_name}") # noqa: S608 - ).fetchall() - names_count = engine.execute( - text(f"SELECT COUNT(*) FROM birth_names") # noqa: F541, S608 - ).first() - assert names_count[0] == len( - data - ) # SQL_MAX_ROW not applied due to the SQLLAB_CTAS_NO_LIMIT set to True + with engine.connect() as conn: + data = conn.execute( + text(f"SELECT * FROM admin_database.{tmp_table_name}") # noqa: S608 + ).fetchall() + names_count = conn.execute( + text(f"SELECT COUNT(*) FROM birth_names") # noqa: F541, S608 + ).first() + assert names_count[0] == len(data) + # SQL_MAX_ROW not applied due to the SQLLAB_CTAS_NO_LIMIT set + # to True # cleanup - engine.execute( - text(f"DROP {ctas_method.name} admin_database.{tmp_table_name}") - ) + with engine.begin() as conn: + conn.execute( + text(f"DROP {ctas_method.name} admin_database.{tmp_table_name}") + ) examples_db.allow_ctas = old_allow_ctas db.session.commit() @@ -294,12 +296,13 @@ class TestSqlLab(SupersetTestCase): ) with examples_db.get_sqla_engine() as engine: - engine.execute( - text(f""" - CREATE TABLE IF NOT EXISTS {CTAS_SCHEMA_NAME}.test_table AS - SELECT 1 as c1, 2 as c2 - """) # noqa: E501 - ) + with engine.begin() as conn: + conn.execute( + text(f""" + CREATE TABLE IF NOT EXISTS {CTAS_SCHEMA_NAME}.test_table AS + SELECT 1 as c1, 2 as c2 + """) # noqa: E501 + ) # SQL Lab raw query access requires datasource_access on a registered # Superset dataset. schema_access alone is no longer sufficient, so @@ -318,7 +321,10 @@ class TestSqlLab(SupersetTestCase): db.session.query(Query).delete() with get_example_database().get_sqla_engine() as engine: - engine.execute(text(f"DROP TABLE IF EXISTS {CTAS_SCHEMA_NAME}.test_table")) + with engine.begin() as conn: + conn.execute( + text(f"DROP TABLE IF EXISTS {CTAS_SCHEMA_NAME}.test_table") + ) db.session.commit() def test_alias_duplicate(self): diff --git a/tests/unit_tests/db_engine_specs/test_datastore.py b/tests/unit_tests/db_engine_specs/test_datastore.py index 39433140da7..a58a80e193a 100644 --- a/tests/unit_tests/db_engine_specs/test_datastore.py +++ b/tests/unit_tests/db_engine_specs/test_datastore.py @@ -974,7 +974,7 @@ def test_get_catalog_names(mocker: MockerFixture) -> None: database = mocker.MagicMock() inspector = mocker.MagicMock() - inspector.bind.execute.return_value = [] + inspector.engine.connect().__enter__().execute.return_value = [] mocker.patch( "superset.db_engine_specs.base.BaseEngineSpec.get_catalog_names", diff --git a/tests/unit_tests/db_engine_specs/test_doris.py b/tests/unit_tests/db_engine_specs/test_doris.py index b0c44b3d478..1e0a2b083cc 100644 --- a/tests/unit_tests/db_engine_specs/test_doris.py +++ b/tests/unit_tests/db_engine_specs/test_doris.py @@ -16,7 +16,7 @@ # under the License. from typing import Any, Optional -from unittest.mock import Mock +from unittest.mock import MagicMock, Mock import pytest from pytest_mock import MockerFixture @@ -230,7 +230,7 @@ def test_get_default_catalog( mocker.MagicMock(IsCurrent=True, CatalogName="catalog2"), ] with database.get_sqla_engine() as engine: - engine.execute.return_value = rows + engine.connect().__enter__().execute.return_value = rows assert DorisEngineSpec.get_default_catalog(database) == expected_catalog @@ -266,16 +266,17 @@ def test_get_catalog_names( from superset.models.core import Database database = Mock(spec=Database) - inspector = Mock() - inspector.bind.execute.return_value = mock_catalogs + inspector = MagicMock() + inspector.engine.connect().__enter__().execute.return_value = mock_catalogs catalogs = DorisEngineSpec.get_catalog_names(database, inspector) # Verify the SQL query - assert_called_once_with_text( - inspector.bind.execute, - "SHOW CATALOGS", - ) + with inspector.engine.connect() as conn: + assert_called_once_with_text( + conn.execute, + "SHOW CATALOGS", + ) # Verify the returned catalog names assert catalogs == expected_result diff --git a/tests/unit_tests/db_engine_specs/test_starrocks.py b/tests/unit_tests/db_engine_specs/test_starrocks.py index e3ef836f36d..2ea2a56f285 100644 --- a/tests/unit_tests/db_engine_specs/test_starrocks.py +++ b/tests/unit_tests/db_engine_specs/test_starrocks.py @@ -242,7 +242,11 @@ def test_get_catalog_names(mocker: MockerFixture) -> None: mock_row_3.keys.return_value = ["Catalog", "Type", "Comment"] mock_row_3.__getitem__ = lambda self, key: "iceberg" if key == "Catalog" else None - inspector.bind.execute.return_value = [mock_row_1, mock_row_2, mock_row_3] + inspector.engine.connect().__enter__().execute.return_value = [ + mock_row_1, + mock_row_2, + mock_row_3, + ] catalogs = StarRocksEngineSpec.get_catalog_names(database, inspector) assert catalogs == {"default_catalog", "hive", "iceberg"} diff --git a/tests/unit_tests/db_engine_specs/test_trino.py b/tests/unit_tests/db_engine_specs/test_trino.py index fcce320208e..462c7efbebc 100644 --- a/tests/unit_tests/db_engine_specs/test_trino.py +++ b/tests/unit_tests/db_engine_specs/test_trino.py @@ -683,11 +683,14 @@ def test_get_columns_error(mocker: MockerFixture): "The specified table does not exist." ) Row = namedtuple("Row", ["Column", "Type"]) - mock_inspector.bind.execute().fetchall.return_value = [ + + mock_connection = mocker.MagicMock() + mock_connection.execute().fetchall.return_value = [ Row("field1", "row(a varchar, b date)"), Row("field2", "row(r1 row(a varchar, b varchar))"), Row("field3", "int"), ] + mock_inspector.engine.connect().__enter__.return_value = mock_connection actual = TrinoEngineSpec.get_columns(mock_inspector, Table("table", "schema")) expected = [ @@ -722,10 +725,11 @@ def test_get_columns_error(mocker: MockerFixture): _assert_columns_equal(actual, expected) - assert_called_with_text( - mock_inspector.bind.execute, - 'SHOW COLUMNS FROM schema."table"', - ) + with mock_inspector.engine.connect() as conn: + assert_called_with_text( + conn.execute, + 'SHOW COLUMNS FROM schema."table"', + ) def test_get_columns_expand_rows(mocker: MockerFixture): diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index b07abe3ac7b..805fec6edbb 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -318,7 +318,10 @@ def test_get_all_catalog_names(mocker: MockerFixture) -> None: get_inspector = mocker.patch.object(database, "get_inspector") with get_inspector() as inspector: - inspector.bind.execute.return_value = [("examples",), ("other",)] + inspector.engine.connect().__enter__().execute.return_value = [ + ("examples",), + ("other",), + ] assert database.get_all_catalog_names(force=True) == {"examples", "other"} get_inspector.assert_called_with()