From a1a4ca023854230c42e347c627ae5d8489d70ca1 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Mon, 2 Feb 2026 10:26:35 -0800 Subject: [PATCH] refactor(db): use Dialect instead of Engine in select_star to avoid SSH tunnels (#35540) Co-authored-by: Claude (cherry picked from commit e9b494163b99872d0163d3d2c5a66742ad44a964) --- superset/db_engine_specs/base.py | 6 ++--- superset/db_engine_specs/bigquery.py | 5 +++-- superset/db_engine_specs/hive.py | 6 ++--- superset/db_engine_specs/presto.py | 6 ++--- superset/models/core.py | 22 +++++++++---------- tests/unit_tests/db_engine_specs/test_base.py | 5 ++--- .../db_engine_specs/test_bigquery.py | 5 ++--- tests/unit_tests/db_engine_specs/test_hive.py | 4 ++-- .../db_engine_specs/test_postgres.py | 4 ++-- .../unit_tests/db_engine_specs/test_presto.py | 4 ++-- 10 files changed, 33 insertions(+), 34 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 5cc2e1a232a..f8e1f2a7b06 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1679,7 +1679,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods cls, database: Database, table: Table, - engine: Engine, + dialect: Dialect, limit: int = 100, show_cols: bool = False, indent: bool = True, @@ -1693,7 +1693,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods :param database: Database instance :param table: Table instance - :param engine: SqlAlchemy Engine instance + :param dialect: SqlAlchemy Dialect instance :param limit: limit to impose on query :param show_cols: Show columns in query; otherwise use "*" :param indent: Add indentation to query @@ -1713,7 +1713,7 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods if show_cols: fields = cls._get_fields(cols) - full_table_name = cls.quote_table(table, engine.dialect) + full_table_name = cls.quote_table(table, dialect) qry = select(fields).select_from(text(full_table_name)) qry = qry.limit(limit) diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py index b69361c2d11..724bf9cf332 100644 --- a/superset/db_engine_specs/bigquery.py +++ b/superset/db_engine_specs/bigquery.py @@ -32,6 +32,7 @@ from marshmallow import fields, Schema from marshmallow.exceptions import ValidationError from sqlalchemy import column, func, types from sqlalchemy.engine.base import Engine +from sqlalchemy.engine.interfaces import Dialect from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.url import URL from sqlalchemy.sql import column as sql_column, select, sqltypes @@ -721,7 +722,7 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met cls, database: Database, table: Table, - engine: Engine, + dialect: Dialect, limit: int = 100, show_cols: bool = False, indent: bool = True, @@ -781,7 +782,7 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met return super().select_star( database, table, - engine, + dialect, limit, show_cols, indent, diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py index e33438338d2..aeb73015eb3 100644 --- a/superset/db_engine_specs/hive.py +++ b/superset/db_engine_specs/hive.py @@ -31,7 +31,7 @@ import pyarrow as pa import pyarrow.parquet as pq from flask import current_app as app, g from sqlalchemy import Column, text, types -from sqlalchemy.engine.base import Engine +from sqlalchemy.engine.interfaces import Dialect from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.url import URL from sqlalchemy.sql.expression import ColumnClause, Select @@ -492,7 +492,7 @@ class HiveEngineSpec(PrestoEngineSpec): cls, database: Database, table: Table, - engine: Engine, + dialect: Dialect, limit: int = 100, show_cols: bool = False, indent: bool = True, @@ -505,7 +505,7 @@ class HiveEngineSpec(PrestoEngineSpec): return super(PrestoEngineSpec, cls).select_star( database, table, - engine, + dialect, limit, show_cols, indent, diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py index 32c82a6af11..49ba7477065 100644 --- a/superset/db_engine_specs/presto.py +++ b/superset/db_engine_specs/presto.py @@ -34,7 +34,7 @@ from flask import current_app as app from flask_babel import gettext as __, lazy_gettext as _ from packaging.version import Version from sqlalchemy import Column, literal_column, types -from sqlalchemy.engine.base import Engine +from sqlalchemy.engine.interfaces import Dialect from sqlalchemy.engine.reflection import Inspector from sqlalchemy.engine.result import Row as ResultRow from sqlalchemy.engine.url import URL @@ -1100,7 +1100,7 @@ class PrestoEngineSpec(PrestoBaseEngineSpec): cls, database: Database, table: Table, - engine: Engine, + dialect: Dialect, limit: int = 100, show_cols: bool = False, indent: bool = True, @@ -1124,7 +1124,7 @@ class PrestoEngineSpec(PrestoBaseEngineSpec): return super().select_star( database, table, - engine, + dialect, limit, show_cols, indent, diff --git a/superset/models/core.py b/superset/models/core.py index 116c2785183..437773648f9 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -828,17 +828,17 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable cols: list[ResultSetColumnType] | None = None, ) -> str: """Generates a ``select *`` statement in the proper dialect""" - with self.get_sqla_engine(catalog=table.catalog, schema=table.schema) as engine: - return self.db_engine_spec.select_star( - self, - table, - engine=engine, - limit=limit, - show_cols=show_cols, - indent=indent, - latest_partition=latest_partition, - cols=cols, - ) + dialect = self.get_dialect() + return self.db_engine_spec.select_star( + self, + table, + dialect=dialect, + limit=limit, + show_cols=show_cols, + indent=indent, + latest_partition=latest_partition, + cols=cols, + ) def apply_limit_to_sql( self, diff --git a/tests/unit_tests/db_engine_specs/test_base.py b/tests/unit_tests/db_engine_specs/test_base.py index bff4c931171..b2de20f8153 100644 --- a/tests/unit_tests/db_engine_specs/test_base.py +++ b/tests/unit_tests/db_engine_specs/test_base.py @@ -229,13 +229,12 @@ def test_select_star(mocker: MockerFixture) -> None: query.compile(dialect=sqlite.dialect()) ) - engine = mocker.MagicMock() - engine.dialect = sqlite.dialect() + dialect = sqlite.dialect() sql = BaseEngineSpec.select_star( database=database, table=Table("my_table", "my_schema", "my_catalog"), - engine=engine, + dialect=dialect, limit=100, show_cols=True, indent=True, diff --git a/tests/unit_tests/db_engine_specs/test_bigquery.py b/tests/unit_tests/db_engine_specs/test_bigquery.py index 7d2e8ef3e41..a7231508e8e 100644 --- a/tests/unit_tests/db_engine_specs/test_bigquery.py +++ b/tests/unit_tests/db_engine_specs/test_bigquery.py @@ -154,13 +154,12 @@ def test_select_star(mocker: MockerFixture) -> None: query.compile(dialect=BigQueryDialect(), compile_kwargs={"literal_binds": True}) ) - engine = mocker.MagicMock() - engine.dialect = BigQueryDialect() + dialect = BigQueryDialect() sql = BigQueryEngineSpec.select_star( database=database, table=Table("my_table"), - engine=engine, + dialect=dialect, limit=100, show_cols=True, indent=True, diff --git a/tests/unit_tests/db_engine_specs/test_hive.py b/tests/unit_tests/db_engine_specs/test_hive.py index 2bbb979782f..ec307154618 100644 --- a/tests/unit_tests/db_engine_specs/test_hive.py +++ b/tests/unit_tests/db_engine_specs/test_hive.py @@ -71,7 +71,7 @@ def test_select_star(mocker: MockerFixture) -> None: from superset.db_engine_specs.hive import HiveEngineSpec database = mocker.MagicMock() - engine = mocker.MagicMock() + dialect = mocker.MagicMock() def quote_table(table: Table, dialect: Dialect) -> str: return ".".join( @@ -83,7 +83,7 @@ def test_select_star(mocker: MockerFixture) -> None: HiveEngineSpec.select_star( database=database, table=Table("my_table", "my_schema", "my_catalog"), - engine=engine, + dialect=dialect, limit=100, show_cols=False, indent=True, diff --git a/tests/unit_tests/db_engine_specs/test_postgres.py b/tests/unit_tests/db_engine_specs/test_postgres.py index 290aa3575c9..ff11dd9aa42 100644 --- a/tests/unit_tests/db_engine_specs/test_postgres.py +++ b/tests/unit_tests/db_engine_specs/test_postgres.py @@ -252,7 +252,7 @@ def test_select_star(mocker: MockerFixture) -> None: Test the ``select_star`` method. """ database = mocker.MagicMock() - engine = mocker.MagicMock() + dialect = mocker.MagicMock() def quote_table(table: Table, dialect: Dialect) -> str: return ".".join( @@ -264,7 +264,7 @@ def test_select_star(mocker: MockerFixture) -> None: spec.select_star( database=database, table=Table("my_table", "my_schema", "my_catalog"), - engine=engine, + dialect=dialect, limit=100, show_cols=False, indent=True, diff --git a/tests/unit_tests/db_engine_specs/test_presto.py b/tests/unit_tests/db_engine_specs/test_presto.py index 857bc19b9be..d73b46f861d 100644 --- a/tests/unit_tests/db_engine_specs/test_presto.py +++ b/tests/unit_tests/db_engine_specs/test_presto.py @@ -314,7 +314,7 @@ def test_select_star(mocker: MockerFixture) -> None: from superset.db_engine_specs.presto import PrestoEngineSpec as spec # noqa: N813 database = mocker.MagicMock() - engine = mocker.MagicMock() + dialect = mocker.MagicMock() def quote_table(table: Table, dialect: Dialect) -> str: return ".".join( @@ -326,7 +326,7 @@ def test_select_star(mocker: MockerFixture) -> None: spec.select_star( database=database, table=Table("my_table", "my_schema", "my_catalog"), - engine=engine, + dialect=dialect, limit=100, show_cols=False, indent=True,