Normalize timezone-aware bounds and stored TinyDB timestamps
This commit is contained in:
parent
ebb3a1fb65
commit
d06e982b46
2 changed files with 27 additions and 2 deletions
|
|
@ -6,6 +6,7 @@ import pandas as pd
|
|||
from filelock import FileLock
|
||||
from tinydb import TinyDB
|
||||
|
||||
from ..helper.normalise_datetime import normalise_datetime
|
||||
from ..views import Filter, SortBy, Trace
|
||||
from .tracing_database_driver import TracingDatabaseDriver
|
||||
|
||||
|
|
@ -52,11 +53,15 @@ class ParallelTinyDbDriver(TracingDatabaseDriver):
|
|||
has_feedback: Optional[bool] = None,
|
||||
sort_by: Sequence[SortBy] = [],
|
||||
) -> Tuple[List[Trace], int]:
|
||||
since = normalise_datetime(since) if since is not None else None
|
||||
until = normalise_datetime(until) if until is not None else None
|
||||
|
||||
def does_match(d: Dict[str, Any]) -> bool:
|
||||
created = normalise_datetime(datetime.fromisoformat(d["created"]))
|
||||
return (
|
||||
not set(conjunctive_tags) - set(d["tags"])
|
||||
and (since is None or datetime.fromisoformat(d["created"]) >= since)
|
||||
and (until is None or datetime.fromisoformat(d["created"]) <= until)
|
||||
and (since is None or created >= since)
|
||||
and (until is None or created <= until)
|
||||
and (
|
||||
has_feedback is None or has_feedback == (d["feedback"] is not None)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import multiprocessing as mp
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest.mock import patch
|
||||
|
|
@ -7,6 +8,7 @@ from unittest.mock import patch
|
|||
from tinydb.table import Table
|
||||
|
||||
from great_ai import ParallelTinyDbDriver, Trace
|
||||
from great_ai.context import get_context
|
||||
|
||||
|
||||
def write_traces(path: str, index: int, barrier: Any) -> None:
|
||||
|
|
@ -54,3 +56,21 @@ def test_spawned_writers_share_database_lock(tmp_path: Path) -> None:
|
|||
process.close()
|
||||
with patch.object(ParallelTinyDbDriver, "path_to_db", path):
|
||||
assert len(ParallelTinyDbDriver().query()[0]) == 6
|
||||
|
||||
|
||||
def test_tinydb_date_queries_normalize_timezones(isolated_context: None) -> None:
|
||||
database = get_context().tracing_database
|
||||
for i, created in enumerate(("2026-01-01T12:00:00", "2026-01-01T14:00:00+02:00")):
|
||||
database.save(
|
||||
Trace(
|
||||
trace_id=str(i),
|
||||
created=created,
|
||||
original_execution_time_ms=0,
|
||||
models=[],
|
||||
tags=[],
|
||||
logged_values={},
|
||||
)
|
||||
)
|
||||
noon = datetime(2026, 1, 1, 12, tzinfo=timezone.utc)
|
||||
assert database.query(since=noon, until=noon)[1] == 2
|
||||
assert database.query(since=noon.replace(hour=13))[1] == 0
|
||||
|
|
|
|||
Loading…
Reference in a new issue