def wraps(self, *args, **kwargs):
try:
return f(self, *args, **kwargs)
except SupersetSecurityException as e:
logging.exception(e)
return json_error_response(
utils.error_msg_from_exception(e),
status=e.status,
stacktrace=utils.get_stacktrace(),
link=e.link,
)
except SupersetException as e:
logging.exception(e)
return json_error_response(
utils.error_msg_from_exception(e),
stacktrace=utils.get_stacktrace(),
status=e.status,
)
except HTTPException as e:
logging.exception(e)
return json_error_response(
utils.error_msg_from_exception(e),
stacktrace=traceback.format_exc(),
status=e.code,
)
except Exception as e: # pylint: disable=broad-except
logging.exception(e)
return json_error_response(
utils.error_msg_from_exception(e), stacktrace=utils.get_stacktrace()
)
def test_get_stacktrace(self):
with app.app_context():
app.config["SHOW_STACKTRACE"] = True
try:
raise Exception("NONONO!")
except Exception:
stacktrace = get_stacktrace()
self.assertIn("NONONO", stacktrace)
app.config["SHOW_STACKTRACE"] = False
try:
raise Exception("NONONO!")
except Exception:
stacktrace = get_stacktrace()
assert stacktrace is None
def json_error_response(msg=None, status=500, stacktrace=None, payload=None, link=None):
if not payload:
payload = {"error": "{}".format(msg)}
payload["stacktrace"] = utils.get_stacktrace()
if link:
payload["link"] = link
return Response(
json.dumps(payload, default=utils.json_iso_dttm_ser, ignore_nan=True),
status=status,
mimetype="application/json",
)
def wraps(self, *args, **kwargs):
try:
return f(self, *args, **kwargs)
except SupersetSecurityException as e:
logging.exception(e)
return json_error_response(
utils.error_msg_from_exception(e),
status=e.status,
stacktrace=utils.get_stacktrace(),
link=e.link,
)
except SupersetException as e:
logging.exception(e)
return json_error_response(
utils.error_msg_from_exception(e),
stacktrace=utils.get_stacktrace(),
status=e.status,
)
except Exception as e:
logging.exception(e)
return json_error_response(utils.error_msg_from_exception(e),
stacktrace=utils.get_stacktrace())
def set_query_result(
self,
key: str,
query_result: QueryResult,
annotation_data: Optional[Dict[str, Any]] = None,
force_query: Optional[bool] = False,
timeout: Optional[int] = None,
datasource_uid: Optional[str] = None,
region: CacheRegion = CacheRegion.DEFAULT,
) -> None:
"""
Set dataframe of query-result to specific cache region
"""
try:
self.status = query_result.status
self.query = query_result.query
self.applied_template_filters = query_result.applied_template_filters
self.error_message = query_result.error_message
self.df = query_result.df
self.annotation_data = {} if annotation_data is None else annotation_data
if self.status != QueryStatus.FAILED:
stats_logger.incr("loaded_from_source")
if not force_query:
stats_logger.incr("loaded_from_source_without_force")
self.is_loaded = True
value = {
"df": self.df,
"query": self.query,
"applied_template_filters": self.applied_template_filters,
"annotation_data": self.annotation_data,
}
if self.is_loaded and key and self.status != QueryStatus.FAILED:
self.set(
key=key,
value=value,
timeout=timeout,
datasource_uid=datasource_uid,
region=region,
)
except Exception as ex: # pylint: disable=broad-except
logger.exception(ex)
if not self.error_message:
self.error_message = str(ex)
self.status = QueryStatus.FAILED
self.stacktrace = get_stacktrace()
def get_df_payload( # pylint: disable=too-many-statements
self, query_obj: QueryObject, **kwargs: Any) -> Dict[str, Any]:
"""Handles caching around the df payload retrieval"""
cache_key = self.cache_key(query_obj, **kwargs)
logger.info("Cache key: %s", cache_key)
is_loaded = False
stacktrace = None
df = pd.DataFrame()
cached_dttm = datetime.utcnow().isoformat().split(".")[0]
cache_value = None
status = None
query = ""
error_message = None
if cache_key and cache and not self.force:
cache_value = cache.get(cache_key)
if cache_value:
stats_logger.incr("loading_from_cache")
try:
df = cache_value["df"]
query = cache_value["query"]
status = utils.QueryStatus.SUCCESS
is_loaded = True
stats_logger.incr("loaded_from_cache")
except KeyError as ex:
logger.exception(ex)
logger.error("Error reading cache: %s",
utils.error_msg_from_exception(ex))
logger.info("Serving from cache")
if query_obj and not is_loaded:
try:
invalid_columns = [
col for col in query_obj.columns + query_obj.groupby +
utils.get_column_names_from_metrics(query_obj.metrics)
if col not in self.datasource.column_names
]
if invalid_columns:
raise QueryObjectValidationError(
_(
"Columns missing in datasource: %(invalid_columns)s",
invalid_columns=invalid_columns,
))
query_result = self.get_query_result(query_obj)
status = query_result["status"]
query = query_result["query"]
error_message = query_result["error_message"]
df = query_result["df"]
if status != utils.QueryStatus.FAILED:
stats_logger.incr("loaded_from_source")
if not self.force:
stats_logger.incr("loaded_from_source_without_force")
is_loaded = True
except QueryObjectValidationError as ex:
error_message = str(ex)
status = utils.QueryStatus.FAILED
except Exception as ex: # pylint: disable=broad-except
logger.exception(ex)
if not error_message:
error_message = str(ex)
status = utils.QueryStatus.FAILED
stacktrace = utils.get_stacktrace()
if is_loaded and cache_key and cache and status != utils.QueryStatus.FAILED:
set_and_log_cache(
cache_key,
df,
query,
cached_dttm,
self.cache_timeout,
self.datasource.uid,
)
return {
"cache_key": cache_key,
"cached_dttm":
cache_value["dttm"] if cache_value is not None else None,
"cache_timeout": self.cache_timeout,
"df": df,
"error": error_message,
"is_cached": cache_value is not None,
"query": query,
"status": status,
"stacktrace": stacktrace,
"rowcount": len(df.index),
}
def get_df_payload( # pylint: disable=too-many-statements,too-many-locals
self, query_obj: QueryObject, force_cached: Optional[bool] = False,
) -> Dict[str, Any]:
"""Handles caching around the df payload retrieval"""
cache_key = self.query_cache_key(query_obj)
logger.info("Cache key: %s", cache_key)
is_loaded = False
stacktrace = None
df = pd.DataFrame()
cache_value = None
status = None
query = ""
annotation_data = {}
error_message = None
if cache_key and cache_manager.data_cache and not self.force:
cache_value = cache_manager.data_cache.get(cache_key)
if cache_value:
stats_logger.incr("loading_from_cache")
try:
df = cache_value["df"]
query = cache_value["query"]
annotation_data = cache_value.get("annotation_data", {})
status = utils.QueryStatus.SUCCESS
is_loaded = True
stats_logger.incr("loaded_from_cache")
except KeyError as ex:
logger.exception(ex)
logger.error(
"Error reading cache: %s", utils.error_msg_from_exception(ex)
)
logger.info("Serving from cache")
if force_cached and not is_loaded:
logger.warning(
"force_cached (QueryContext): value not found for key %s", cache_key
)
raise CacheLoadError("Error loading data from cache")
if query_obj and not is_loaded:
try:
invalid_columns = [
col
for col in query_obj.columns
+ query_obj.groupby
+ utils.get_column_names_from_metrics(query_obj.metrics)
if col not in self.datasource.column_names and col != DTTM_ALIAS
]
if invalid_columns:
raise QueryObjectValidationError(
_(
"Columns missing in datasource: %(invalid_columns)s",
invalid_columns=invalid_columns,
)
)
query_result = self.get_query_result(query_obj)
status = query_result["status"]
query = query_result["query"]
error_message = query_result["error_message"]
df = query_result["df"]
annotation_data = self.get_annotation_data(query_obj)
if status != utils.QueryStatus.FAILED:
stats_logger.incr("loaded_from_source")
if not self.force:
stats_logger.incr("loaded_from_source_without_force")
is_loaded = True
except QueryObjectValidationError as ex:
error_message = str(ex)
status = utils.QueryStatus.FAILED
except Exception as ex: # pylint: disable=broad-except
logger.exception(ex)
if not error_message:
error_message = str(ex)
status = utils.QueryStatus.FAILED
stacktrace = utils.get_stacktrace()
if is_loaded and cache_key and status != utils.QueryStatus.FAILED:
set_and_log_cache(
cache_manager.data_cache,
cache_key,
{"df": df, "query": query, "annotation_data": annotation_data},
self.cache_timeout,
self.datasource.uid,
)
return {
"cache_key": cache_key,
"cached_dttm": cache_value["dttm"] if cache_value is not None else None,
"cache_timeout": self.cache_timeout,
"df": df,
"annotation_data": annotation_data,
"error": error_message,
"is_cached": cache_value is not None,
"query": query,
"status": status,
"stacktrace": stacktrace,
"rowcount": len(df.index),
}
def get_df_payload( # pylint: disable=too-many-locals,too-many-statements
self, query_obj: QueryObject, **kwargs: Any) -> Dict[str, Any]:
"""Handles caching around the df payload retrieval"""
cache_key = self.cache_key(query_obj, **kwargs)
logger.info("Cache key: %s", cache_key)
is_loaded = False
stacktrace = None
df = pd.DataFrame()
cached_dttm = datetime.utcnow().isoformat().split(".")[0]
cache_value = None
status = None
query = ""
error_message = None
if cache_key and cache and not self.force:
cache_value = cache.get(cache_key)
if cache_value:
stats_logger.incr("loading_from_cache")
try:
cache_value = pkl.loads(cache_value)
df = cache_value["df"]
query = cache_value["query"]
status = utils.QueryStatus.SUCCESS
is_loaded = True
stats_logger.incr("loaded_from_cache")
except Exception as ex: # pylint: disable=broad-except
logger.exception(ex)
logger.error("Error reading cache: %s",
utils.error_msg_from_exception(ex))
logger.info("Serving from cache")
if query_obj and not is_loaded:
try:
query_result = self.get_query_result(query_obj)
status = query_result["status"]
query = query_result["query"]
error_message = query_result["error_message"]
df = query_result["df"]
if status != utils.QueryStatus.FAILED:
stats_logger.incr("loaded_from_source")
if not self.force:
stats_logger.incr("loaded_from_source_without_force")
is_loaded = True
except Exception as ex: # pylint: disable=broad-except
logger.exception(ex)
if not error_message:
error_message = "{}".format(ex)
status = utils.QueryStatus.FAILED
stacktrace = utils.get_stacktrace()
if is_loaded and cache_key and cache and status != utils.QueryStatus.FAILED:
try:
cache_value = dict(dttm=cached_dttm, df=df, query=query)
cache_binary = pkl.dumps(cache_value,
protocol=pkl.HIGHEST_PROTOCOL)
logger.info("Caching %d chars at key %s",
len(cache_binary), cache_key)
stats_logger.incr("set_cache_key")
cache.set(cache_key,
cache_binary,
timeout=self.cache_timeout)
except Exception as ex: # pylint: disable=broad-except
# cache.set call can fail if the backend is down or if
# the key is too large or whatever other reasons
logger.warning("Could not cache key %s", cache_key)
logger.exception(ex)
cache.delete(cache_key)
return {
"cache_key": cache_key,
"cached_dttm":
cache_value["dttm"] if cache_value is not None else None,
"cache_timeout": self.cache_timeout,
"df": df,
"error": error_message,
"is_cached": cache_key is not None,
"query": query,
"status": status,
"stacktrace": stacktrace,
"rowcount": len(df.index),
}
def get_df_payload(self, query_obj: QueryObject, **kwargs):
"""Handles caching around the df paylod retrieval"""
extra_cache_keys = self.datasource.get_extra_cache_keys(query_obj.to_dict())
cache_key = (
query_obj.cache_key(
datasource=self.datasource.uid,
extra_cache_keys=extra_cache_keys,
**kwargs
)
if query_obj
else None
)
logging.info("Cache key: {}".format(cache_key))
is_loaded = False
stacktrace = None
df = None
cached_dttm = datetime.utcnow().isoformat().split(".")[0]
cache_value = None
status = None
query = ""
error_message = None
if cache_key and cache and not self.force:
cache_value = cache.get(cache_key)
if cache_value:
stats_logger.incr("loaded_from_cache")
try:
cache_value = pkl.loads(cache_value)
df = cache_value["df"]
query = cache_value["query"]
status = utils.QueryStatus.SUCCESS
is_loaded = True
except Exception as e:
logging.exception(e)
logging.error(
"Error reading cache: " + utils.error_msg_from_exception(e)
)
logging.info("Serving from cache")
if query_obj and not is_loaded:
try:
query_result = self.get_query_result(query_obj)
status = query_result["status"]
query = query_result["query"]
error_message = query_result["error_message"]
df = query_result["df"]
if status != utils.QueryStatus.FAILED:
stats_logger.incr("loaded_from_source")
is_loaded = True
except Exception as e:
logging.exception(e)
if not error_message:
error_message = "{}".format(e)
status = utils.QueryStatus.FAILED
stacktrace = utils.get_stacktrace()
if is_loaded and cache_key and cache and status != utils.QueryStatus.FAILED:
try:
cache_value = dict(
dttm=cached_dttm, df=df if df is not None else None, query=query
)
cache_binary = pkl.dumps(cache_value, protocol=pkl.HIGHEST_PROTOCOL)
logging.info(
"Caching {} chars at key {}".format(
len(cache_binary), cache_key
)
)
stats_logger.incr("set_cache_key")
cache.set(cache_key, cache_binary, timeout=self.cache_timeout)
except Exception as e:
# cache.set call can fail if the backend is down or if
# the key is too large or whatever other reasons
logging.warning("Could not cache key {}".format(cache_key))
logging.exception(e)
cache.delete(cache_key)
return {
"cache_key": cache_key,
"cached_dttm": cache_value["dttm"] if cache_value is not None else None,
"cache_timeout": self.cache_timeout,
"df": df,
"error": error_message,
"is_cached": cache_key is not None,
"query": query,
"status": status,
"stacktrace": stacktrace,
"rowcount": len(df.index) if df is not None else 0,
}