GitHub user Dhruvin4530 added a comment to the discussion: Superset is not able
to load the result from hydrolix
@dosu
This is my query mutator script.
```
"""Superset configurations, including extra Jinja macros and SQL_QUERY_MUTATOR
Updated for Superset v5.0 compatibility:
- SQL_QUERY_MUTATOR signature changed to (sql: str, **kwargs) -> str
- DatasetDAO moved from superset.datasets.dao to superset.daos.dataset
- get_username() removed; use Flask g.user instead
- merge_extra_filters / convert_legacy_filters_into_adhoc removed (legacy
filters dropped in v4)
"""
import json
import re
import shlex
from datetime import datetime, timezone, timedelta
from typing import Any, Dict, List, Optional, Tuple
import logging
log = logging.getLogger(__name__)
def _get_username() -> Optional[str]:
"""Get the current username from Flask's g object.
Replaces the removed superset.utils.core.get_username().
"""
try:
from flask import g
if hasattr(g, "user") and g.user and hasattr(g.user, "username"):
return g.user.username
except Exception:
pass
return None
############################## Jinja Macros ##############################
def events_dataset_id() -> int:
"""Returns the ID of the events dataset in Superset.
Potentially needs to have a different value between Superset instances.
"""
return 2
def dataset_columns(dataset_id: Optional[int] = None) -> List:
"""List of columns available for a dataset, according to Superset.
Example usage:
```
SELECT
name
{% for column in dataset_columns(10) %}
, "column"
{% endfor %}
```
Args:
dataset_id: ID of the dataset in Superset.
"""
# pylint: disable=import-outside-toplevel
# v5: DatasetDAO moved from superset.datasets.dao to superset.daos.dataset
from superset.daos.dataset import DatasetDAO
if dataset_id is None:
dataset_id = datasource_id()
if dataset_id is None:
return []
dataset = DatasetDAO.find_by_id(dataset_id)
if not dataset:
from superset.datasets.commands.exceptions import DatasetNotFoundError
raise DatasetNotFoundError(f"Dataset {dataset_id} not found!")
return dataset.columns
def events_column_names() -> List[str]:
"""List of (non-calculated) colummns available in the "events" table,
according to Superset.
Note that for this to work properly, events_dataset_id() needs to return
the correct dataset ID for the events table
in this particular Superset instance.
"""
return [
col.column_name
for col in dataset_columns(events_dataset_id())
if not col.expression
]
def dimensions(dataset_id: Optional[int] = None) -> List[str]:
"""List of (non-calculated) dimensions (columns available for GROUP BY)
available for a dataset, according to Superset.
Args:
dataset_id: ID of the dataset in Superset.
"""
return [
col.column_name
for col in dataset_columns(dataset_id)
if not col.expression and col.groupby
]
def slice_id() -> Optional[int]:
"""ID of the "slice" (a.k.a. a "Chart") in Superset associated with the
current query, if any."""
# pylint: disable=import-outside-toplevel
from superset.views.utils import get_form_data
try:
return get_form_data()[0].get("slice_id")
except Exception:
return None
def datasource_id() -> Optional[int]:
"""Returns the ID of the "datasource" (source dataset) in Superset
associated with the current query, if any."""
# pylint: disable=import-outside-toplevel
from superset.views.utils import get_form_data
try:
form_data, slc = get_form_data(use_slice_data=True)
except Exception:
return None
dataset_id = form_data.get("url_params", {}).get("dataset_id")
if dataset_id is not None:
return dataset_id
if slc is not None:
return slc.datasource_id
return None
def dashboard_id() -> Optional[int]:
"""Returns the ID of the dashboard in Superset associated with the current
query, if any."""
import flask
try:
return flask.request.values.get("dashboard_id")
except Exception:
return None
def full_form_data() -> Dict:
"""Returns the full form data for the current query.
This includes various settings passed on from the Superset UI, e.g. filter
configurations,
which are used internally in Superset as part of automatic query generation.
Note: In Superset v4+, legacy filters were removed. The calls to
convert_legacy_filters_into_adhoc() and merge_extra_filters() have been
removed as they no longer exist.
"""
# pylint: disable=import-outside-toplevel
from superset.views.utils import get_form_data
form_data, _ = get_form_data()
# v5: These two functions were removed along with legacy filter support:
# convert_legacy_filters_into_adhoc(form_data)
# merge_extra_filters(form_data)
return form_data
def time_filter(
field_name: str = "timestamp",
default_minutes_since: Optional[int] = 60,
default_minutes_until: Optional[int] = 0,
) -> str:
"""Returns an SQL-formatted time filter condition for the selected time
range.
This makes it more convenient to write more efficient inner queries in
virtual datasets.
Example usage:
```
WHERE {# time_filter("s_timestamp", 10) #}
AND x = 7
```
Args:
field_name: field to filter on
default_minutes_since: default minutes (back from now) to cover, if no
time range is specified via the Superset UI.
Set this to None to allow omitting a timerange start.
default_minutes_until: default minutes (forward from now) to cover, if
no time range is specified via the Superset UI.
Set this to None to allow omitting a timerange start.
"""
since, until = since_until()
time_format = "%Y-%m-%d %H:%M:%S"
now = datetime.now(tz=timezone.utc)
if since is None and default_minutes_since is not None:
since = now - timedelta(minutes=default_minutes_since)
if until is None and default_minutes_until is not None:
until = now + timedelta(minutes=default_minutes_until)
conditions = []
if since is not None:
conditions.append(f"{field_name} >= '{since.strftime(time_format)}'")
if until is not None:
conditions.append(f"{field_name} < '{until.strftime(time_format)}'")
if not conditions:
conditions.append("1")
return combine_sql_conditions(conditions)
def where_filters(
restricted_columns: Optional[List[str]] = None, include_extras: bool = True
) -> str:
"""Returns an SQL-formatted combination of WHERE conditions for the current
query.
This makes it more convenient to write more efficient inner queries in
virtual datasets.
Example usage:
```
WHERE {# time_filter("timestamp", 30) #}
AND {# where_filters(events_column_names(), False) #}
```
Args:
restricted_columns: list of fields to allow (only for "Simple"
filters); when present, filters that reference unlisted fields
will be excluded. This is useful for excluding e.g. filters on
computed fields that are only present in the outer query.
include_extras: whether to include any extra filters specified via
"Custom SQL" in Superset - these can't be filtered using
restricted_columns because Superset doesn't parse field names out
of them.
"""
form_data = full_form_data()
filter_conditions = []
if include_extras:
base_condition = form_data.get("extras", {}).get("where")
if base_condition:
filter_conditions.append(base_condition)
for f in form_data.get("adhoc_filters", []):
if f["clause"] != "WHERE":
continue
if restricted_columns is not None and f["subject"] not in
restricted_columns:
continue
filter_conditions.append(adhoc_filter_to_sql(f))
if not filter_conditions:
filter_conditions = ["1"]
return combine_sql_conditions(filter_conditions, joiner="\nAND ")
JINJA_CONTEXT_ADDONS = {
"events_dataset_id": events_dataset_id,
"dataset_columns": dataset_columns,
"events_column_names": events_column_names,
"dimensions": dimensions,
"slice_id": slice_id,
"datasource_id": datasource_id,
"dashboard_id": dashboard_id,
"full_form_data": full_form_data,
"time_filter": time_filter,
"where_filters": where_filters,
}
"""See the official documentation here:
https://superset.apache.org/docs/installation/sql-templating/#jinja-templates
"""
############################ Jinja Macro Helpers ############################
def adhoc_filter_to_sql(filter: Dict) -> str:
col = filter.get("subject")
op = filter.get("operator")
val = filter.get("comparator")
if val is None:
return f"{col} {op}"
else:
return f"{col} {op} {val!r}"
def combine_sql_conditions(conditions: List[str], joiner: str = " AND ") -> str:
return "(" + joiner.join(f"({c})" for c in conditions) + ")"
def since_until() -> Tuple[Optional[datetime], Optional[datetime]]:
# pylint: disable=import-outside-toplevel
from superset.utils.date_parser import get_since_until
form_data = full_form_data()
time_range = form_data.get("time_range")
if time_range is not None:
return get_since_until(time_range)
else:
return None, None
########################### Query Mutator Helpers ###########################
def load_setting(token: str) -> Any:
if token[0] == "'":
token = '"' + token[1:-1] + '"'
return json.loads(token)
def dump_setting(value: Any) -> str:
token = json.dumps(value)
if token[0] == '"':
token = "'" + token[1:-1] + "'"
return token
def split_settings(settings_str: str) -> Dict[str, Any]:
"""Parses the provided settings string into a dictionary of settings
key:value pairs"""
result = {}
tokens = list(shlex.shlex(settings_str))
it = iter(tokens)
try:
while True:
key = next(it)
if next(it) != "=":
break
result[key] = load_setting(next(it))
if next(it) != ",":
break
except StopIteration:
pass
return result
def join_settings(settings_dict: Dict[str, Any]) -> str:
"""Joins settings using function `dump_settings()`"""
return ", ".join(
f"{key}={dump_setting(value)}" for key, value in settings_dict.items()
)
def extract_settings(query: str) -> Tuple[str, Dict[str, Any]]:
"""Extracts currently applied query settings from the query string.
Args:
query (str): The query string to extract from.
Returns:
Tuple[str, Dict[str, Any]]: A tuple containing the query and a
dictionary of
the extracted query `SETTINGS`.
"""
settings_regex = r"SETTINGS\s+(hdx_[^\)\n]+)"
flags = re.IGNORECASE | re.MULTILINE
results = re.findall(settings_regex, query, flags=flags)
current_settings = {}
if results:
query = re.sub(settings_regex, "", query, flags=flags)
for result in results:
current_settings.update(split_settings(result))
return query, current_settings
def hdx_with_settings(query: str, settings: Dict[str, Any]) -> str:
"""Applies provided Hdyrolix query SETTINGS to the query string
https://docs.hydrolix.io/docs/query-api-http-options#using-query-options
"""
settings_str = join_settings(settings)
return f"{query}\nSETTINGS {settings_str}"
def athena_with_settings(query: str, settings: Dict[str, Any]) -> str:
"""This is just to give some tagging to the query in the Athena log.
This functionality to be improved in JIRA ticket
[DATA-2357](https://arkoselabs.atlassian.net/browse/DATA-2357)
"""
settings_str = join_settings(settings)
return f"/*\nSuperset {settings_str}\n*/\n{query}"
def replace_series_limit(sql: str) -> str:
if "_no_replace_series_limit" in sql:
return sql
inner_query_regex = (
r"INNER JOIN.*?\(SELECT.*?LIMIT ([0-9]+)\) AS [`]?\w+[`]?.*?(?=WHERE)"
)
flags = re.IGNORECASE | re.DOTALL
match = re.search(inner_query_regex, sql, flags=flags)
if not match:
return sql
series_limit = int(match.group(1))
sql = re.sub(inner_query_regex, "", sql, flags=flags)
# Extract the first column from GROUP BY to use in LIMIT BY
# v4 used __timestamp, v5 uses hashed aliases like timestamp_d7e6d5
limit_by_col = "__timestamp"
group_by_match = re.search(r"GROUP BY\s+`?(\w+)`?", sql,
flags=re.IGNORECASE)
if group_by_match:
limit_by_col = group_by_match.group(1)
limit_regex = "LIMIT [0-9]+$"
sql = re.sub(limit_regex, "", sql, flags=flags)
sql = f"{sql}\nLIMIT {series_limit} BY {limit_by_col}"
return sql
def rewrite_array_join_filter_conditions(sql: str) -> str:
if "_no_array_join_rewrite" in sql:
return sql
flags = re.MULTILINE
array_join_equals_regex = (
r"(\n[ \t]*| +)AND arrayJoin\(([^\n]+?)\) = ([^\n]*?)(?=[ \t]*(?:\n|$))"
)
array_join_equals_replacement = r"\1AND has(\2, \3)"
sql = re.sub(
array_join_equals_regex, array_join_equals_replacement, sql, flags=flags
)
# Special case for single element + IN clause
array_join_one_element_in_regex = (
r"(\n[ \t]*| +)AND arrayJoin\(([^\n]+?)\) IN \(([^,\n]*?)\)"
)
array_join_one_element_in_replacement = r"\1AND has(\2, \3)"
sql = re.sub(
array_join_one_element_in_regex,
array_join_one_element_in_replacement,
sql,
flags=flags,
)
# Alias with _AND -> use hasAll(...)
array_join_in_AND_regex = (
r"(\n[ \t]*| +)AND \(arrayJoin\(([^\n]+?)\) AS __\w*_AND\) IN
\(([^)]*?)\)"
)
array_join_in_AND_replacement = r"\1AND hasAll(\2, [\3])"
sql = re.sub(
array_join_in_AND_regex, array_join_in_AND_replacement, sql, flags=flags
)
# Alias with _OR -> use hasAny(...)
array_join_in_OR_regex = (
r"(\n[ \t]*| +)AND \(arrayJoin\(([^\n]+?)\) AS __\w*_OR\) IN
\(([^)]*?)\)"
)
array_join_in_OR_replacement = r"\1AND hasAny(\2, [\3])"
sql = re.sub(array_join_in_OR_regex, array_join_in_OR_replacement, sql,
flags=flags)
# Any other alias => won't be replaced (Like _no_array_join_rewrite, but
for one field instead of the whole query)
# Default (no alias) -> use hasAny(...)
array_join_in_regex = r"(\n[ \t]*| +)AND arrayJoin\(([^\n]+?)\) IN
\(([^)]*?)\)"
array_join_in_replacement = r"\1AND hasAny(\2, [\3])"
sql = re.sub(array_join_in_regex, array_join_in_replacement, sql,
flags=flags)
return sql
def get_dashboard_certification(dashboard_id: int) -> bool:
# from superset import db
from superset.models.dashboard import Dashboard
CERTIFIED_DASHBOARD_NAMES = [""]
results = Dashboard.query.filter_by(id=dashboard_id).first()
# session = db.session
# result = session.execute(select(Dashboard).order_by(User.id))
if results.dashboard_title in CERTIFIED_DASHBOARD_NAMES:
return True
return results is not None
def get_settings(
user_name: str,
security_manager: Any,
query_settings: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
SECONDS_PER_DAY = 60 * 60 * 24
BINARY_GIGA = 2**30
# Mutable settings can be changed on a per-query basis by using a SETTINGS
statement in the query.
# query_settings can then override any mutable default settings for that
user.
mutable_settings = set(
[
"hdx_query_timerange_required",
"hdx_query_max_timerange_sec",
"hdx_query_max_execution_time",
"hdx_query_pool_name",
]
)
# Default settings
settings = {
"hdx_query_pool_name": "query-peer",
"hdx_query_max_memory_usage": 38 * BINARY_GIGA,
"hdx_query_timerange_required": True,
"hdx_query_max_timerange_sec": 62 * SECONDS_PER_DAY,
"hdx_query_max_execution_time": 120,
}
# Role overrides are processed in order, and will override values in the
default settings.
role_overrides = [
(
"SOC",
{
"hdx_query_pool_name": "soc",
"hdx_query_max_timerange_sec": 93 * SECONDS_PER_DAY,
},
),
(
"CSOPS",
{
"hdx_query_pool_name": "soc",
"hdx_query_max_timerange_sec": 93 * SECONDS_PER_DAY,
},
),
(
"SRSOC",
{
"hdx_query_pool_name": "soc",
"hdx_query_max_timerange_sec": 93 * SECONDS_PER_DAY,
},
),
(
"PORTAL",
{
"hdx_query_pool_name": "portal",
},
),
]
# Apply role overrides to the default settings, based on the user's roles.
user = security_manager.find_user(username=user_name)
log.debug(f"user: {user}")
user_roles = security_manager.get_user_roles(user)
user_role_names = set(user_role.name for user_role in user_roles)
for role_name, override_dict in role_overrides:
if role_name in user_role_names:
settings.update(override_dict)
# Update any mutable / absent settings with the query settings, if present.
if query_settings is not None:
for key, value in query_settings.items():
if key in mutable_settings or key not in settings:
settings[key] = value
# Add an admin comment to the SETTINGS
admin_comment = f"User: {user}"
s_id = slice_id()
if s_id is not None:
admin_comment += f"; slice_id: {s_id}"
ds_id = datasource_id()
if ds_id is not None:
admin_comment += f"; datasource_id: {ds_id}"
d_id = dashboard_id()
if d_id is not None:
admin_comment += f"; dashboard_id: {d_id}"
settings["hdx_query_admin_comment"] = admin_comment
return settings
############################### Query Mutator ###############################
def SQL_QUERY_MUTATOR(sql, **kwargs):
"""Query mutator compatible with Superset v5.0+.
In Superset v2, the signature was:
SQL_QUERY_MUTATOR(sql, security_manager, database, **kwargs)
Since Superset v2.0 (PR #19083), the signature changed to:
SQL_QUERY_MUTATOR(sql: str, **kwargs: Any) -> str
where security_manager, database, user_name, etc. are passed via kwargs.
"""
# Extract arguments from kwargs (v5 passes everything via kwargs)
security_manager = kwargs.get("security_manager")
database = kwargs.get("database")
# Sometimes we get a different object here.
# But it should at least be convertible to a string that represents the
actual query.
log.debug(f"SQL_QUERY_MUTATOR")
log.debug(f"kwargs: {kwargs}")
query = str(sql)
# Don't mutate if this isn't a SELECT query
select_search = re.search(r"^\s*(SELECT|WITH)", query, flags=re.IGNORECASE)
if not select_search:
return query
# Extract any SETTINGS already set in the query
query, query_settings = extract_settings(query)
# Get username - v5: get_username() was removed, use Flask g.user instead
user_name = _get_username()
# Determine the SETTINGS for this query
new_settings = get_settings(
user_name=user_name,
security_manager=security_manager,
query_settings=query_settings,
)
database_name = database.database_name
log.debug(f"database_name: {database_name}")
# Add the SETTINGS to the query and return it.
if "hydrolix" in database_name.lower():
# Fix series limits to use LIMIT BY instead of subquery.
query = replace_series_limit(query)
query = rewrite_array_join_filter_conditions(query)
query = hdx_with_settings(query, new_settings)
if "athena" in database_name.lower():
query = athena_with_settings(query, new_settings)
return query
```
GitHub link:
https://github.com/apache/superset/discussions/41395#discussioncomment-17424927
----
This is an automatically sent email for [email protected].
To unsubscribe, please send an email to:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
