(superset) 01/01: fix: Chart execution for Databricks

[va...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

vavila pushed a commit to branch fix/query-execution-cursor
in repository https://gitbox.apache.org/repos/asf/superset.git

commit ef815ff89ce9c31a2fa32ad08e18cfb205b68974
Author: Vitor Avila <vi...@preset.io>
AuthorDate: Fri Aug 29 14:59:48 2025 -0300

    fix: Chart execution for Databricks
---
 superset/models/core.py | 14 +++++++++-----
 1 file changed, 9 insertions(+), 5 deletions(-)

diff --git a/superset/models/core.py b/superset/models/core.py
index 0bb5ddf2be..3daabc9dfa 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -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)