test(traces): align native tests with named read queries

tests/test_litellm_rust/test_traces.py

@@ -1,6 +1,7 @@
 import base64
 import gzip
 import json
+import time
 from typing import Final
 from urllib.parse import parse_qs, urlsplit
 
@@ -14,16 +15,24 @@ pytestmark = pytest.mark.requires_rust_extension
 
 @pytest.mark.asyncio
 async def test_trace_reader_projects_connection_and_parameters(recording_server: RecordingServer) -> None:
-    recording_server.enqueue(ResponseSpec(body={"data": [{"trace_id": "trace-1"}]}))
+    recording_server.enqueue(ResponseSpec(body={"data": [{"span_id": "span-1"}]}))
     reader_url: Final = recording_server.base_url.replace("http://", "http://reader:p%40ss%2Fword%25@")
     storage: Final = NativeTraceStorage("trace_test", recording_server.base_url, reader_url + "?database=wrong")
-    rows: Final = json.loads(await storage.query("SELECT {trace_id:String} AS trace_id", {"trace_id": "trace-1"}))
+    rows: Final = json.loads(
+        await storage.query(
+            "trace_spans", {"trace_id": "trace-1", "team_ids": [], "api_key_hash": "", "trace_ref": ""}
+        )
+    )
     request: Final = recording_server.requests[0]
-    parameters: Final = parse_qs(urlsplit(request.path).query)
-    assert rows == [{"trace_id": "trace-1"}]
-    assert request.raw_body == b"SELECT {trace_id:String} AS trace_id"
+    parameters: Final = parse_qs(urlsplit(request.path).query, keep_blank_values=True)
+    assert rows == {"data": [{"span_id": "span-1"}]}
+    assert b"FROM otel_traces AS o" in request.raw_body
+    assert b"WHERE o.TraceId = {trace_id:String}" in request.raw_body
     assert parameters["database"] == ["trace_test"]
     assert parameters["param_trace_id"] == ["trace-1"]
+    assert parameters["param_team_ids"] == ["[]"]
+    assert parameters["param_api_key_hash"] == [""]
+    assert parameters["param_trace_ref"] == [""]
     assert parameters["readonly"] == ["1"]
     assert "user" not in parameters
     assert "password" not in parameters
@@ -35,6 +44,16 @@ async def test_trace_reader_rejects_success_status_with_embedded_error(recording
     recording_server.enqueue(ResponseSpec(body={"data": [], "exception": "query failed"}))
     storage: Final = NativeTraceStorage("trace_test", recording_server.base_url, recording_server.base_url)
     with pytest.raises(RuntimeError, match="invalid or failed JSON"):
+        await storage.query(
+            "trace_spans", {"trace_id": "trace-1", "team_ids": [], "api_key_hash": "", "trace_ref": ""}
+        )
+
+
+@pytest.mark.asyncio
+async def test_trace_reader_rejects_arbitrary_sql_before_sending(recording_server: RecordingServer) -> None:
+    recording_server.expected_requests = 0
+    storage: Final = NativeTraceStorage("trace_test", recording_server.base_url, recording_server.base_url)
+    with pytest.raises(ValueError, match="unknown ClickHouse read query"):
         await storage.query("SELECT 1", {})
 
 
@@ -73,9 +92,15 @@ async def test_schema_setup_uses_writer_credentials_and_rejects_failed_statement
 async def test_insert_encodes_and_sends_rows(recording_server: RecordingServer) -> None:
     recording_server.enqueue(ResponseSpec(body=""))
     storage: Final = NativeTraceStorage("trace_test", recording_server.base_url)
-    await storage.insert_rows("otel_traces", [{"Timestamp": 1_234_567_890, "Input": "hello"}])
+    before_insert_ms: Final = time.time_ns() // 1_000_000
+    await storage.insert_rows("otel_traces", [{"Timestamp": 1_234_567_890, "Input": "hello", "EngineReceivedMs": 0}])
+    after_insert_ms: Final = time.time_ns() // 1_000_000
     request: Final = recording_server.requests[0]
-    assert json.loads(gzip.decompress(request.raw_body)) == {
+    row: Final = json.loads(gzip.decompress(request.raw_body))
+    assert type(row["EngineReceivedMs"]) is int
+    assert before_insert_ms <= row["EngineReceivedMs"] <= after_insert_ms
+    assert row == {
+        "EngineReceivedMs": row["EngineReceivedMs"],
         "Input": "hello",
         "Timestamp": "1970-01-01T00:00:01.23456789Z",
     }