chore: refactor, add typing and fix uncovered errors (apache#8900)

* Add type annotations and fix inconsistencies

* Address review comments

* Remove incorrect typing of jsonable obj

Author: villebro

superset/app.py

@@ -221,7 +221,7 @@ def configure_middlewares(self):
 
         if self.config["ENABLE_CHUNK_ENCODING"]:
 
-            class ChunkedEncodingFix(object):  # pylint: disable=too-few-public-methods
+            class ChunkedEncodingFix:  # pylint: disable=too-few-public-methods
                 def __init__(self, app):
                     self.app = app
 

superset/common/query_context.py

@@ -17,7 +17,7 @@
 import logging
 import pickle as pkl
 from datetime import datetime, timedelta
-from typing import Any, Dict, List, Optional
+from typing import Any, ClassVar, Dict, List, Optional
 
 import numpy as np
 import pandas as pd
@@ -41,8 +41,8 @@ class QueryContext:
     to retrieve the data payload for a given viz.
     """
 
-    cache_type: str = "df"
-    enforce_numerical_metrics: bool = True
+    cache_type: ClassVar[str] = "df"
+    enforce_numerical_metrics: ClassVar[bool] = True
 
     datasource: BaseDatasource
     queries: List[QueryObject]
@@ -53,20 +53,16 @@ class QueryContext:
     # a vanilla python type https://github.com/python/mypy/issues/5288
     def __init__(
         self,
-        datasource: Dict,
-        queries: List[Dict],
+        datasource: Dict[str, Any],
+        queries: List[Dict[str, Any]],
         force: bool = False,
         custom_cache_timeout: Optional[int] = None,
     ) -> None:
-        self.datasource = ConnectorRegistry.get_datasource(  # type: ignore
-            datasource.get("type"),  # type: ignore
-            int(datasource.get("id")),  # type: ignore
-            db.session,
+        self.datasource = ConnectorRegistry.get_datasource(
+            str(datasource["type"]), int(datasource["id"]), db.session
         )
-        self.queries = list(map(lambda query_obj: QueryObject(**query_obj), queries))
-
+        self.queries = [QueryObject(**query_obj) for query_obj in queries]
         self.force = force
-
         self.custom_cache_timeout = custom_cache_timeout
 
     def get_query_result(self, query_object: QueryObject) -> Dict[str, Any]:
@@ -78,7 +74,7 @@ def get_query_result(self, query_object: QueryObject) -> Dict[str, Any]:
 
         timestamp_format = None
         if self.datasource.type == "table":
-            dttm_col = self.datasource.get_col(query_object.granularity)
+            dttm_col = self.datasource.get_column(query_object.granularity)
             if dttm_col:
                 timestamp_format = dttm_col.python_date_format
 
@@ -115,17 +111,18 @@ def get_query_result(self, query_object: QueryObject) -> Dict[str, Any]:
             "df": df,
         }
 
+    @staticmethod
     def df_metrics_to_num(  # pylint: disable=invalid-name,no-self-use
-        self, df: pd.DataFrame, query_object: QueryObject
+        df: pd.DataFrame, query_object: QueryObject
     ) -> None:
         """Converting metrics to numeric when pandas.read_sql cannot"""
-        metrics = [metric for metric in query_object.metrics]
         for col, dtype in df.dtypes.items():
-            if dtype.type == np.object_ and col in metrics:
+            if dtype.type == np.object_ and col in query_object.metrics:
                 df[col] = pd.to_numeric(df[col], errors="coerce")
 
+    @staticmethod
     def get_data(  # pylint: disable=invalid-name,no-self-use
-        self, df: pd.DataFrame
+        df: pd.DataFrame
     ) -> List[Dict]:
         return df.to_dict(orient="records")
 

superset/config.py

@@ -449,7 +449,7 @@ def _try_json_readsha(filepath, length):  # pylint: disable=unused-argument
 # http://docs.celeryproject.org/en/latest/getting-started/brokers/index.html
 
 
-class CeleryConfig(object):  # pylint: disable=too-few-public-methods
+class CeleryConfig:  # pylint: disable=too-few-public-methods
     BROKER_URL = "sqla+sqlite:///celerydb.sqlite"
     CELERY_IMPORTS = ("superset.sql_lab", "superset.tasks")
     CELERY_RESULT_BACKEND = "db+sqlite:///celery_results.sqlite"

superset/connectors/base/models.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import json
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, Hashable, List, Optional, Type
 
 from flask_appbuilder.security.sqla.models import User
 from sqlalchemy import and_, Boolean, Column, Integer, String, Text
@@ -44,7 +44,7 @@ class BaseDatasource(
     baselink: Optional[str] = None  # url portion pointing to ModelView endpoint
     column_class: Optional[Type] = None  # link to derivative of BaseColumn
     metric_class: Optional[Type] = None  # link to derivative of BaseMetric
-    owner_class = None
+    owner_class: Optional[User] = None
 
     # Used to do code highlighting when displaying the query in the UI
     query_language: Optional[str] = None
@@ -342,11 +342,14 @@ def update_from_object(self, obj) -> None:
             obj.get("columns"), self.columns, self.column_class, "column_name"
         )
 
-    def get_extra_cache_keys(  # pylint: disable=unused-argument,no-self-use
-        self, query_obj: Dict
-    ) -> List[Any]:
+    def get_extra_cache_keys(  # pylint: disable=no-self-use
+        self, query_obj: Dict[str, Any]  # pylint: disable=unused-argument
+    ) -> List[Hashable]:
         """ If a datasource needs to provide additional keys for calculation of
         cache keys, those can be provided via this method
+
+        :param query_obj: The dict representation of a query object
+        :return: list of keys
         """
         return []
 
@@ -403,6 +406,10 @@ def is_string(self) -> bool:
     def expression(self):
         raise NotImplementedError()
 
+    @property
+    def python_date_format(self):
+        raise NotImplementedError()
+
     @property
     def data(self) -> Dict[str, Any]:
         attrs = (
@@ -415,7 +422,6 @@ def data(self) -> Dict[str, Any]:
             "groupby",
             "is_dttm",
             "type",
-            "python_date_format",
         )
         return {s: getattr(self, s) for s in attrs if hasattr(self, s)}
 

superset/connectors/connector_registry.py

@@ -26,7 +26,7 @@
     from superset.connectors.base.models import BaseDatasource
 
 
-class ConnectorRegistry(object):
+class ConnectorRegistry:
     """ Central Registry for all available datasource engines"""
 
     sources: Dict[str, Type["BaseDatasource"]] = {}
@@ -43,11 +43,11 @@ def register_sources(cls, datasource_config: OrderedDict) -> None:
     @classmethod
     def get_datasource(
         cls, datasource_type: str, datasource_id: int, session: Session
-    ) -> Optional["BaseDatasource"]:
+    ) -> "BaseDatasource":
         return (
             session.query(cls.sources[datasource_type])
             .filter_by(id=datasource_id)
-            .first()
+            .one()
         )
 
     @classmethod

superset/connectors/sqla/models.py

@@ -19,7 +19,7 @@
 import re
 from collections import OrderedDict
 from datetime import datetime
-from typing import Any, Dict, List, NamedTuple, Optional, Tuple, Union
+from typing import Any, Dict, Hashable, List, NamedTuple, Optional, Tuple, Union
 
 import pandas as pd
 import sqlalchemy as sa
@@ -84,7 +84,7 @@ class AnnotationDatasource(BaseDatasource):
 
     cache_timeout = 0
 
-    def query(self, query_obj: Dict) -> QueryResult:
+    def query(self, query_obj: Dict[str, Any]) -> QueryResult:
         df = None
         error_message = None
         qry = db.session.query(Annotation)
@@ -537,16 +537,9 @@ def select_star(self) -> str:
             latest_partition=False,
         )
 
-    def get_col(self, col_name: str) -> Optional[Column]:
-        columns = self.columns
-        for col in columns:
-            if col_name == col.column_name:
-                return col
-        return None
-
     @property
     def data(self) -> Dict:
-        d = super(SqlaTable, self).data
+        d = super().data
         if self.type == "table":
             grains = self.database.grains() or []
             if grains:
@@ -598,7 +591,7 @@ def mutate_query_from_config(self, sql: str) -> str:
     def get_template_processor(self, **kwargs):
         return get_template_processor(table=self, database=self.database, **kwargs)
 
-    def get_query_str_extended(self, query_obj: Dict) -> QueryStringExtended:
+    def get_query_str_extended(self, query_obj: Dict[str, Any]) -> QueryStringExtended:
         sqlaq = self.get_sqla_query(**query_obj)
         sql = self.database.compile_sqla_query(sqlaq.sqla_query)
         logging.info(sql)
@@ -608,7 +601,7 @@ def get_query_str_extended(self, query_obj: Dict) -> QueryStringExtended:
             labels_expected=sqlaq.labels_expected, sql=sql, prequeries=sqlaq.prequeries
         )
 
-    def get_query_str(self, query_obj: Dict) -> str:
+    def get_query_str(self, query_obj: Dict[str, Any]) -> str:
         query_str_ext = self.get_query_str_extended(query_obj)
         all_queries = query_str_ext.prequeries + [query_str_ext.sql]
         return ";\n\n".join(all_queries) + ";"
@@ -976,14 +969,23 @@ def _get_top_groups(
 
         return or_(*groups)
 
-    def query(self, query_obj: Dict) -> QueryResult:
+    def query(self, query_obj: Dict[str, Any]) -> QueryResult:
         qry_start_dttm = datetime.now()
         query_str_ext = self.get_query_str_extended(query_obj)
         sql = query_str_ext.sql
         status = utils.QueryStatus.SUCCESS
         error_message = None
 
-        def mutator(df):
+        def mutator(df: pd.DataFrame) -> None:
+            """
+            Some engines change the case or generate bespoke column names, either by
+            default or due to lack of support for aliasing. This function ensures that
+            the column names in the DataFrame correspond to what is expected by
+            the viz components.
+
+            :param df: Original DataFrame returned by the engine
+            """
+
             labels_expected = query_str_ext.labels_expected
             if df is not None and not df.empty:
                 if len(df.columns) != len(labels_expected):
@@ -993,7 +995,6 @@ def mutator(df):
                     )
                 else:
                     df.columns = labels_expected
-            return df
 
         try:
             df = self.database.get_df(sql, self.schema, mutator)
@@ -1135,13 +1136,16 @@ def query_datasources_by_name(
     def default_query(qry) -> Query:
         return qry.filter_by(is_sqllab_view=False)
 
-    def has_extra_cache_keys(self, query_obj: Dict) -> bool:
+    def has_calls_to_cache_key_wrapper(self, query_obj: Dict[str, Any]) -> bool:
         """
-        Detects the presence of calls to cache_key_wrapper in items in query_obj that can
-        be templated.
+        Detects the presence of calls to `cache_key_wrapper` in items in query_obj that
+        can be templated. If any are present, the query must be evaluated to extract
+        additional keys for the cache key. This method is needed to avoid executing
+        the template code unnecessarily, as it may contain expensive calls, e.g. to
+        extract the latest partition of a database.
 
         :param query_obj: query object to analyze
-        :return: True if at least one item calls cache_key_wrapper, otherwise False
+        :return: True if at least one item calls `cache_key_wrapper`, otherwise False
         """
         regex = re.compile(r"\{\{.*cache_key_wrapper\(.*\).*\}\}")
         templatable_statements: List[str] = []
@@ -1159,12 +1163,19 @@ def has_extra_cache_keys(self, query_obj: Dict) -> bool:
                 return True
         return False
 
-    def get_extra_cache_keys(self, query_obj: Dict) -> List[Any]:
-        if self.has_extra_cache_keys(query_obj):
+    def get_extra_cache_keys(self, query_obj: Dict[str, Any]) -> List[Hashable]:
+        """
+        The cache key of a SqlaTable needs to consider any keys added by the parent class
+        and any keys added via `cache_key_wrapper`.
+
+        :param query_obj: query object to analyze
+        :return: True if at least one item calls `cache_key_wrapper`, otherwise False
+        """
+        extra_cache_keys = super().get_extra_cache_keys(query_obj)
+        if self.has_calls_to_cache_key_wrapper(query_obj):
             sqla_query = self.get_sqla_query(**query_obj)
-            extra_cache_keys = sqla_query.extra_cache_keys
-            return extra_cache_keys
-        return []
+            extra_cache_keys += sqla_query.extra_cache_keys
+        return extra_cache_keys
 
 
 sa.event.listen(SqlaTable, "after_insert", security_manager.set_perm)

superset/dataframe.py

@@ -68,7 +68,7 @@ def is_numeric(dtype):
     return np.issubdtype(dtype, np.number)
 
 
-class SupersetDataFrame(object):
+class SupersetDataFrame:
     # Mapping numpy dtype.char to generic database types
     type_map = {
         "b": "BOOL",  # boolean

superset/db_engine_specs/base.py

@@ -104,7 +104,7 @@ def compile_timegrain_expression(
     return element.name.replace("{col}", compiler.process(element.col, **kw))
 
 
-class LimitMethod(object):  # pylint: disable=too-few-public-methods
+class LimitMethod:  # pylint: disable=too-few-public-methods
     """Enum the ways that limits can be applied"""
 
     FETCH_MANY = "fetch_many"

superset/migrations/versions/258b5280a45e_form_strip_leading_and_trailing_whitespace.py

@@ -33,20 +33,20 @@
 Base = declarative_base()
 
 
-class BaseColumnMixin(object):
+class BaseColumnMixin:
     id = Column(Integer, primary_key=True)
     column_name = Column(String(255))
     description = Column(Text)
     type = Column(String(32))
     verbose_name = Column(String(1024))
 
 
-class BaseDatasourceMixin(object):
+class BaseDatasourceMixin:
     id = Column(Integer, primary_key=True)
     description = Column(Text)
 
 
-class BaseMetricMixin(object):
+class BaseMetricMixin:
     id = Column(Integer, primary_key=True)
     d3format = Column(String(128))
     description = Column(Text)

superset/migrations/versions/c617da68de7d_form_nullable.py

@@ -36,20 +36,20 @@
 Base = declarative_base()
 
 
-class BaseColumnMixin(object):
+class BaseColumnMixin:
     id = Column(Integer, primary_key=True)
     column_name = Column(String(255))
     description = Column(Text)
     type = Column(String(32))
     verbose_name = Column(String(1024))
 
 
-class BaseDatasourceMixin(object):
+class BaseDatasourceMixin:
     id = Column(Integer, primary_key=True)
     description = Column(Text)
 
 
-class BaseMetricMixin(object):
+class BaseMetricMixin:
     id = Column(Integer, primary_key=True)
     d3format = Column(String(128))
     description = Column(Text)

superset/migrations/versions/d94d33dbe938_form_strip.py

@@ -36,20 +36,20 @@
 Base = declarative_base()
 
 
-class BaseColumnMixin(object):
+class BaseColumnMixin:
     id = Column(Integer, primary_key=True)
     column_name = Column(String(255))
     description = Column(Text)
     type = Column(String(32))
     verbose_name = Column(String(1024))
 
 
-class BaseDatasourceMixin(object):
+class BaseDatasourceMixin:
     id = Column(Integer, primary_key=True)
     description = Column(Text)
 
 
-class BaseMetricMixin(object):
+class BaseMetricMixin:
     id = Column(Integer, primary_key=True)
     d3format = Column(String(128))
     description = Column(Text)