QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-07 18:35:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: Chart execution for Databricks (#34906)

Browse Source
This commit is contained in:
Vitor Avila
2025-08-29 15:35:00 -03:00
committed by GitHub
parent 812374b31b
commit 54f071138c
1 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
superset/models/core.py
View File

@@ -675,7 +675,7 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
catalog: str | None = None,
schema: str | None = None,
fetch_last_result: bool = False,
) -> tuple[Any, list[tuple[Any, ...]] | None]:
) -> tuple[Any, list[tuple[Any, ...]] | None, DbapiDescription | None]:
"""
Internal method to execute SQL with mutation and logging.
@@ -683,7 +683,8 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
:param catalog: Optional catalog name
:param schema: Optional schema name
:param fetch_last_result: Whether to fetch results from last statement
:return: Tuple of (cursor, rows) where rows is None if not fetching
:return: Tuple of (cursor, rows, description) where rows and description
are None if not fetching.
"""
script = SQLScript(sql, self.db_engine_spec.engine)
@@ -705,6 +706,7 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
with self.get_raw_connection(catalog=catalog, schema=schema) as conn:
cursor = conn.cursor()
rows = None
description = None
for i, statement in enumerate(script.statements):
sql_ = self.mutate_sql_based_on_config(
@@ -722,12 +724,14 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
# Fetch results from last statement if requested
if fetch_last_result and i == len(script.statements) - 1:
# Capture cursor.description while it's still valid
description = cursor.description
rows = self.db_engine_spec.fetch_data(cursor)
else:
# Consume results without storing
cursor.fetchall()
return cursor, rows
return cursor, rows, description
def execute_sql_statements(
self,
@@ -763,13 +767,13 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
schema: str | None = None,
mutator: Callable[[pd.DataFrame], None] | None = None,
) -> pd.DataFrame:
cursor, rows = self._execute_sql_with_mutation_and_logging(
cursor, rows, description = self._execute_sql_with_mutation_and_logging(
sql, catalog, schema, fetch_last_result=True
)
df = None
if rows is not None:
df = self.load_into_dataframe(cursor.description, rows)
df = self.load_into_dataframe(description, rows)
if mutator:
df = mutator(df)