refactor: use Dialect instead of Engine in select_star to avoid SSH tunnels
Change the select_star method signature across all database engine specs
to use dialect: Dialect instead of engine: Engine. This prevents
unnecessary SSH tunnel connections from being created when only the
database dialect is needed for SQL generation.
Changes:
- Update select_star signature in BaseEngineSpec and subclasses
- Change Database.select_star to use get_dialect() instead of get_sqla_engine()
- Update all related unit tests
- Remove unused Engine imports
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude <noreply@anthropic.com>
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 74ba20e..41908bf 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1651,7 +1651,7 @@
cls,
database: Database,
table: Table,
- engine: Engine,
+ dialect: Dialect,
limit: int = 100,
show_cols: bool = False,
indent: bool = True,
@@ -1665,7 +1665,7 @@
: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
@@ -1685,7 +1685,7 @@
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 9ff78de..43352f8 100644
--- a/superset/db_engine_specs/bigquery.py
+++ b/superset/db_engine_specs/bigquery.py
@@ -32,6 +32,7 @@
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 @@
cls,
database: Database,
table: Table,
- engine: Engine,
+ dialect: Dialect,
limit: int = 100,
show_cols: bool = False,
indent: bool = True,
@@ -781,7 +782,7 @@
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 e334383..aeb7301 100644
--- a/superset/db_engine_specs/hive.py
+++ b/superset/db_engine_specs/hive.py
@@ -31,7 +31,7 @@
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 @@
cls,
database: Database,
table: Table,
- engine: Engine,
+ dialect: Dialect,
limit: int = 100,
show_cols: bool = False,
indent: bool = True,
@@ -505,7 +505,7 @@
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 32c82a6..49ba747 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -34,7 +34,7 @@
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 @@
cls,
database: Database,
table: Table,
- engine: Engine,
+ dialect: Dialect,
limit: int = 100,
show_cols: bool = False,
indent: bool = True,
@@ -1124,7 +1124,7 @@
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 f6643d1..5e7b2ea 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -833,17 +833,17 @@
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 6db0edb..aea1186 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 @@
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 7d2e8ef..a723150 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 @@
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 2bbb979..ec30715 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 @@
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 @@
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 290aa35..ff11dd9 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 @@
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 @@
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 857bc19..d73b46f 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 @@
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 @@
spec.select_star(
database=database,
table=Table("my_table", "my_schema", "my_catalog"),
- engine=engine,
+ dialect=dialect,
limit=100,
show_cols=False,
indent=True,
