def handle_query_error(
ex: Exception,
query: Query,
session: Session,
payload: Optional[Dict[str, Any]] = None,
prefix_message: str = "",
) -> Dict[str, Any]:
"""Local method handling error while processing the SQL"""
payload = payload or {}
msg = f"{prefix_message} {str(ex)}".strip()
troubleshooting_link = config["TROUBLESHOOTING_LINK"]
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
# extract DB-specific errors (invalid column, eg)
if isinstance(ex, SupersetErrorException):
errors = [ex.error]
elif isinstance(ex, SupersetErrorsException):
errors = ex.errors
else:
errors = query.database.db_engine_spec.extract_errors(str(ex))
errors_payload = [dataclasses.asdict(error) for error in errors]
if errors:
query.set_extra_json_key("errors", errors_payload)
session.commit()
payload.update({"status": query.status, "error": msg, "errors": errors_payload})
if troubleshooting_link:
payload["link"] = troubleshooting_link
return payload
def handle_query_error(
msg: str,
query: Query,
session: Session,
payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
"""Local method handling error while processing the SQL"""
payload = payload or {}
troubleshooting_link = config["TROUBLESHOOTING_LINK"]
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
# extract DB-specific errors (invalid column, eg)
errors = [
dataclasses.asdict(error)
for error in query.database.db_engine_spec.extract_errors(msg)
]
if errors:
query.set_extra_json_key("errors", errors)
session.commit()
payload.update({"status": query.status, "error": msg, "errors": errors})
if troubleshooting_link:
payload["link"] = troubleshooting_link
return payload
def handle_query_error(
msg: str, query: Query, session: Session, payload: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""Local method handling error while processing the SQL"""
payload = payload or {}
troubleshooting_link = config["TROUBLESHOOTING_LINK"]
query.error_message = msg
query.status = QueryStatus.FAILED
query.tmp_table_name = None
session.commit()
payload.update({"status": query.status, "error": msg})
if troubleshooting_link:
payload["link"] = troubleshooting_link
return payload
def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-statements
sql_statement: str,
query: Query,
session: Session,
cursor: Any,
log_params: Optional[Dict[str, Any]],
apply_ctas: bool = False,
) -> SupersetResultSet:
"""Executes a single SQL statement"""
database: Database = query.database
db_engine_spec = database.db_engine_spec
parsed_query = ParsedQuery(sql_statement)
if is_feature_enabled("RLS_IN_SQLLAB"):
# Insert any applicable RLS predicates
parsed_query = ParsedQuery(
str(
insert_rls(
parsed_query._parsed[0], # pylint: disable=protected-access
database.id,
query.schema,
)))
sql = parsed_query.stripped()
# This is a test to see if the query is being
# limited by either the dropdown or the sql.
# We are testing to see if more rows exist than the limit.
increased_limit = None if query.limit is None else query.limit + 1
if not db_engine_spec.is_readonly_query(
parsed_query) and not database.allow_dml:
raise SupersetErrorException(
SupersetError(
message=__(
"Only SELECT statements are allowed against this database."
),
error_type=SupersetErrorType.DML_NOT_ALLOWED_ERROR,
level=ErrorLevel.ERROR,
))
if apply_ctas:
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = "tmp_{}_table_{}".format(
query.user_id, start_dttm.strftime("%Y_%m_%d_%H_%M_%S"))
sql = parsed_query.as_create_table(
query.tmp_table_name,
schema_name=query.tmp_schema_name,
method=query.ctas_method,
)
query.select_as_cta_used = True
# Do not apply limit to the CTA queries when SQLLAB_CTAS_NO_LIMIT is set to true
if db_engine_spec.is_select_query(parsed_query) and not (
query.select_as_cta_used and SQLLAB_CTAS_NO_LIMIT):
if SQL_MAX_ROW and (not query.limit or query.limit > SQL_MAX_ROW):
query.limit = SQL_MAX_ROW
sql = apply_limit_if_exists(database, increased_limit, query, sql)
# Hook to allow environment-specific mutation (usually comments) to the SQL
sql = SQL_QUERY_MUTATOR(
sql,
user_name=get_username(), # TODO(john-bodley): Deprecate in 3.0.
security_manager=security_manager,
database=database,
)
try:
query.executed_sql = sql
if log_query:
log_query(
query.database.sqlalchemy_uri,
query.executed_sql,
query.schema,
get_username(),
__name__,
security_manager,
log_params,
)
session.commit()
with stats_timing("sqllab.query.time_executing_query", stats_logger):
logger.debug("Query %d: Running query: %s", query.id, sql)
db_engine_spec.execute(cursor, sql, async_=True)
logger.debug("Query %d: Handling cursor", query.id)
db_engine_spec.handle_cursor(cursor, query, session)
with stats_timing("sqllab.query.time_fetching_results", stats_logger):
logger.debug(
"Query %d: Fetching data for query object: %s",
query.id,
str(query.to_dict()),
)
data = db_engine_spec.fetch_data(cursor, increased_limit)
if query.limit is None or len(data) <= query.limit:
query.limiting_factor = LimitingFactor.NOT_LIMITED
else:
# return 1 row less than increased_query
data = data[:-1]
except SoftTimeLimitExceeded as ex:
query.status = QueryStatus.TIMED_OUT
logger.warning("Query %d: Time limit exceeded", query.id)
logger.debug("Query %d: %s", query.id, ex)
raise SupersetErrorException(
SupersetError(
message=__(
"The query was killed after %(sqllab_timeout)s seconds. It might "
"be too complex, or the database might be under heavy load.",
sqllab_timeout=SQLLAB_TIMEOUT,
),
error_type=SupersetErrorType.SQLLAB_TIMEOUT_ERROR,
level=ErrorLevel.ERROR,
)) from ex
except Exception as ex:
# query is stopped in another thread/worker
# stopping raises expected exceptions which we should skip
session.refresh(query)
if query.status == QueryStatus.STOPPED:
raise SqlLabQueryStoppedException() from ex
logger.error("Query %d: %s", query.id, type(ex), exc_info=True)
logger.debug("Query %d: %s", query.id, ex)
raise SqlLabException(db_engine_spec.extract_error_message(ex)) from ex
logger.debug("Query %d: Fetching cursor description", query.id)
cursor_description = cursor.description
return SupersetResultSet(data, cursor_description, db_engine_spec)
