[fix] SQL query source

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 @@ def get_sqla_engine(
         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 @@ def get_sqla_engine(
         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 @@ def get_df(  # pylint: disable=too-many-locals
         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 @@ def select_star(  # pylint: disable=too-many-arguments
         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.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 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 @@ def validate(
             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 @@
 
 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

@@ -1981,7 +1981,7 @@ def estimate_query_cost(
         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)