[datasets] Add strict type annotation (#9437)

* [datasets] Add strict type annotation

* Fix refresh endpoint

* Improve logic on update

setup.cfg

@@ -53,7 +53,7 @@ order_by_type = false
 ignore_missing_imports = true
 no_implicit_optional = true
 
-[mypy-superset.bin.*,superset.charts.*,superset.dashboards.*,superset.commands.*,superset.common.*,superset.dao.*,superset.db_engine_specs.*,superset.db_engines.*,superset.examples.*]
+[mypy-superset.bin.*,superset.charts.*,superset.datasets.*,superset.dashboards.*,superset.commands.*,superset.common.*,superset.dao.*,superset.db_engine_specs.*,superset.db_engines.*,superset.examples.*]
 check_untyped_defs = true
 disallow_untyped_calls = true
 disallow_untyped_defs = true

superset/datasets/api.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 import logging
+from typing import Any
 
 import yaml
 from flask import g, request, Response
@@ -288,7 +289,7 @@ class DatasetRestApi(BaseSupersetModelRestApi):
     @protect()
     @safe
     @rison(get_export_ids_schema)
-    def export(self, **kwargs):
+    def export(self, **kwargs: Any) -> Response:
         """Export dashboards
         ---
         get:

superset/datasets/commands/create.py

@@ -17,6 +17,7 @@
 import logging
 from typing import Dict, List, Optional
 
+from flask_appbuilder.models.sqla import Model
 from flask_appbuilder.security.sqla.models import User
 from marshmallow import ValidationError
 from sqlalchemy.exc import SQLAlchemyError
@@ -42,7 +43,7 @@ class CreateDatasetCommand(BaseCommand):
         self._actor = user
         self._properties = data.copy()
 
-    def run(self):
+    def run(self) -> Model:
         self.validate()
         try:
             # Creates SqlaTable (Dataset)

superset/datasets/commands/delete.py

@@ -17,6 +17,7 @@
 import logging
 from typing import Optional
 
+from flask_appbuilder.models.sqla import Model
 from flask_appbuilder.security.sqla.models import User
 from sqlalchemy.exc import SQLAlchemyError
 
@@ -42,7 +43,7 @@ class DeleteDatasetCommand(BaseCommand):
         self._model_id = model_id
         self._model: Optional[SqlaTable] = None
 
-    def run(self):
+    def run(self) -> Model:
         self.validate()
         try:
             dataset = DatasetDAO.delete(self._model, commit=False)

superset/datasets/commands/exceptions.py

@@ -33,7 +33,7 @@ class DatabaseNotFoundValidationError(ValidationError):
     Marshmallow validation error for database does not exist
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(_("Database does not exist"), field_names=["database"])
 
 
@@ -42,7 +42,7 @@ class DatabaseChangeValidationError(ValidationError):
     Marshmallow validation error database changes are not allowed on update
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(_("Database not allowed to change"), field_names=["database"])
 
 
@@ -51,7 +51,7 @@ class DatasetExistsValidationError(ValidationError):
     Marshmallow validation error for dataset already exists
     """
 
-    def __init__(self, table_name: str):
+    def __init__(self, table_name: str) -> None:
         super().__init__(
             get_datasource_exist_error_msg(table_name), field_names=["table_name"]
         )
@@ -62,7 +62,7 @@ class DatasetColumnNotFoundValidationError(ValidationError):
     Marshmallow validation error when dataset column for update does not exist
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(_("One or more columns do not exist"), field_names=["columns"])
 
 
@@ -71,7 +71,7 @@ class DatasetColumnsDuplicateValidationError(ValidationError):
     Marshmallow validation error when dataset columns have a duplicate on the list
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             _("One or more columns are duplicated"), field_names=["columns"]
         )
@@ -82,7 +82,7 @@ class DatasetColumnsExistsValidationError(ValidationError):
     Marshmallow validation error when dataset columns already exist
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             _("One or more columns already exist"), field_names=["columns"]
         )
@@ -93,7 +93,7 @@ class DatasetMetricsNotFoundValidationError(ValidationError):
     Marshmallow validation error when dataset metric for update does not exist
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(_("One or more metrics do not exist"), field_names=["metrics"])
 
 
@@ -102,7 +102,7 @@ class DatasetMetricsDuplicateValidationError(ValidationError):
     Marshmallow validation error when dataset metrics have a duplicate on the list
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             _("One or more metrics are duplicated"), field_names=["metrics"]
         )
@@ -113,7 +113,7 @@ class DatasetMetricsExistsValidationError(ValidationError):
     Marshmallow validation error when dataset metrics already exist
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(
             _("One or more metrics already exist"), field_names=["metrics"]
         )
@@ -124,7 +124,7 @@ class TableNotFoundValidationError(ValidationError):
     Marshmallow validation error when a table does not exist on the database
     """
 
-    def __init__(self, table_name: str):
+    def __init__(self, table_name: str) -> None:
         super().__init__(
             _(
                 f"Table [{table_name}] could not be found, "
@@ -137,7 +137,7 @@ class TableNotFoundValidationError(ValidationError):
 
 
 class OwnersNotFoundValidationError(ValidationError):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(_("Owners are invalid"), field_names=["owners"])
 
 

superset/datasets/commands/refresh.py

@@ -17,6 +17,7 @@
 import logging
 from typing import Optional
 
+from flask_appbuilder.models.sqla import Model
 from flask_appbuilder.security.sqla.models import User
 
 from superset.commands.base import BaseCommand
@@ -39,15 +40,16 @@ class RefreshDatasetCommand(BaseCommand):
         self._model_id = model_id
         self._model: Optional[SqlaTable] = None
 
-    def run(self):
+    def run(self) -> Model:
         self.validate()
-        try:
-            # Updates columns and metrics from the dataset
-            self._model.fetch_metadata()
-        except Exception as e:
-            logger.exception(e)
-            raise DatasetRefreshFailedError()
-        return self._model
+        if self._model:
+            try:
+                self._model.fetch_metadata()
+                return self._model
+            except Exception as e:
+                logger.exception(e)
+                raise DatasetRefreshFailedError()
+        raise DatasetRefreshFailedError()
 
     def validate(self) -> None:
         # Validate/populate model exists

superset/datasets/commands/update.py

@@ -18,6 +18,7 @@ import logging
 from collections import Counter
 from typing import Dict, List, Optional
 
+from flask_appbuilder.models.sqla import Model
 from flask_appbuilder.security.sqla.models import User
 from marshmallow import ValidationError
 
@@ -53,14 +54,16 @@ class UpdateDatasetCommand(BaseCommand):
         self._properties = data.copy()
         self._model: Optional[SqlaTable] = None
 
-    def run(self):
+    def run(self) -> Model:
         self.validate()
-        try:
-            dataset = DatasetDAO.update(self._model, self._properties)
-        except DAOUpdateFailedError as e:
-            logger.exception(e.exception)
-            raise DatasetUpdateFailedError()
-        return dataset
+        if self._model:
+            try:
+                dataset = DatasetDAO.update(self._model, self._properties)
+                return dataset
+            except DAOUpdateFailedError as e:
+                logger.exception(e.exception)
+                raise DatasetUpdateFailedError()
+        raise DatasetUpdateFailedError()
 
     def validate(self) -> None:
         exceptions = list()
@@ -107,7 +110,9 @@ class UpdateDatasetCommand(BaseCommand):
             exception.add_list(exceptions)
             raise exception
 
-    def _validate_columns(self, columns: List[Dict], exceptions: List[ValidationError]):
+    def _validate_columns(
+        self, columns: List[Dict], exceptions: List[ValidationError]
+    ) -> None:
         # Validate duplicates on data
         if self._get_duplicates(columns, "column_name"):
             exceptions.append(DatasetColumnsDuplicateValidationError())
@@ -127,7 +132,9 @@ class UpdateDatasetCommand(BaseCommand):
             ):
                 exceptions.append(DatasetColumnsExistsValidationError())
 
-    def _validate_metrics(self, metrics: List[Dict], exceptions: List[ValidationError]):
+    def _validate_metrics(
+        self, metrics: List[Dict], exceptions: List[ValidationError]
+    ) -> None:
         if self._get_duplicates(metrics, "metric_name"):
             exceptions.append(DatasetMetricsDuplicateValidationError())
         else:
@@ -145,7 +152,7 @@ class UpdateDatasetCommand(BaseCommand):
                 exceptions.append(DatasetMetricsExistsValidationError())
 
     @staticmethod
-    def _get_duplicates(data: List[Dict], key: str):
+    def _get_duplicates(data: List[Dict], key: str) -> List[str]:
         duplicates = [
             name
             for name, count in Counter([item[key] for item in data]).items()

superset/datasets/dao.py

@@ -42,7 +42,7 @@ class DatasetDAO(BaseDAO):
         )
 
     @staticmethod
-    def get_database_by_id(database_id) -> Optional[Database]:
+    def get_database_by_id(database_id: int) -> Optional[Database]:
         try:
             return db.session.query(Database).filter_by(id=database_id).one_or_none()
         except SQLAlchemyError as e:  # pragma: no cover
@@ -116,7 +116,7 @@ class DatasetDAO(BaseDAO):
 
     @classmethod
     def update(
-        cls, model: SqlaTable, properties: Dict, commit=True
+        cls, model: SqlaTable, properties: Dict, commit: bool = True
     ) -> Optional[SqlaTable]:
         """
         Updates a Dataset model on the metadata DB
@@ -151,12 +151,14 @@ class DatasetDAO(BaseDAO):
 
     @classmethod
     def update_column(
-        cls, model: TableColumn, properties: Dict, commit=True
+        cls, model: TableColumn, properties: Dict, commit: bool = True
     ) -> Optional[TableColumn]:
         return DatasetColumnDAO.update(model, properties, commit=commit)
 
     @classmethod
-    def create_column(cls, properties: Dict, commit=True) -> Optional[TableColumn]:
+    def create_column(
+        cls, properties: Dict, commit: bool = True
+    ) -> Optional[TableColumn]:
         """
         Creates a Dataset model on the metadata DB
         """
@@ -164,12 +166,14 @@ class DatasetDAO(BaseDAO):
 
     @classmethod
     def update_metric(
-        cls, model: SqlMetric, properties: Dict, commit=True
+        cls, model: SqlMetric, properties: Dict, commit: bool = True
     ) -> Optional[SqlMetric]:
         return DatasetMetricDAO.update(model, properties, commit=commit)
 
     @classmethod
-    def create_metric(cls, properties: Dict, commit=True) -> Optional[SqlMetric]:
+    def create_metric(
+        cls, properties: Dict, commit: bool = True
+    ) -> Optional[SqlMetric]:
         """
         Creates a Dataset model on the metadata DB
         """

superset/datasets/schemas.py

@@ -23,7 +23,7 @@ from marshmallow.validate import Length
 get_export_ids_schema = {"type": "array", "items": {"type": "integer"}}
 
 
-def validate_python_date_format(value):
+def validate_python_date_format(value: str) -> None:
     regex = re.compile(
         r"""
         ^(

superset/views/base.py

@@ -144,7 +144,7 @@ def handle_api_exception(f):
     return functools.update_wrapper(wraps, f)
 
 
-def get_datasource_exist_error_msg(full_name):
+def get_datasource_exist_error_msg(full_name: str) -> str:
     return __("Datasource %(name)s already exists", name=full_name)