Refactor/sql json view14

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

@@ -69,6 +69,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,
@@ -150,12 +151,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/dao/base.py

@@ -89,6 +89,19 @@ def find_all(cls) -> List[Model]:
             ).apply(query, None)
         return query.all()
 
+    @classmethod
+    def find_one_or_none(cls, **filter_by: Any) -> Optional[Model]:
+        """
+        Get the first that fit the `base_filter`
+        """
+        query = db.session.query(cls.model_cls)
+        if cls.base_filter:
+            data_model = SQLAInterface(cls.model_cls, db.session)
+            query = cls.base_filter(  # pylint: disable=not-callable
+                "id", data_model
+            ).apply(query, None)
+        return query.filter_by(**filter_by).one_or_none()
+
     @classmethod
     def create(cls, properties: Dict[str, Any], commit: bool = True) -> Model:
         """
@@ -109,6 +122,29 @@ def create(cls, properties: Dict[str, Any], commit: bool = True) -> Model:
             raise DAOCreateFailedError(exception=ex) from ex
         return model
 
+    @classmethod
+    def save(cls, instance_model: Model, commit: bool = True) -> Model:
+        """
+        Generic for saving models
+        :raises: DAOCreateFailedError
+        """
+        if cls.model_cls is None:
+            raise DAOConfigError()
+        if not isinstance(  # pylint: disable=isinstance-second-argument-not-valid-type
+            instance_model, cls.model_cls
+        ):
+            raise DAOCreateFailedError(
+                "the instance model is not a type of the model class"
+            )
+        try:
+            db.session.add(instance_model)
+            if commit:
+                db.session.commit()
+        except SQLAlchemyError as ex:  # pragma: no cover
+            db.session.rollback()
+            raise DAOCreateFailedError(exception=ex) from ex
+        return instance_model
+
     @classmethod
     def update(
         cls, model: Model, properties: Dict[str, Any], commit: bool = 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/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/exceptions.py

@@ -28,17 +28,25 @@ class SupersetException(Exception):
     message = ""
 
     def __init__(
-        self, message: str = "", exception: Optional[Exception] = None,
+        self,
+        message: str = "",
+        exception: Optional[Exception] = None,
+        error_type: Optional[SupersetErrorType] = None,
     ) -> None:
         if message:
             self.message = message
         self._exception = exception
+        self._error_type = error_type
         super().__init__(self.message)
 
     @property
     def exception(self) -> Optional[Exception]:
         return self._exception
 
+    @property
+    def error_type(self) -> Optional[SupersetErrorType]:
+        return self._error_type
+
 
 class SupersetErrorException(SupersetException):
     """Exceptions with a single SupersetErrorType associated with them"""

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