feat: improve event logging for queries + refactor (#27943)

Co-authored-by: Beto Dealmeida <roberto@dealmeida.net>
This commit is contained in:
Maxime Beauchemin
2024-04-22 16:30:50 -07:00
committed by GitHub
parent 83fedcc9ea
commit cfc440c56c
14 changed files with 169 additions and 132 deletions

View File

@@ -46,7 +46,7 @@ from superset.exceptions import (
SupersetErrorException,
SupersetErrorsException,
)
from superset.extensions import celery_app
from superset.extensions import celery_app, event_logger
from superset.models.core import Database
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
@@ -73,7 +73,6 @@ SQLLAB_TIMEOUT = config["SQLLAB_ASYNC_TIME_LIMIT_SEC"]
SQLLAB_HARD_TIMEOUT = SQLLAB_TIMEOUT + 60
SQL_MAX_ROW = config["SQL_MAX_ROW"]
SQLLAB_CTAS_NO_LIMIT = config["SQLLAB_CTAS_NO_LIMIT"]
SQL_QUERY_MUTATOR = config["SQL_QUERY_MUTATOR"]
log_query = config["QUERY_LOGGER"]
logger = logging.getLogger(__name__)
@@ -264,11 +263,7 @@ def execute_sql_statement( # pylint: disable=too-many-statements
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,
security_manager=security_manager,
database=database,
)
sql = database.mutate_sql_based_on_config(sql)
try:
query.executed_sql = sql
if log_query:
@@ -281,21 +276,26 @@ def execute_sql_statement( # pylint: disable=too-many-statements
log_params,
)
db.session.commit()
with stats_timing("sqllab.query.time_executing_query", stats_logger):
db_engine_spec.execute_with_cursor(cursor, sql, query)
with event_logger.log_context(
action="execute_sql",
database=database,
object_ref=__name__,
):
with stats_timing("sqllab.query.time_executing_query", stats_logger):
db_engine_spec.execute_with_cursor(cursor, sql, query)
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]
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