refactor(ChartData): move ChartDataResult enums to common

superset/charts/api.py

@@ -67,17 +67,14 @@
 )
 from superset.commands.importers.exceptions import NoValidFilesFoundError
 from superset.commands.importers.v1.utils import get_contents_from_bundle
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.constants import MODEL_API_RW_METHOD_PERMISSION_MAP, RouteMethod
 from superset.exceptions import QueryObjectValidationError
 from superset.extensions import event_logger, security_manager
 from superset.models.slice import Slice
 from superset.tasks.thumbnails import cache_chart_thumbnail
 from superset.utils.async_query_manager import AsyncQueryTokenException
-from superset.utils.core import (
- ChartDataResultFormat,
- ChartDataResultType,
- json_int_dttm_ser,
-)
+from superset.utils.core import json_int_dttm_ser
 from superset.utils.screenshots import ChartScreenshot
 from superset.utils.urls import get_url_path
 from superset.views.base_api import (

superset/charts/post_processing.py

@@ -31,12 +31,8 @@
 
 import pandas as pd
 
-from superset.utils.core import (
- ChartDataResultFormat,
- DTTM_ALIAS,
- extract_dataframe_dtypes,
- get_metric_name,
-)
+from superset.common.chart_data import ChartDataResultFormat
+from superset.utils.core import DTTM_ALIAS, extract_dataframe_dtypes, get_metric_name
 
 
 def get_column_key(label: Tuple[str, ...], metrics: List[str]) -> Tuple[Any, ...]:

superset/charts/schemas.py

@@ -23,13 +23,12 @@
 from marshmallow_enum import EnumField
 
 from superset import app
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.common.query_context import QueryContext
 from superset.db_engine_specs.base import builtin_time_grains
 from superset.utils import schema as utils
 from superset.utils.core import (
  AnnotationType,
- ChartDataResultFormat,
- ChartDataResultType,
  FilterOperator,
  PostProcessingBoxplotWhiskerType,
  PostProcessingContributionOrientation,

superset/common/chart_data.py

@@ -0,0 +1,40 @@
+# 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.
+from enum import Enum
+
+
+class ChartDataResultFormat(str, Enum):
+ """
+ Chart data response format
+ """
+
+ CSV = "csv"
+ JSON = "json"
+
+
+class ChartDataResultType(str, Enum):
+ """
+ Chart data response type
+ """
+
+ COLUMNS = "columns"
+ FULL = "full"
+ QUERY = "query"
+ RESULTS = "results"
+ SAMPLES = "samples"
+ TIMEGRAINS = "timegrains"
+ POST_PROCESSED = "post_processed"

superset/common/query_actions.py

@@ -20,11 +20,11 @@
 from flask_babel import _
 
 from superset import app
+from superset.common.chart_data import ChartDataResultType
 from superset.common.db_query_status import QueryStatus
 from superset.connectors.base.models import BaseDatasource
 from superset.exceptions import QueryObjectValidationError
 from superset.utils.core import (
- ChartDataResultType,
  extract_column_dtype,
  extract_dataframe_dtypes,
  ExtraFiltersReasonType,

superset/common/query_context.py

@@ -29,6 +29,7 @@
 from superset import app, db, is_feature_enabled
 from superset.annotation_layers.dao import AnnotationLayerDAO
 from superset.charts.dao import ChartDAO
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.common.db_query_status import QueryStatus
 from superset.common.query_actions import get_query_results
 from superset.common.query_object import QueryObject
@@ -42,8 +43,6 @@
 from superset.utils import csv
 from superset.utils.cache import generate_cache_key, set_and_log_cache
 from superset.utils.core import (
- ChartDataResultFormat,
- ChartDataResultType,
  DatasourceDict,
  DTTM_ALIAS,
  error_msg_from_exception,

superset/common/query_object.py

@@ -23,14 +23,14 @@
 from pandas import DataFrame
 
 from superset import app, db
+from superset.common.chart_data import ChartDataResultType
 from superset.connectors.base.models import BaseDatasource
 from superset.connectors.connector_registry import ConnectorRegistry
 from superset.exceptions import QueryObjectValidationError
 from superset.typing import Metric, OrderBy
 from superset.utils import pandas_postprocessing
 from superset.utils.core import (
  apply_max_row_limit,
- ChartDataResultType,
  DatasourceDict,
  DTTM_ALIAS,
  find_duplicates,

superset/reports/commands/execute.py

@@ -28,6 +28,7 @@
 from superset import app
 from superset.commands.base import BaseCommand
 from superset.commands.exceptions import CommandException
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.extensions import feature_flag_manager, machine_auth_provider_factory
 from superset.models.reports import (
  ReportDataFormat,
@@ -64,7 +65,6 @@
 from superset.reports.notifications.base import NotificationContent
 from superset.reports.notifications.exceptions import NotificationError
 from superset.utils.celery import session_scope
-from superset.utils.core import ChartDataResultFormat, ChartDataResultType
 from superset.utils.csv import get_chart_csv_data, get_chart_dataframe
 from superset.utils.screenshots import (
  BaseScreenshot,

superset/utils/core.py

@@ -174,29 +174,6 @@ class GenericDataType(IntEnum):
  # ROW = 7
 
 
-class ChartDataResultFormat(str, Enum):
- """
- Chart data response format
- """
-
- CSV = "csv"
- JSON = "json"
-
-
-class ChartDataResultType(str, Enum):
- """
- Chart data response type
- """
-
- COLUMNS = "columns"
- FULL = "full"
- QUERY = "query"
- RESULTS = "results"
- SAMPLES = "samples"
- TIMEGRAINS = "timegrains"
- POST_PROCESSED = "post_processed"
-
-
 class DatasourceDict(TypedDict):
  type: str
  id: int

superset/views/core.py

@@ -60,6 +60,7 @@
  viz,
 )
 from superset.charts.dao import ChartDAO
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.common.db_query_status import QueryStatus
 from superset.connectors.base.models import BaseDatasource
 from superset.connectors.connector_registry import ConnectorRegistry
@@ -459,18 +460,18 @@ def send_data_payload_response(viz_obj: BaseViz, payload: Any) -> FlaskResponse:
  def generate_json(
  self, viz_obj: BaseViz, response_type: Optional[str] = None
  ) -> FlaskResponse:
- if response_type == utils.ChartDataResultFormat.CSV:
+ if response_type == ChartDataResultFormat.CSV:
  return CsvResponse(
  viz_obj.get_csv(), headers=generate_download_headers("csv")
  )
 
- if response_type == utils.ChartDataResultType.QUERY:
+ if response_type == ChartDataResultType.QUERY:
  return self.get_query_string_response(viz_obj)
 
- if response_type == utils.ChartDataResultType.RESULTS:
+ if response_type == ChartDataResultType.RESULTS:
  return self.get_raw_results(viz_obj)
 
- if response_type == utils.ChartDataResultType.SAMPLES:
+ if response_type == ChartDataResultType.SAMPLES:
  return self.get_samples(viz_obj)
 
  payload = viz_obj.get_payload()
@@ -598,19 +599,19 @@ def explore_json(
 
  TODO: break into one endpoint for each return shape"""
 
- response_type = utils.ChartDataResultFormat.JSON.value
- responses: List[
- Union[utils.ChartDataResultFormat, utils.ChartDataResultType]
- ] = list(utils.ChartDataResultFormat)
- responses.extend(list(utils.ChartDataResultType))
+ response_type = ChartDataResultFormat.JSON.value
+ responses: List[Union[ChartDataResultFormat, ChartDataResultType]] = list(
+ ChartDataResultFormat
+ )
+ responses.extend(list(ChartDataResultType))
  for response_option in responses:
  if request.args.get(response_option) == "true":
  response_type = response_option
  break
 
  # Verify user has permission to export CSV file
  if (
- response_type == utils.ChartDataResultFormat.CSV
+ response_type == ChartDataResultFormat.CSV
  and not security_manager.can_access("can_csv", "Superset")
  ):
  return json_error_response(
@@ -628,7 +629,7 @@ def explore_json(
  # TODO: support CSV, SQL query and other non-JSON types
  if (
  is_feature_enabled("GLOBAL_ASYNC_QUERIES")
- and response_type == utils.ChartDataResultFormat.JSON
+ and response_type == ChartDataResultFormat.JSON
  ):
  # First, look for the chart query results in the cache.
  try:

tests/integration_tests/charts/api_tests.py

@@ -51,15 +51,13 @@
 from superset.models.dashboard import Dashboard
 from superset.models.reports import ReportSchedule, ReportScheduleType
 from superset.models.slice import Slice
-from superset.utils import core as utils
 from superset.utils.core import (
  AnnotationType,
- ChartDataResultFormat,
  get_example_database,
  get_example_default_schema,
  get_main_database,
 )
-
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 
 from tests.integration_tests.base_api_tests import ApiOwnersTestCaseMixin
 from tests.integration_tests.base_tests import (
@@ -1239,7 +1237,7 @@ def test_chart_data_sample_default_limit(self):
  """
  self.login(username="admin")
  request_payload = get_query_context("birth_names")
- request_payload["result_type"] = utils.ChartDataResultType.SAMPLES
+ request_payload["result_type"] = ChartDataResultType.SAMPLES
  del request_payload["queries"][0]["row_limit"]
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  response_payload = json.loads(rv.data.decode("utf-8"))
@@ -1258,7 +1256,7 @@ def test_chart_data_sample_custom_limit(self):
  """
  self.login(username="admin")
  request_payload = get_query_context("birth_names")
- request_payload["result_type"] = utils.ChartDataResultType.SAMPLES
+ request_payload["result_type"] = ChartDataResultType.SAMPLES
  request_payload["queries"][0]["row_limit"] = 10
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  response_payload = json.loads(rv.data.decode("utf-8"))
@@ -1276,7 +1274,7 @@ def test_chart_data_sql_max_row_sample_limit(self):
  """
  self.login(username="admin")
  request_payload = get_query_context("birth_names")
- request_payload["result_type"] = utils.ChartDataResultType.SAMPLES
+ request_payload["result_type"] = ChartDataResultType.SAMPLES
  request_payload["queries"][0]["row_limit"] = 10000000
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  response_payload = json.loads(rv.data.decode("utf-8"))
@@ -1326,7 +1324,7 @@ def test_chart_data_query_result_type(self):
  """
  self.login(username="admin")
  request_payload = get_query_context("birth_names")
- request_payload["result_type"] = utils.ChartDataResultType.QUERY
+ request_payload["result_type"] = ChartDataResultType.QUERY
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  self.assertEqual(rv.status_code, 200)
 
@@ -1453,7 +1451,7 @@ def test_chart_data_query_missing_filter(self):
  request_payload["queries"][0]["filters"] = [
  {"col": "non_existent_filter", "op": "==", "val": "foo"},
  ]
- request_payload["result_type"] = utils.ChartDataResultType.QUERY
+ request_payload["result_type"] = ChartDataResultType.QUERY
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  self.assertEqual(rv.status_code, 200)
  response_payload = json.loads(rv.data.decode("utf-8"))
@@ -1532,7 +1530,7 @@ def test_chart_data_jinja_filter_request(self):
  """
  self.login(username="admin")
  request_payload = get_query_context("birth_names")
- request_payload["result_type"] = utils.ChartDataResultType.QUERY
+ request_payload["result_type"] = ChartDataResultType.QUERY
  request_payload["queries"][0]["filters"] = [
  {"col": "gender", "op": "==", "val": "boy"}
  ]
@@ -1574,7 +1572,7 @@ def test_chart_data_async_cached_sync_response(self):
 
  class QueryContext:
  result_format = ChartDataResultFormat.JSON
- result_type = utils.ChartDataResultType.FULL
+ result_type = ChartDataResultType.FULL
 
  cmd_run_val = {
  "query_context": QueryContext(),
@@ -1585,7 +1583,7 @@ class QueryContext:
  ChartDataCommand, "run", return_value=cmd_run_val
  ) as patched_run:
  request_payload = get_query_context("birth_names")
- request_payload["result_type"] = utils.ChartDataResultType.FULL
+ request_payload["result_type"] = ChartDataResultType.FULL
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  self.assertEqual(rv.status_code, 200)
  data = json.loads(rv.data.decode("utf-8"))
@@ -1997,8 +1995,8 @@ def test_chart_data_timegrains(self):
  self.login(username="admin")
  request_payload = get_query_context("birth_names")
  request_payload["queries"] = [
- {"result_type": utils.ChartDataResultType.TIMEGRAINS},
- {"result_type": utils.ChartDataResultType.COLUMNS},
+ {"result_type": ChartDataResultType.TIMEGRAINS},
+ {"result_type": ChartDataResultType.COLUMNS},
  ]
  rv = self.post_assert_metric(CHART_DATA_URI, request_payload, "data")
  response_payload = json.loads(rv.data.decode("utf-8"))

tests/integration_tests/query_context_tests.py

@@ -24,18 +24,13 @@
 
 from superset import db
 from superset.charts.schemas import ChartDataQueryContextSchema
+from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.common.query_context import QueryContext
 from superset.common.query_object import QueryObject
 from superset.connectors.connector_registry import ConnectorRegistry
 from superset.connectors.sqla.models import SqlMetric
 from superset.extensions import cache_manager
-from superset.utils.core import (
- AdhocMetricExpressionType,
- backend,
- ChartDataResultFormat,
- ChartDataResultType,
- TimeRangeEndpoint,
-)
+from superset.utils.core import AdhocMetricExpressionType, backend, TimeRangeEndpoint
 from tests.integration_tests.base_tests import SupersetTestCase
 from tests.integration_tests.fixtures.birth_names_dashboard import (
  load_birth_names_dashboard_with_slices,