feat(pipeline): Capture Trace flag (#2293)

Added capture_trace flag to RequestSettings and schema
Updated execution pipeline to propagate flag through to ClickHouse execute() methods

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,
         )

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

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 = "<unknown>",
         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

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,
     },

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")

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")
 

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"