From 16e012fe5d791ed6aa5ab755a6386b1581fc5ebd Mon Sep 17 00:00:00 2001 From: ofekisr <35701650+ofekisr@users.noreply.github.com> Date: Mon, 22 Nov 2021 19:57:52 +0200 Subject: [PATCH] refactor(chart.commands): separate commands into two different modules (#17509) refactor: move imports under TYPE_CHECKING --- superset/charts/data/api.py | 4 +- superset/charts/data/commands/__init__.py | 16 ++++++++ .../data/commands/create_async_job_command.py | 38 +++++++++++++++++++ .../get_data_command.py} | 19 +--------- superset/charts/schemas.py | 5 ++- superset/commands/utils.py | 8 +++- superset/tasks/async_queries.py | 2 +- .../charts/data/api_tests.py | 2 +- .../tasks/async_queries_tests.py | 2 +- 9 files changed, 70 insertions(+), 26 deletions(-) create mode 100644 superset/charts/data/commands/__init__.py create mode 100644 superset/charts/data/commands/create_async_job_command.py rename superset/charts/data/{commands.py => commands/get_data_command.py} (77%) diff --git a/superset/charts/data/api.py b/superset/charts/data/api.py index ffa04abae4c46..6d964091ad553 100644 --- a/superset/charts/data/api.py +++ b/superset/charts/data/api.py @@ -32,10 +32,10 @@ ChartDataCacheLoadError, ChartDataQueryFailedError, ) -from superset.charts.data.commands import ( - ChartDataCommand, +from superset.charts.data.commands.create_async_job_command import ( CreateAsyncChartDataJobCommand, ) +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.charts.data.query_context_cache_loader import QueryContextCacheLoader from superset.charts.post_processing import apply_post_process from superset.charts.schemas import ChartDataQueryContextSchema diff --git a/superset/charts/data/commands/__init__.py b/superset/charts/data/commands/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/superset/charts/data/commands/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/superset/charts/data/commands/create_async_job_command.py b/superset/charts/data/commands/create_async_job_command.py new file mode 100644 index 0000000000000..98a0174e6252f --- /dev/null +++ b/superset/charts/data/commands/create_async_job_command.py @@ -0,0 +1,38 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +import logging +from typing import Any, Dict, Optional + +from flask import Request + +from superset.extensions import async_query_manager +from superset.tasks.async_queries import load_chart_data_into_cache + +logger = logging.getLogger(__name__) + + +class CreateAsyncChartDataJobCommand: + _async_channel_id: str + + def validate(self, request: Request) -> None: + jwt_data = async_query_manager.parse_jwt_from_request(request) + self._async_channel_id = jwt_data["channel"] + + def run(self, form_data: Dict[str, Any], user_id: Optional[str]) -> Dict[str, Any]: + job_metadata = async_query_manager.init_job(self._async_channel_id, user_id) + load_chart_data_into_cache.delay(job_metadata, form_data) + return job_metadata diff --git a/superset/charts/data/commands.py b/superset/charts/data/commands/get_data_command.py similarity index 77% rename from superset/charts/data/commands.py rename to superset/charts/data/commands/get_data_command.py index 3fc02e260f350..95f7513f253f6 100644 --- a/superset/charts/data/commands.py +++ b/superset/charts/data/commands/get_data_command.py @@ -15,9 +15,7 @@ # specific language governing permissions and limitations # under the License. import logging -from typing import Any, Dict, Optional - -from flask import Request +from typing import Any, Dict from superset.charts.commands.exceptions import ( ChartDataCacheLoadError, @@ -26,8 +24,6 @@ from superset.commands.base import BaseCommand from superset.common.query_context import QueryContext from superset.exceptions import CacheLoadError -from superset.extensions import async_query_manager -from superset.tasks.async_queries import load_chart_data_into_cache logger = logging.getLogger(__name__) @@ -66,16 +62,3 @@ def run(self, **kwargs: Any) -> Dict[str, Any]: def validate(self) -> None: self._query_context.raise_for_access() - - -class CreateAsyncChartDataJobCommand: - _async_channel_id: str - - def validate(self, request: Request) -> None: - jwt_data = async_query_manager.parse_jwt_from_request(request) - self._async_channel_id = jwt_data["channel"] - - def run(self, form_data: Dict[str, Any], user_id: Optional[str]) -> Dict[str, Any]: - job_metadata = async_query_manager.init_job(self._async_channel_id, user_id) - load_chart_data_into_cache.delay(job_metadata, form_data) - return job_metadata diff --git a/superset/charts/schemas.py b/superset/charts/schemas.py index 46aa9a48af0e5..51411318e2474 100644 --- a/superset/charts/schemas.py +++ b/superset/charts/schemas.py @@ -26,7 +26,6 @@ from superset import app from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType -from superset.common.query_context_factory import QueryContextFactory from superset.db_engine_specs.base import builtin_time_grains from superset.utils import schema as utils from superset.utils.core import ( @@ -39,6 +38,7 @@ if TYPE_CHECKING: from superset.common.query_context import QueryContext + from superset.common.query_context_factory import QueryContextFactory config = app.config @@ -1153,6 +1153,9 @@ def make_query_context(self, data: Dict[str, Any], **kwargs: Any) -> QueryContex def get_query_context_factory(self) -> QueryContextFactory: if self.query_context_factory is None: + # pylint: disable=import-outside-toplevel + from superset.common.query_context_factory import QueryContextFactory + self.query_context_factory = QueryContextFactory() return self.query_context_factory diff --git a/superset/commands/utils.py b/superset/commands/utils.py index 070039874eb17..c68dde3d5a758 100644 --- a/superset/commands/utils.py +++ b/superset/commands/utils.py @@ -14,7 +14,9 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -from typing import List, Optional +from __future__ import annotations + +from typing import List, Optional, TYPE_CHECKING from flask_appbuilder.security.sqla.models import Role, User @@ -23,11 +25,13 @@ OwnersNotFoundValidationError, RolesNotFoundValidationError, ) -from superset.connectors.base.models import BaseDatasource from superset.connectors.connector_registry import ConnectorRegistry from superset.datasets.commands.exceptions import DatasetNotFoundError from superset.extensions import db, security_manager +if TYPE_CHECKING: + from superset.connectors.base.models import BaseDatasource + def populate_owners( user: User, owner_ids: Optional[List[int]], default_to_user: bool, diff --git a/superset/tasks/async_queries.py b/superset/tasks/async_queries.py index e916028b12a86..fcd6f91ebf754 100644 --- a/superset/tasks/async_queries.py +++ b/superset/tasks/async_queries.py @@ -70,7 +70,7 @@ def load_chart_data_into_cache( job_metadata: Dict[str, Any], form_data: Dict[str, Any], ) -> None: # pylint: disable=import-outside-toplevel - from superset.charts.data.commands import ChartDataCommand + from superset.charts.data.commands.get_data_command import ChartDataCommand try: ensure_user_is_set(job_metadata.get("user_id")) diff --git a/tests/integration_tests/charts/data/api_tests.py b/tests/integration_tests/charts/data/api_tests.py index 9a93cb479b5e1..12d667f1611b8 100644 --- a/tests/integration_tests/charts/data/api_tests.py +++ b/tests/integration_tests/charts/data/api_tests.py @@ -37,7 +37,7 @@ import pytest -from superset.charts.data.commands import ChartDataCommand +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.connectors.sqla.models import TableColumn, SqlaTable from superset.errors import SupersetErrorType from superset.extensions import async_query_manager, db diff --git a/tests/integration_tests/tasks/async_queries_tests.py b/tests/integration_tests/tasks/async_queries_tests.py index e2cf21c552624..6acdd88b78838 100644 --- a/tests/integration_tests/tasks/async_queries_tests.py +++ b/tests/integration_tests/tasks/async_queries_tests.py @@ -23,7 +23,7 @@ from flask import g from superset.charts.commands.exceptions import ChartDataQueryFailedError -from superset.charts.data.commands import ChartDataCommand +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.exceptions import SupersetException from superset.extensions import async_query_manager, security_manager from superset.tasks import async_queries