[fix] SQL query source (#9173)

UPDATING.md

@@ -22,17 +22,18 @@ This file documents any backwards-incompatible changes in Superset and
 assists people when migrating to a new version.
 
 ## Next
+* [9173](https://github.com/apache/incubator-superset/pull/9173): Changes the encoding of the query source from an int to an enum.
 * [9120](https://github.com/apache/incubator-superset/pull/9120): Changes the default behavior of ad-hoc sharing of
 queries in SQLLab to one that links to the saved query rather than one that copies the query data into the KVStore
 model and links to the record there. This is a security-related change that makes SQLLab query
 sharing respect the existing role-based access controls. Should you wish to retain the existing behavior, set two feature flags:
 `"KV_STORE": True` will re-enable the `/kv/` and `/kv/store/` endpoints, and `"SHARE_QUERIES_VIA_KV_STORE": True`
 will tell the front-end to utilize them for query sharing.
-* [9109](https://github.com/apache/incubator-superset/pull/9109): Expire `filter_immune_slices` and 
+* [9109](https://github.com/apache/incubator-superset/pull/9109): Expire `filter_immune_slices` and
 `filter_immune_filter_fields` to favor dashboard scoped filter metadata `filter_scopes`.
 
 * [9046](https://github.com/apache/incubator-superset/pull/9046): Replaces `can_only_access_owned_queries` by
-`all_query_access` favoring a white list approach. Since a new permission is introduced use `superset init` 
+`all_query_access` favoring a white list approach. Since a new permission is introduced use `superset init`
 to create and associate it by default to the `Admin` role. Note that, by default, all non `Admin` users will
 not be able to access queries they do not own.
 

superset/models/core.py

@@ -278,7 +278,7 @@ class Database(
         schema: Optional[str] = None,
         nullpool: bool = True,
         user_name: Optional[str] = None,
-        source: Optional[int] = None,
+        source: Optional[utils.QuerySource] = None,
     ) -> Engine:
         extra = self.get_extra()
         sqlalchemy_url = make_url(self.sqlalchemy_uri_decrypted)
@@ -314,6 +314,12 @@ class Database(
         params.update(self.get_encrypted_extra())
 
         if DB_CONNECTION_MUTATOR:
+            if not source and request and request.referrer:
+                if "/superset/dashboard/" in request.referrer:
+                    source = utils.QuerySource.DASHBOARD
+                elif "/superset/explore/" in request.referrer:
+                    source = utils.QuerySource.CHART
+
             sqlalchemy_url, params = DB_CONNECTION_MUTATOR(
                 sqlalchemy_url, params, effective_username, security_manager, source
             )
@@ -330,15 +336,8 @@ class Database(
         self, sql: str, schema: Optional[str] = None, mutator: Optional[Callable] = None
     ) -> pd.DataFrame:
         sqls = [str(s).strip(" ;") for s in sqlparse.parse(sql)]
-        source_key = None
-        if request and request.referrer:
-            if "/superset/dashboard/" in request.referrer:
-                source_key = "dashboard"
-            elif "/superset/explore/" in request.referrer:
-                source_key = "chart"
-        engine = self.get_sqla_engine(
-            schema=schema, source=utils.sources[source_key] if source_key else None
-        )
+
+        engine = self.get_sqla_engine(schema=schema)
         username = utils.get_username()
 
         def needs_conversion(df_series: pd.Series) -> bool:
@@ -398,9 +397,7 @@ class Database(
         cols: Optional[List[Dict[str, Any]]] = None,
     ):
         """Generates a ``select *`` statement in the proper dialect"""
-        eng = self.get_sqla_engine(
-            schema=schema, source=utils.sources.get("sql_lab", None)
-        )
+        eng = self.get_sqla_engine(schema=schema, source=utils.QuerySource.SQL_LAB)
         return self.db_engine_spec.select_star(
             self,
             table_name,

superset/sql_lab.py

@@ -45,7 +45,12 @@ from superset.extensions import celery_app
 from superset.models.sql_lab import Query
 from superset.result_set import SupersetResultSet
 from superset.sql_parse import ParsedQuery
-from superset.utils.core import json_iso_dttm_ser, QueryStatus, sources, zlib_compress
+from superset.utils.core import (
+    json_iso_dttm_ser,
+    QuerySource,
+    QueryStatus,
+    zlib_compress,
+)
 from superset.utils.dates import now_as_float
 from superset.utils.decorators import stats_timing
 
@@ -341,7 +346,7 @@ def execute_sql_statements(
         schema=query.schema,
         nullpool=True,
         user_name=user_name,
-        source=sources.get("sql_lab", None),
+        source=QuerySource.SQL_LAB,
     )
     # Sharing a single connection and cursor across the
     # execution of all statements (if many)

superset/sql_validators/presto_db.py

@@ -25,7 +25,7 @@ from flask import g
 from superset import app, security_manager
 from superset.sql_parse import ParsedQuery
 from superset.sql_validators.base import BaseSQLValidator, SQLValidationAnnotation
-from superset.utils.core import sources
+from superset.utils.core import QuerySource
 
 MAX_ERROR_ROWS = 10
 
@@ -160,7 +160,7 @@ class PrestoDBSQLValidator(BaseSQLValidator):
             schema=schema,
             nullpool=True,
             user_name=user_name,
-            source=sources.get("sql_lab", None),
+            source=QuerySource.SQL_LAB,
         )
         # Sharing a single connection and cursor across the
         # execution of all statements (if many)

superset/utils/core.py

@@ -72,8 +72,6 @@ ADHOC_METRIC_EXPRESSION_TYPES = {"SIMPLE": "SIMPLE", "SQL": "SQL"}
 
 JS_MAX_INTEGER = 9007199254740991  # Largest int Java Script can handle 2^53-1
 
-sources = {"chart": 0, "dashboard": 1, "sql_lab": 2}
-
 try:
     # Having might not have been imported.
     class DimSelector(Having):
@@ -1235,3 +1233,13 @@ class ReservedUrlParameters(Enum):
 
     STANDALONE = "standalone"
     EDIT_MODE = "edit"
+
+
+class QuerySource(Enum):
+    """
+    The source of a SQL query.
+    """
+
+    CHART = 0
+    DASHBOARD = 1
+    SQL_LAB = 2

superset/views/core.py

@@ -1978,7 +1978,7 @@ class Superset(BaseSupersetView):
         try:
             with utils.timeout(seconds=timeout, error_message=timeout_msg):
                 cost = mydb.db_engine_spec.estimate_query_cost(
-                    mydb, schema, sql, utils.sources.get("sql_lab")
+                    mydb, schema, sql, utils.QuerySource.SQL_LAB
                 )
         except SupersetTimeoutException as e:
             logger.exception(e)