[fix] Enforce the query result to contain a data-frame (#8935)

superset/connectors/druid/models.py

@@ -988,11 +988,9 @@ def values_for_column(self, column_name: str, limit: int = 10000) -> List:
     def get_query_str(self, query_obj, phase=1, client=None):
         return self.run_query(client=client, phase=phase, **query_obj)
 
-    def _add_filter_from_pre_query_data(
-        self, df: Optional[pd.DataFrame], dimensions, dim_filter
-    ):
+    def _add_filter_from_pre_query_data(self, df: pd.DataFrame, dimensions, dim_filter):
         ret = dim_filter
-        if df is not None and not df.empty:
+        if not df.empty:
             new_filters = []
             for unused, row in df.iterrows():
                 fields = []
@@ -1379,7 +1377,7 @@ def query(self, query_obj: Dict) -> QueryResult:
 
         if df is None or df.size == 0:
             return QueryResult(
-                df=pd.DataFrame([]),
+                df=pd.DataFrame(),
                 query=query_str,
                 duration=datetime.now() - qry_start_dttm,
             )

superset/connectors/sqla/models.py

@@ -85,7 +85,6 @@ class AnnotationDatasource(BaseDatasource):
     cache_timeout = 0
 
     def query(self, query_obj: Dict[str, Any]) -> QueryResult:
-        df = None
         error_message = None
         qry = db.session.query(Annotation)
         qry = qry.filter(Annotation.layer_id == query_obj["filter"][0]["val"])
@@ -97,6 +96,7 @@ def query(self, query_obj: Dict[str, Any]) -> QueryResult:
         try:
             df = pd.read_sql_query(qry.statement, db.engine)
         except Exception as e:
+            df = pd.DataFrame()
             status = utils.QueryStatus.FAILED
             logging.exception(e)
             error_message = utils.error_msg_from_exception(e)
@@ -995,7 +995,7 @@ def mutator(df: pd.DataFrame) -> None:
         try:
             df = self.database.get_df(sql, self.schema, mutator)
         except Exception as e:
-            df = None
+            df = pd.DataFrame()
             status = utils.QueryStatus.FAILED
             logging.exception(f"Query {sql} on schema {self.schema} failed")
             db_engine_spec = self.database.db_engine_spec

superset/models/helpers.py

@@ -24,6 +24,7 @@
 # isort and pylint disagree, isort should win
 # pylint: disable=ungrouped-imports
 import humanize
+import pandas as pd
 import sqlalchemy as sa
 import yaml
 from flask import escape, g, Markup
@@ -368,11 +369,11 @@ class QueryResult:  # pylint: disable=too-few-public-methods
     def __init__(  # pylint: disable=too-many-arguments
         self, df, query, duration, status=QueryStatus.SUCCESS, error_message=None
     ):
-        self.df = df  # pylint: disable=invalid-name
-        self.query = query
-        self.duration = duration
-        self.status = status
-        self.error_message = error_message
+        self.df: pd.DataFrame = df  # pylint: disable=invalid-name
+        self.query: str = query
+        self.duration: int = duration
+        self.status: str = status
+        self.error_message: Optional[str] = error_message
 
 
 class ExtraJSONMixin: