Protect trace fields from logged input collisions
This commit is contained in:
parent
1bfdf05081
commit
1d088669d1
2 changed files with 29 additions and 0 deletions
|
|
@ -76,6 +76,7 @@ class Trace(HashableBaseModel, Generic[T]):
|
||||||
return ",\n".join(self.tags)
|
return ",\n".join(self.tags)
|
||||||
|
|
||||||
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
|
def to_flat_dict(self, include_original: bool = True) -> Dict[str, Any]:
|
||||||
|
reserved = set(type(self).model_fields) | {"_id"}
|
||||||
return {
|
return {
|
||||||
**(
|
**(
|
||||||
self.model_dump()
|
self.model_dump()
|
||||||
|
|
@ -93,6 +94,7 @@ class Trace(HashableBaseModel, Generic[T]):
|
||||||
else pformat(v, indent=2, compact=True)
|
else pformat(v, indent=2, compact=True)
|
||||||
)
|
)
|
||||||
for k, v in self.logged_values.items()
|
for k, v in self.logged_values.items()
|
||||||
|
if k not in reserved
|
||||||
},
|
},
|
||||||
"models_flat": self.models_flat,
|
"models_flat": self.models_flat,
|
||||||
"exception_flat": self.exception_flat,
|
"exception_flat": self.exception_flat,
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,4 @@
|
||||||
|
from datetime import datetime
|
||||||
from typing import Any, Dict
|
from typing import Any, Dict
|
||||||
from unittest.mock import patch
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
|
@ -7,6 +8,32 @@ from pymongo import MongoClient
|
||||||
from pymongo.synchronous.collection import Collection
|
from pymongo.synchronous.collection import Collection
|
||||||
|
|
||||||
|
|
||||||
|
def test_flattening_preserves_reserved_trace_fields() -> None:
|
||||||
|
trace = Trace(
|
||||||
|
trace_id="actual-id",
|
||||||
|
created=datetime(2026, 1, 1).isoformat(),
|
||||||
|
original_execution_time_ms=0,
|
||||||
|
models=[],
|
||||||
|
tags=["online"],
|
||||||
|
output="prediction",
|
||||||
|
logged_values={
|
||||||
|
"trace_id": "input-id",
|
||||||
|
"created": "input-created",
|
||||||
|
"output": 123,
|
||||||
|
"tags": ["input"],
|
||||||
|
"models": "input",
|
||||||
|
"_id": "input-id",
|
||||||
|
"feature": 5,
|
||||||
|
},
|
||||||
|
)
|
||||||
|
flat = trace.to_flat_dict()
|
||||||
|
for name, value in trace.model_dump().items():
|
||||||
|
assert flat[name] == value
|
||||||
|
assert "_id" not in flat
|
||||||
|
assert flat["feature"] == 5
|
||||||
|
assert Trace.model_validate(flat) == trace
|
||||||
|
|
||||||
|
|
||||||
def test_mongo_update_uses_valid_replacement() -> None:
|
def test_mongo_update_uses_valid_replacement() -> None:
|
||||||
trace = Trace(
|
trace = Trace(
|
||||||
trace_id="id",
|
trace_id="id",
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue