refactor(chart.commands): separate commands into two different modules (#17509)

refactor: move imports under TYPE_CHECKING

superset/charts/data/api.py

@@ -32,10 +32,10 @@ from superset.charts.commands.exceptions import (
     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

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.

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

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.charts.commands.exceptions import (
 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 @@ class ChartDataCommand(BaseCommand):
 
     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

superset/charts/schemas.py

@@ -26,7 +26,6 @@ from marshmallow_enum import EnumField
 
 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 @@ from superset.utils.core import (
 
 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 @@ class ChartDataQueryContextSchema(Schema):
 
     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
 

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 @@ from superset.commands.exceptions import (
     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,

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"))

tests/integration_tests/charts/data/api_tests.py

@@ -37,7 +37,7 @@ from tests.integration_tests.test_app import app
 
 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

tests/integration_tests/tasks/async_queries_tests.py

@@ -23,7 +23,7 @@ from celery.exceptions import SoftTimeLimitExceeded
 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