diff --git a/setup.cfg b/setup.cfg index bff4adf2ac2a..b535c63a76f8 100644 --- a/setup.cfg +++ b/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 diff --git a/superset/datasets/api.py b/superset/datasets/api.py index cab625b16995..c12f39024857 100644 --- a/superset/datasets/api.py +++ b/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 @@ def delete(self, pk: int) -> Response: # pylint: disable=arguments-differ @protect() @safe @rison(get_export_ids_schema) - def export(self, **kwargs): + def export(self, **kwargs: Any) -> Response: """Export dashboards --- get: diff --git a/superset/datasets/commands/create.py b/superset/datasets/commands/create.py index 71d622569e3a..9200a1371e22 100644 --- a/superset/datasets/commands/create.py +++ b/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 @@ def __init__(self, user: User, data: Dict): self._actor = user self._properties = data.copy() - def run(self): + def run(self) -> Model: self.validate() try: # Creates SqlaTable (Dataset) diff --git a/superset/datasets/commands/delete.py b/superset/datasets/commands/delete.py index 668d63a90248..4ee3ed76c1c2 100644 --- a/superset/datasets/commands/delete.py +++ b/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 @@ def __init__(self, user: User, model_id: int): 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) diff --git a/superset/datasets/commands/exceptions.py b/superset/datasets/commands/exceptions.py index 35c6106da35b..e5f3ae292801 100644 --- a/superset/datasets/commands/exceptions.py +++ b/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 @@ def __init__(self, table_name: str): class OwnersNotFoundValidationError(ValidationError): - def __init__(self): + def __init__(self) -> None: super().__init__(_("Owners are invalid"), field_names=["owners"]) diff --git a/superset/datasets/commands/refresh.py b/superset/datasets/commands/refresh.py index 2c58cc4b2afc..aed5efc29417 100644 --- a/superset/datasets/commands/refresh.py +++ b/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 @@ def __init__(self, user: User, model_id: int): 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 diff --git a/superset/datasets/commands/update.py b/superset/datasets/commands/update.py index 0b40952836ce..00626b45a7be 100644 --- a/superset/datasets/commands/update.py +++ b/superset/datasets/commands/update.py @@ -18,6 +18,7 @@ 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 @@ def __init__(self, user: User, model_id: int, data: Dict): 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 @@ def validate(self) -> None: 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 @@ def _validate_columns(self, columns: List[Dict], exceptions: List[ValidationErro ): 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 @@ def _validate_metrics(self, metrics: List[Dict], exceptions: List[ValidationErro 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() diff --git a/superset/datasets/dao.py b/superset/datasets/dao.py index 7a35eaed4ce8..ae5d4061c473 100644 --- a/superset/datasets/dao.py +++ b/superset/datasets/dao.py @@ -42,7 +42,7 @@ def get_owner_by_id(owner_id: int) -> Optional[object]: ) @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 @@ def validate_metrics_uniqueness(dataset_id: int, metrics_names: List[str]) -> bo @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 @@ def update( @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 @@ def create_column(cls, properties: Dict, commit=True) -> Optional[TableColumn]: @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 """ diff --git a/superset/datasets/schemas.py b/superset/datasets/schemas.py index 0c1e67586f7a..b27df9328c30 100644 --- a/superset/datasets/schemas.py +++ b/superset/datasets/schemas.py @@ -23,7 +23,7 @@ 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""" ^( diff --git a/superset/views/base.py b/superset/views/base.py index 5f40d34a857d..f169a79f59a2 100644 --- a/superset/views/base.py +++ b/superset/views/base.py @@ -144,7 +144,7 @@ def wraps(self, *args, **kwargs): 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)