You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2025/08/25 21:27:56 UTC
(superset) branch standardize-db-engine-spec-execute updated: chore: standardize DB engine spec query execution
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch standardize-db-engine-spec-execute
in repository https://gitbox.apache.org/repos/asf/superset.git
The following commit(s) were added to refs/heads/standardize-db-engine-spec-execute by this push:
new 8d15e70bb0 chore: standardize DB engine spec query execution
8d15e70bb0 is described below
commit 8d15e70bb02829c918d65fe2e4291ca9459d8e7c
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Mon Aug 25 17:27:41 2025 -0400
chore: standardize DB engine spec query execution
---
superset/db_engine_specs/base.py | 33 +++++++++++++++++++++++++++++++++
superset/db_engine_specs/bigquery.py | 12 +++++-------
superset/db_engine_specs/databricks.py | 6 ++++--
superset/db_engine_specs/hive.py | 11 ++++++-----
superset/db_engine_specs/postgres.py | 13 ++++++-------
superset/db_engine_specs/presto.py | 3 ++-
superset/utils/core.py | 1 +
7 files changed, 57 insertions(+), 22 deletions(-)
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 4eb30f40e1..6a7f2e2ca7 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1854,6 +1854,39 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
cls.start_oauth2_dance(database)
raise cls.get_dbapi_mapped_exception(ex) from ex
+ @classmethod
+ def execute_metadata_query(
+ cls,
+ database: Database,
+ query: str,
+ catalog: str | None = None,
+ schema: str | None = None,
+ ) -> list[tuple[Any, ...]]:
+ """
+ Standardized method for executing metadata queries.
+
+ This method provides a unified interface for all metadata query operations
+ across different database engines using SQLAlchemy connections.
+
+ For single-row results, add "LIMIT 1" to your query rather than using
+ separate fetch parameters.
+
+ :param database: Database instance
+ :param query: SQL query to execute for metadata
+ :param catalog: Optional catalog/database name
+ :param schema: Optional schema name
+ :return: Query results as list of tuples
+ """
+ with cls.get_engine(
+ database,
+ catalog=catalog,
+ schema=schema,
+ source=utils.QuerySource.METADATA,
+ ) as engine:
+ with engine.connect() as conn:
+ result = conn.execute(text(query))
+ return result.fetchall()
+
@classmethod
def needs_oauth2(cls, ex: Exception) -> bool:
"""
diff --git a/superset/db_engine_specs/bigquery.py b/superset/db_engine_specs/bigquery.py
index 9aa223bbcb..f2cc920903 100644
--- a/superset/db_engine_specs/bigquery.py
+++ b/superset/db_engine_specs/bigquery.py
@@ -342,15 +342,13 @@ class BigQueryEngineSpec(BaseEngineSpec): # pylint: disable=too-many-public-met
)
# Run the query and handle result
- with database.get_raw_connection(
+ results = cls.execute_metadata_query(
+ database,
+ f"{str(compiled_query)} LIMIT 1",
catalog=table.catalog,
schema=table.schema,
- ) as conn:
- cursor = conn.cursor()
- cursor.execute(str(compiled_query))
- if row := cursor.fetchone():
- return row[0]
- return None
+ )
+ return results[0][0] if results else None
@classmethod
def get_time_partition_column(
diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py
index a885684408..752e9df503 100644
--- a/superset/db_engine_specs/databricks.py
+++ b/superset/db_engine_specs/databricks.py
@@ -532,7 +532,8 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec):
database: Database,
inspector: Inspector,
) -> set[str]:
- return {catalog for (catalog,) in inspector.bind.execute("SHOW CATALOGS")}
+ results = cls.execute_metadata_query(database, "SHOW CATALOGS")
+ return {catalog for (catalog,) in results}
class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec):
@@ -624,7 +625,8 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec):
database: Database,
inspector: Inspector,
) -> set[str]:
- return {catalog for (catalog,) in inspector.bind.execute("SHOW CATALOGS")}
+ results = cls.execute_metadata_query(database, "SHOW CATALOGS")
+ return {catalog for (catalog,) in results}
@classmethod
def adjust_engine_params(
diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py
index e33438338d..77342d846d 100644
--- a/superset/db_engine_specs/hive.py
+++ b/superset/db_engine_specs/hive.py
@@ -614,8 +614,9 @@ class HiveEngineSpec(PrestoEngineSpec):
if schema:
sql += f" IN `{schema}`"
- with database.get_raw_connection(schema=schema) as conn:
- cursor = conn.cursor()
- cursor.execute(sql)
- results = cursor.fetchall()
- return {row[0] for row in results}
+ results = cls.execute_metadata_query(
+ database,
+ sql,
+ schema=schema,
+ )
+ return {row[0] for row in results}
diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py
index b259164e4f..670d69c65c 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -393,15 +393,14 @@ class PostgresEngineSpec(BasicParametersMixin, PostgresBaseEngineSpec):
In Postgres, a catalog is called a "database".
"""
- return {
- catalog
- for (catalog,) in inspector.bind.execute(
- """
+ results = cls.execute_metadata_query(
+ database,
+ """
SELECT datname FROM pg_database
WHERE datistemplate = false;
- """
- )
- }
+ """,
+ )
+ return {catalog for (catalog,) in results}
@classmethod
def get_table_names(
diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
index 32c82a6af1..4066ae1e1b 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -321,7 +321,8 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta):
"""
Get all catalogs.
"""
- return {catalog for (catalog,) in inspector.bind.execute("SHOW CATALOGS")}
+ results = cls.execute_metadata_query(database, "SHOW CATALOGS")
+ return {catalog for (catalog,) in results}
@classmethod
def adjust_engine_params(
diff --git a/superset/utils/core.py b/superset/utils/core.py
index 3185148ff8..8827f8d421 100644
--- a/superset/utils/core.py
+++ b/superset/utils/core.py
@@ -320,6 +320,7 @@ class QuerySource(Enum):
CHART = 0
DASHBOARD = 1
SQL_LAB = 2
+ METADATA = 3
class QueryStatus(StrEnum):