This is an automated email from the ASF dual-hosted git repository.
rusackas pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/master by this push:
new 68ebc22e1ee refactor: replace Engine.execute() with
Connection.execute() (#41917)
68ebc22e1ee is described below
commit 68ebc22e1eea2ae6da93d5896f209d2e564f9d4f
Author: Hans Yu <[email protected]>
AuthorDate: Mon Jul 13 18:40:33 2026 +0200
refactor: replace Engine.execute() with Connection.execute() (#41917)
---
pytest.ini | 2 +-
scripts/benchmark_migration.py | 9 ++-
superset/cli/test_db.py | 14 ++--
superset/db_engine_specs/databricks.py | 17 ++--
superset/db_engine_specs/doris.py | 12 +--
superset/db_engine_specs/duckdb.py | 13 ++--
superset/db_engine_specs/hive.py | 36 +++++----
superset/db_engine_specs/impala.py | 11 +--
superset/db_engine_specs/postgres.py | 15 ++--
superset/db_engine_specs/presto.py | 8 +-
superset/db_engine_specs/snowflake.py | 13 ++--
superset/db_engine_specs/starrocks.py | 42 +++++-----
superset/utils/mock_data.py | 16 ++--
.../data_loading/pandas/pandas_data_loader.py | 3 +-
tests/integration_tests/celery_tests.py | 6 +-
tests/integration_tests/conftest.py | 90 +++++++++++-----------
.../databases/commands/upload_test.py | 55 +++++++------
tests/integration_tests/datasets/api_tests.py | 27 ++++---
tests/integration_tests/datasets/commands_tests.py | 20 +++--
tests/integration_tests/datasource_tests.py | 36 +++++----
.../db_engine_specs/hive_tests.py | 20 +++--
.../db_engine_specs/presto_tests.py | 30 ++++----
.../integration_tests/fixtures/energy_dashboard.py | 3 +-
.../fixtures/unicode_dashboard.py | 3 +-
.../fixtures/world_bank_dashboard.py | 3 +-
tests/integration_tests/reports/commands_tests.py | 20 ++---
tests/integration_tests/sqllab_tests.py | 44 ++++++-----
tests/unit_tests/db_engine_specs/test_datastore.py | 2 +-
tests/unit_tests/db_engine_specs/test_doris.py | 17 ++--
tests/unit_tests/db_engine_specs/test_starrocks.py | 6 +-
tests/unit_tests/db_engine_specs/test_trino.py | 14 ++--
tests/unit_tests/models/core_test.py | 5 +-
32 files changed, 344 insertions(+), 268 deletions(-)
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()
-
- 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
+ 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
+
+ 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()
