diff --git a/snuba/admin/views.py b/snuba/admin/views.py index 0b68ed96c4..01007ac097 100644 --- a/snuba/admin/views.py +++ b/snuba/admin/views.py @@ -2,10 +2,11 @@ import io import sys +import time from contextlib import redirect_stdout from dataclasses import asdict from datetime import datetime -from typing import Any, List, Mapping, Optional, Sequence, Tuple, cast +from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple, cast import sentry_sdk import simplejson as json @@ -501,11 +502,43 @@ def clickhouse_trace_query() -> Response: for query_trace_data in parse_trace_for_query_ids(query_trace, storage): sql = profile_events_raw_sql.format(query_trace_data.query_id) logger.info( - "Profile event gathering host: {}, port = {}, storage = {}, sql = {}, g.user = {}".format( + "Gathering profile event using host: {}, port = {}, storage = {}, sql = {}, g.user = {}".format( query_trace_data.host, query_trace_data.port, storage, sql, g.user ) ) - # TODO: Onkar to add the profile event logic later. + system_query_result, counter = None, 0 + while counter < 60: + # There is a race between the trace query and the 'SELECT ProfileEvents...' query. ClickHouse does not immediately + # return the rows for 'SELECT ProfileEvents...' query. To make it return rows, sleep between the query executions. + system_query_result = run_system_query_on_host_with_sql( + query_trace_data.host, + int(query_trace_data.port), + storage, + sql, + False, + g.user, + ) + if not system_query_result.results: + time.sleep(1) + counter += 1 + else: + break + + if system_query_result is not None and len(system_query_result.results) > 0: + query_trace.profile_events_meta.append(system_query_result.meta) + query_trace.profile_events_profile = cast( + Dict[str, int], system_query_result.profile + ) + columns = system_query_result.meta + if columns: + res = {} + res["column_names"] = [name for name, _ in columns] + res["rows"] = [] + for query_result in system_query_result.results: + if query_result[0]: + res["rows"].append(json.dumps(query_result[0])) + query_trace.profile_events_results[query_trace_data.node_name] = res + return make_response(jsonify(asdict(query_trace)), 200) except InvalidCustomQuery as err: return make_response( @@ -520,6 +553,7 @@ def clickhouse_trace_query() -> Response: 400, ) except ClickhouseError as err: + logger.error(err, exc_info=True) details = { "type": "clickhouse", "message": str(err), @@ -527,6 +561,7 @@ def clickhouse_trace_query() -> Response: } return make_response(jsonify({"error": details}), 400) except Exception as err: + logger.error(err, exc_info=True) return make_response( jsonify({"error": {"type": "unknown", "message": str(err)}}), 500, @@ -536,7 +571,7 @@ def clickhouse_trace_query() -> Response: def parse_trace_for_query_ids( trace_output: TraceOutput, storage_key: str ) -> List[QueryTraceData]: - result = [] + traces = [] summarized_trace_output = trace_output.summarized_trace_output storage_info = get_storage_info() matched = next( @@ -545,16 +580,28 @@ def parse_trace_for_query_ids( if matched is not None: local_nodes = matched.get("local_nodes", []) query_node = matched.get("query_node", None) - result = [ - QueryTraceData( - host=local_nodes[0].get("host") if local_nodes else query_node.get("host"), # type: ignore - port=local_nodes[0].get("port") if local_nodes else query_node.get("port"), # type: ignore - query_id=query_summary.query_id, - node_name=node_name, - ) - for node_name, query_summary in summarized_trace_output.query_summaries.items() - ] - return result + for node_name, query_summary in summarized_trace_output.query_summaries.items(): + if local_nodes: + for local_node in local_nodes: + traces.append( + QueryTraceData( + host=local_node.get("host"), # type: ignore + port=local_node.get("port"), # type: ignore + query_id=query_summary.query_id, + node_name=node_name, + ) + ) + if query_node: + traces.append( + QueryTraceData( + host=query_node.get("host"), # type: ignore + port=query_node.get("port"), # type: ignore + query_id=query_summary.query_id, + node_name=node_name, + ) + ) + + return traces @application.route("/clickhouse_querylog_query", methods=["POST"]) diff --git a/tests/admin/test_api.py b/tests/admin/test_api.py index 014c302fc7..ca04de2f27 100644 --- a/tests/admin/test_api.py +++ b/tests/admin/test_api.py @@ -259,6 +259,7 @@ def test_query_trace(admin_api: FlaskClient) -> None: data = json.loads(response.data) assert " executeQuery" in data["trace_output"] assert "summarized_trace_output" in data + assert "profile_events_results" in data @pytest.mark.redis_db