diff --git a/superset/models/core.py b/superset/models/core.py index 0bb5ddf2be3..3daabc9dfab 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)