diff --git a/snuba/clickhouse/native.py b/snuba/clickhouse/native.py index 6baf267861..11d387273b 100644 --- a/snuba/clickhouse/native.py +++ b/snuba/clickhouse/native.py @@ -72,6 +72,7 @@ def execute( settings: Optional[Mapping[str, Any]] = None, types_check: bool = False, columnar: bool = False, + capture_trace: bool = False, ) -> ClickhouseResult: """ Execute a clickhouse query with a single quick retry in case of @@ -142,6 +143,7 @@ def execute_robust( settings: Optional[Mapping[str, Any]] = None, types_check: bool = False, columnar: bool = False, + capture_trace: bool = False, ) -> ClickhouseResult: """ Execute a clickhouse query with a bit more tenacity. Make more retry @@ -164,6 +166,7 @@ def execute_robust( settings=settings, types_check=types_check, columnar=columnar, + capture_trace=capture_trace, ) except (errors.NetworkError, errors.SocketTimeoutError, EOFError) as e: # Try 3 times on connection issues. @@ -308,6 +311,7 @@ def execute( settings: Optional[Mapping[str, str]] = None, with_totals: bool = False, robust: bool = False, + capture_trace: bool = False, ) -> Result: settings = {**settings} if settings is not None else {} @@ -325,6 +329,7 @@ def execute( with_column_types=True, query_id=query_id, settings=settings, + capture_trace=capture_trace, ), with_totals=with_totals, ) diff --git a/snuba/reader.py b/snuba/reader.py index 60ad7f0074..66284632a6 100644 --- a/snuba/reader.py +++ b/snuba/reader.py @@ -118,6 +118,7 @@ def execute( settings: Optional[Mapping[str, str]] = None, with_totals: bool = False, robust: bool = False, + capture_trace: bool = False, ) -> Result: """Execute a query.""" raise NotImplementedError diff --git a/snuba/request/request_settings.py b/snuba/request/request_settings.py index 284d5b3506..05c0565f61 100644 --- a/snuba/request/request_settings.py +++ b/snuba/request/request_settings.py @@ -50,6 +50,10 @@ def get_rate_limit_params(self) -> Sequence[RateLimitParameters]: def add_rate_limit(self, rate_limit_param: RateLimitParameters) -> None: pass + @abstractmethod + def get_capture_trace(self) -> bool: + pass + class HTTPRequestSettings(RequestSettings): """ @@ -67,6 +71,7 @@ def __init__( parent_api: str = "", dry_run: bool = False, legacy: bool = False, + capture_trace: bool = False, ) -> None: super().__init__(referrer=referrer) self.__turbo = turbo @@ -76,6 +81,7 @@ def __init__( self.__dry_run = dry_run self.__legacy = legacy self.__rate_limit_params: List[RateLimitParameters] = [] + self.__capture_trace = capture_trace def get_turbo(self) -> bool: return self.__turbo @@ -101,6 +107,9 @@ def get_rate_limit_params(self) -> Sequence[RateLimitParameters]: def add_rate_limit(self, rate_limit_param: RateLimitParameters) -> None: self.__rate_limit_params.append(rate_limit_param) + def get_capture_trace(self) -> bool: + return self.__capture_trace + class SubscriptionRequestSettings(RequestSettings): """ @@ -137,3 +146,6 @@ def get_rate_limit_params(self) -> Sequence[RateLimitParameters]: def add_rate_limit(self, rate_limit_param: RateLimitParameters) -> None: pass + + def get_capture_trace(self) -> bool: + return False diff --git a/snuba/request/schema.py b/snuba/request/schema.py index 74b37757c4..e18d22ca62 100644 --- a/snuba/request/schema.py +++ b/snuba/request/schema.py @@ -133,6 +133,10 @@ def generate_template(self) -> Any: "dry_run": {"type": "boolean", "default": False}, # Flags if this a legacy query that was automatically generated by the SnQL SDK "legacy": {"type": "boolean", "default": False}, + # Flag to retrieve only ClickHouse tracing data. If this is True, + # the query execution should not return the results of the query, but + # return the tracing data of the query execution within ClickHouse + "capture_trace": {"type": "boolean", "default": False}, }, "additionalProperties": False, }, diff --git a/snuba/web/db_query.py b/snuba/web/db_query.py index 0d9157d55b..93c54c4731 100644 --- a/snuba/web/db_query.py +++ b/snuba/web/db_query.py @@ -238,6 +238,7 @@ def execute_query( query_settings, with_totals=clickhouse_query.has_totals(), robust=robust, + capture_trace=request_settings.get_capture_trace(), ) timer.mark("execute") diff --git a/tests/clusters/fake_cluster.py b/tests/clusters/fake_cluster.py index 612995ebd0..90488b823e 100644 --- a/tests/clusters/fake_cluster.py +++ b/tests/clusters/fake_cluster.py @@ -27,6 +27,7 @@ def execute( settings: Optional[Mapping[str, Any]] = None, types_check: bool = False, columnar: bool = False, + capture_trace: bool = False, ) -> ClickhouseResult: self.__queries.append(query) return ClickhouseResult([[1]]) @@ -49,6 +50,7 @@ def execute( settings: Optional[Mapping[str, Any]] = None, types_check: bool = False, columnar: bool = False, + capture_trace: bool = False, ) -> ClickhouseResult: raise ServerExplodedException("The server exploded") diff --git a/tests/test_snql_api.py b/tests/test_snql_api.py index 594a9664c1..05ee99c626 100644 --- a/tests/test_snql_api.py +++ b/tests/test_snql_api.py @@ -596,6 +596,25 @@ def test_valid_columns_composite_query(self) -> None: ) assert response.status_code == 200 + def test_capture_trace_flag(self) -> None: + response = self.post( + "/events/snql", + data=json.dumps( + { + "query": """ + MATCH (discover) + SELECT count() AS count + WHERE + timestamp >= toDateTime('2021-01-01') AND + timestamp < toDateTime('2022-01-01') AND + project_id IN tuple(5433960) + """, + "capture_trace": True, + } + ), + ) + assert response.status_code == 200 + MATCH = "MATCH (e: events) -[grouped]-> (gm: groupedmessage)" SELECT = "SELECT e.group_id, gm.status, avg(e.retention_days) AS avg BY e.group_id, gm.status" WHERE = "WHERE e.project_id = 1 AND gm.project_id = 1"