From acbf3edb149b0b2fd330fca40708bd96f5068fcd Mon Sep 17 00:00:00 2001 From: gtarpenning Date: Mon, 6 Jan 2025 11:31:51 -0800 Subject: [PATCH 1/5] chore(weave): add query and parameters to failed calls query stream requests --- .../clickhouse_trace_server_batched.py | 40 ++++++++++--------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py index e97e020b1abe..307366426c09 100644 --- a/weave/trace_server/clickhouse_trace_server_batched.py +++ b/weave/trace_server/clickhouse_trace_server_batched.py @@ -1603,24 +1603,28 @@ def _query_stream( summary = None parameters = _process_parameters(parameters) - with self.ch_client.query_rows_stream( - query, - parameters=parameters, - column_formats=column_formats, - use_none=True, - settings=settings, - ) as stream: - if isinstance(stream.source, QueryResult): - summary = stream.source.summary - logger.info( - "clickhouse_stream_query", - extra={ - "query": query, - "parameters": parameters, - "summary": summary, - }, - ) - yield from stream + try: + with self.ch_client.query_rows_stream( + query, + parameters=parameters, + column_formats=column_formats, + use_none=True, + settings=settings, + ) as stream: + if isinstance(stream.source, QueryResult): + summary = stream.source.summary + logger.info( + "clickhouse_stream_query", + extra={ + "query": query, + "parameters": parameters, + "summary": summary, + }, + ) + yield from stream + except Exception as e: + logger.exception(f"{e}", extra={"query": query, "parameters": parameters}) + raise e def _query( self, From 0ae2913827eb3aa15907d21c9b9f30f7143dcc88 Mon Sep 17 00:00:00 2001 From: gtarpenning Date: Mon, 6 Jan 2025 17:16:39 -0800 Subject: [PATCH 2/5] slightly better --- weave/trace_server/clickhouse_trace_server_batched.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py index 307366426c09..9ff4a78d4a17 100644 --- a/weave/trace_server/clickhouse_trace_server_batched.py +++ b/weave/trace_server/clickhouse_trace_server_batched.py @@ -1623,8 +1623,11 @@ def _query_stream( ) yield from stream except Exception as e: - logger.exception(f"{e}", extra={"query": query, "parameters": parameters}) - raise e + logger.exception( + f"_query_stream exception: {e}", + extra={"query": query, "parameters": parameters}, + ) + raise def _query( self, From 34e13d3b2d0d618f6e7dafa80fab500bfe7379b3 Mon Sep 17 00:00:00 2001 From: gtarpenning Date: Thu, 16 Jan 2025 11:15:30 -0800 Subject: [PATCH 3/5] 1 mib --- weave/trace_server/clickhouse_trace_server_batched.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py index 9ff4a78d4a17..0c0fd9328526 100644 --- a/weave/trace_server/clickhouse_trace_server_batched.py +++ b/weave/trace_server/clickhouse_trace_server_batched.py @@ -173,7 +173,8 @@ class NotFoundError(Exception): ENTITY_TOO_LARGE_PAYLOAD = '{"_weave": {"error":""}}' CLICKHOUSE_DEFAULT_QUERY_SETTINGS = { - "max_memory_usage": 16 * 1024 * 1024 * 1024, # 16 GiB + # "max_memory_usage": 16 * 1024 * 1024 * 1024, # 16 GiB + "max_memory_usage": 1 * 1024 * , # } From 22d4cb48af743e22d73fe55a3b3cceb61a6cd3dc Mon Sep 17 00:00:00 2001 From: gtarpenning Date: Thu, 16 Jan 2025 11:33:18 -0800 Subject: [PATCH 4/5] lint --- weave/trace_server/clickhouse_trace_server_batched.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py index 0c0fd9328526..3eab1aa40358 100644 --- a/weave/trace_server/clickhouse_trace_server_batched.py +++ b/weave/trace_server/clickhouse_trace_server_batched.py @@ -174,7 +174,7 @@ class NotFoundError(Exception): CLICKHOUSE_DEFAULT_QUERY_SETTINGS = { # "max_memory_usage": 16 * 1024 * 1024 * 1024, # 16 GiB - "max_memory_usage": 1 * 1024 * , # + "max_memory_usage": 1 * 1024 * 1024, # 1 MiB } From 9341c62a0e3110d69d340e37a01a3b08db593d63 Mon Sep 17 00:00:00 2001 From: gtarpenning Date: Thu, 16 Jan 2025 13:51:10 -0800 Subject: [PATCH 5/5] explicit error --- weave/trace_server/clickhouse_trace_server_batched.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/weave/trace_server/clickhouse_trace_server_batched.py b/weave/trace_server/clickhouse_trace_server_batched.py index 8aa1965becb1..3aebed8e4868 100644 --- a/weave/trace_server/clickhouse_trace_server_batched.py +++ b/weave/trace_server/clickhouse_trace_server_batched.py @@ -1716,11 +1716,9 @@ def _query_stream( ) yield from stream except Exception as e: - logger.exception( - f"_query_stream exception: {e}", - extra={"query": query, "parameters": parameters}, - ) - raise + # wrap exception with query and parameters + extra = {"query": query, "parameters": parameters} + raise type(e)(f"{str(e)} - Query: {query}, Parameters: {parameters}") from None def _query( self,