refactor: sqllab: move sqllab ralated enumns and utils to more logical place

superset/common/db_query_status.py

@@ -0,0 +1,30 @@
+# 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 QueryStatus(str, Enum):
+    """Enum-type class for query statuses"""
+
+    STOPPED: str = "stopped"
+    FAILED: str = "failed"
+    PENDING: str = "pending"
+    RUNNING: str = "running"
+    SCHEDULED: str = "scheduled"
+    SUCCESS: str = "success"
+    FETCHING: str = "fetching"
+    TIMED_OUT: str = "timed_out"

superset/common/query_actions.py

@@ -20,6 +20,7 @@
 from flask_babel import _
 
 from superset import app
+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 (
@@ -28,7 +29,6 @@
     extract_dataframe_dtypes,
     ExtraFiltersReasonType,
     get_time_filter_status,
-    QueryStatus,
 )
 
 if TYPE_CHECKING:

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.db_query_status import QueryStatus
 from superset.common.query_actions import get_query_results
 from superset.common.query_object import QueryObject
 from superset.common.utils import QueryCacheManager
@@ -49,7 +50,6 @@
     get_column_names_from_metrics,
     get_metric_names,
     normalize_dttm_col,
-    QueryStatus,
     TIME_COMPARISION,
 )
 from superset.utils.date_parser import get_past_or_future, normalize_time_delta

superset/common/utils.py

@@ -21,13 +21,14 @@
 from pandas import DataFrame
 
 from superset import app
+from superset.common.db_query_status import QueryStatus
 from superset.constants import CacheRegion
 from superset.exceptions import CacheLoadError
 from superset.extensions import cache_manager
 from superset.models.helpers import QueryResult
 from superset.stats_logger import BaseStatsLogger
 from superset.utils.cache import set_and_log_cache
-from superset.utils.core import error_msg_from_exception, get_stacktrace, QueryStatus
+from superset.utils.core import error_msg_from_exception, get_stacktrace
 
 config = app.config
 stats_logger: BaseStatsLogger = config["STATS_LOGGER"]

superset/connectors/sqla/models.py

@@ -70,6 +70,7 @@
 from sqlalchemy.sql.selectable import Alias, TableClause
 
 from superset import app, db, is_feature_enabled, security_manager
+from superset.common.db_query_status import QueryStatus
 from superset.connectors.base.models import BaseColumn, BaseDatasource, BaseMetric
 from superset.connectors.sqla.utils import (
     get_physical_table_metadata,
@@ -151,12 +152,12 @@ def query(self, query_obj: QueryObjectDict) -> QueryResult:
             qry = qry.filter(Annotation.start_dttm >= query_obj["from_dttm"])
         if query_obj["to_dttm"]:
             qry = qry.filter(Annotation.end_dttm <= query_obj["to_dttm"])
-        status = utils.QueryStatus.SUCCESS
+        status = QueryStatus.SUCCESS
         try:
             df = pd.read_sql_query(qry.statement, db.engine)
         except Exception as ex:  # pylint: disable=broad-except
             df = pd.DataFrame()
-            status = utils.QueryStatus.FAILED
+            status = QueryStatus.FAILED
             logger.exception(ex)
             error_message = utils.error_msg_from_exception(ex)
         return QueryResult(
@@ -1444,7 +1445,7 @@ def query(self, query_obj: QueryObjectDict) -> QueryResult:
         qry_start_dttm = datetime.now()
         query_str_ext = self.get_query_str_extended(query_obj)
         sql = query_str_ext.sql
-        status = utils.QueryStatus.SUCCESS
+        status = QueryStatus.SUCCESS
         errors = None
         error_message = None
 
@@ -1477,7 +1478,7 @@ def assign_column_label(df: pd.DataFrame) -> Optional[pd.DataFrame]:
             df = self.database.get_df(sql, self.schema, mutator=assign_column_label)
         except Exception as ex:  # pylint: disable=broad-except
             df = pd.DataFrame()
-            status = utils.QueryStatus.FAILED
+            status = QueryStatus.FAILED
             logger.warning(
                 "Query %s on schema %s failed", sql, self.schema, exc_info=True
             )

superset/db_engine_specs/base.py

@@ -82,8 +82,6 @@ class TimeGrain(NamedTuple):
     duration: Optional[str]
 
 
-QueryStatus = utils.QueryStatus
-
 builtin_time_grains: Dict[Optional[str], str] = {
     None: __("Original value"),
     "PT1S": __("Second"),

superset/db_engine_specs/druid.py

@@ -31,7 +31,7 @@
 logger = logging.getLogger()
 
 
-class DruidEngineSpec(BaseEngineSpec):  # pylint: disable=abstract-method
+class DruidEngineSpec(BaseEngineSpec):
     """Engine spec for Druid.io"""
 
     engine = "druid"

superset/db_engine_specs/hive.py

@@ -35,6 +35,7 @@
 from sqlalchemy.orm import Session
 from sqlalchemy.sql.expression import ColumnClause, Select
 
+from superset.common.db_query_status import QueryStatus
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.presto import PrestoEngineSpec
 from superset.exceptions import SupersetException
@@ -48,7 +49,6 @@
     from superset.models.core import Database
 
 
-QueryStatus = utils.QueryStatus
 logger = logging.getLogger(__name__)
 
 

superset/db_engine_specs/presto.py

@@ -52,6 +52,7 @@
 from sqlalchemy.types import TypeEngine
 
 from superset import cache_manager, is_feature_enabled
+from superset.common.db_query_status import QueryStatus
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.errors import SupersetErrorType
 from superset.exceptions import SupersetTemplateException
@@ -95,7 +96,6 @@
 )
 
 
-QueryStatus = utils.QueryStatus
 logger = logging.getLogger(__name__)
 
 

superset/models/helpers.py

@@ -38,7 +38,7 @@
 from sqlalchemy.orm.exc import MultipleResultsFound
 from sqlalchemy_utils import UUIDType
 
-from superset.utils.core import QueryStatus
+from superset.common.db_query_status import QueryStatus
 
 logger = logging.getLogger(__name__)
 

superset/models/sql_lab.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 """A collection of ORM sqlalchemy models for SQL Lab"""
-import enum
 import re
 from datetime import datetime
 from typing import Any, Dict, List
@@ -48,17 +47,10 @@
 )
 from superset.models.tags import QueryUpdater
 from superset.sql_parse import CtasMethod, ParsedQuery, Table
+from superset.sqllab.limiting_factor import LimitingFactor
 from superset.utils.core import QueryStatus, user_label
 
 
-class LimitingFactor(str, enum.Enum):
-    QUERY = "QUERY"
-    DROPDOWN = "DROPDOWN"
-    QUERY_AND_DROPDOWN = "QUERY_AND_DROPDOWN"
-    NOT_LIMITED = "NOT_LIMITED"
-    UNKNOWN = "UNKNOWN"
-
-
 class Query(Model, ExtraJSONMixin):
     """ORM model for SQL query
 

superset/sql_lab.py

@@ -32,21 +32,18 @@
 from sqlalchemy.orm import Session
 
 from superset import app, results_backend, results_backend_use_msgpack, security_manager
+from superset.common.db_query_status import QueryStatus
 from superset.dataframe import df_to_records
 from superset.db_engine_specs import BaseEngineSpec
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import SupersetErrorException, SupersetErrorsException
 from superset.extensions import celery_app
-from superset.models.sql_lab import LimitingFactor, Query
+from superset.models.sql_lab import Query
 from superset.result_set import SupersetResultSet
 from superset.sql_parse import CtasMethod, ParsedQuery
+from superset.sqllab.limiting_factor import LimitingFactor
 from superset.utils.celery import session_scope
-from superset.utils.core import (
-    json_iso_dttm_ser,
-    QuerySource,
-    QueryStatus,
-    zlib_compress,
-)
+from superset.utils.core import json_iso_dttm_ser, QuerySource, zlib_compress
 from superset.utils.dates import now_as_float
 from superset.utils.decorators import stats_timing
 

superset/sqllab/command.py

@@ -31,6 +31,7 @@
 
 from superset import app, db, is_feature_enabled, sql_lab
 from superset.commands.base import BaseCommand
+from superset.common.db_query_status import QueryStatus
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import (
     SupersetErrorException,
@@ -43,16 +44,16 @@
 )
 from superset.jinja_context import BaseTemplateProcessor, get_template_processor
 from superset.models.core import Database
-from superset.models.sql_lab import LimitingFactor, Query
+from superset.models.sql_lab import Query
 from superset.queries.dao import QueryDAO
 from superset.sqllab.command_status import SqlJsonExecutionStatus
+from superset.sqllab.limiting_factor import LimitingFactor
+from superset.sqllab.utils import apply_display_max_row_configuration_if_require
 from superset.utils import core as utils
 from superset.utils.dates import now_as_float
 from superset.utils.sqllab_execution_context import SqlJsonExecutionContext
-from superset.views.utils import apply_display_max_row_limit
 
 config = app.config
-QueryStatus = utils.QueryStatus
 logger = logging.getLogger(__name__)
 
 PARAMETER_MISSING_ERR = (
@@ -397,7 +398,9 @@ def _to_payload_results_based(  # pylint: disable=no-self-use
     ) -> str:
         display_max_row = config["DISPLAY_MAX_ROW"]
         return json.dumps(
-            apply_display_max_row_limit(execution_result, display_max_row),
+            apply_display_max_row_configuration_if_require(
+                execution_result, display_max_row
+            ),
             default=utils.pessimistic_json_iso_dttm_ser,
             ignore_nan=True,
             encoding=None,

superset/sqllab/limiting_factor.py

@@ -0,0 +1,25 @@
+# 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 enum
+
+
+class LimitingFactor(str, enum.Enum):
+    QUERY = "QUERY"
+    DROPDOWN = "DROPDOWN"
+    QUERY_AND_DROPDOWN = "QUERY_AND_DROPDOWN"
+    NOT_LIMITED = "NOT_LIMITED"
+    UNKNOWN = "UNKNOWN"

superset/sqllab/utils.py

@@ -0,0 +1,47 @@
+# 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 typing import Any, Dict
+
+from superset.common.db_query_status import QueryStatus
+
+
+def apply_display_max_row_configuration_if_require(  # pylint: disable=invalid-name
+    sql_results: Dict[str, Any], max_rows_in_result: int
+) -> Dict[str, Any]:
+    """
+    Given a `sql_results` nested structure, applies a limit to the number of rows
+
+    `sql_results` here is the nested structure coming out of sql_lab.get_sql_results, it
+    contains metadata about the query, as well as the data set returned by the query.
+    This method limits the number of rows adds a `displayLimitReached: True` flag to the
+    metadata.
+
+    :param max_rows_in_result:
+    :param sql_results: The results of a sql query from sql_lab.get_sql_results
+    :returns: The mutated sql_results structure
+    """
+
+    def is_require_to_apply() -> bool:
+        return (
+            sql_results["status"] == QueryStatus.SUCCESS
+            and sql_results["query"]["rows"] > max_rows_in_result
+        )
+
+    if is_require_to_apply():
+        sql_results["data"] = sql_results["data"][:max_rows_in_result]
+        sql_results["displayLimitReached"] = True
+    return sql_results

superset/utils/sqllab_execution_context.py

@@ -34,7 +34,6 @@
 if TYPE_CHECKING:
     from superset.connectors.sqla.models import Database
 
-QueryStatus = utils.QueryStatus
 logger = logging.getLogger(__name__)
 
 SqlResults = Dict[str, Any]