Skip to content

[datasets] Add strict type annotation #9437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dpgaspar merged 4 commits into from
Apr 7, 2020
Merged

[datasets] Add strict type annotation #9437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
46b2091
[datasets] Add strict type annotation
dpgaspar Apr 1, 2020
10c318f
Merge remote-tracking branch 'upstream/master' into fix/mypy-datasets
dpgaspar Apr 7, 2020
d29d011
Fix refresh endpoint
dpgaspar Apr 7, 2020
4586a91
Improve logic on update
dpgaspar Apr 7, 2020
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion setup.cfg
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
3 changes: 2 additions & 1 deletion superset/datasets/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
3 changes: 2 additions & 1 deletion superset/datasets/commands/create.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand Down
3 changes: 2 additions & 1 deletion superset/datasets/commands/delete.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
Expand Down
22 changes: 11 additions & 11 deletions superset/datasets/commands/exceptions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"])


Expand All @@ -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"])


Expand All @@ -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"]
)
Expand All @@ -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"])


Expand All @@ -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"]
)
Expand All @@ -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"]
)
Expand All @@ -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"])


Expand All @@ -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"]
)
Expand All @@ -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"]
)
Expand All @@ -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, "
Expand All @@ -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"])


Expand Down
18 changes: 10 additions & 8 deletions superset/datasets/commands/refresh.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand Down
27 changes: 17 additions & 10 deletions superset/datasets/commands/update.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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())
Expand All @@ -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:
Expand All @@ -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()
Expand Down
16 changes: 10 additions & 6 deletions superset/datasets/dao.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -151,25 +151,29 @@ 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
"""
return DatasetColumnDAO.create(properties, commit=commit)

@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
"""
Expand Down
2 changes: 1 addition & 1 deletion superset/datasets/schemas.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
^(
Expand Down
2 changes: 1 addition & 1 deletion superset/views/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down