Query MongoDB timestamps using normalized ISO string bounds
This commit is contained in:
parent
757eb97be0
commit
ebb3a1fb65
3 changed files with 47 additions and 2 deletions
8
great_ai/helper/normalise_datetime.py
Normal file
8
great_ai/helper/normalise_datetime.py
Normal file
|
|
@ -0,0 +1,8 @@
|
||||||
|
from datetime import datetime, timezone
|
||||||
|
|
||||||
|
|
||||||
|
def normalise_datetime(value: datetime) -> datetime:
|
||||||
|
"""Interpret naive timestamps as UTC and normalize aware timestamps to UTC."""
|
||||||
|
return (
|
||||||
|
value.astimezone(timezone.utc).replace(tzinfo=None) if value.tzinfo else value
|
||||||
|
)
|
||||||
|
|
@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple
|
||||||
|
|
||||||
from pymongo import ASCENDING, DESCENDING, MongoClient
|
from pymongo import ASCENDING, DESCENDING, MongoClient
|
||||||
|
|
||||||
|
from ..helper.normalise_datetime import normalise_datetime
|
||||||
from ..utilities import chunk
|
from ..utilities import chunk
|
||||||
from ..views import Filter, SortBy, Trace
|
from ..views import Filter, SortBy, Trace
|
||||||
from .tracing_database_driver import TracingDatabaseDriver
|
from .tracing_database_driver import TracingDatabaseDriver
|
||||||
|
|
@ -125,10 +126,14 @@ class MongoDbDriver(TracingDatabaseDriver):
|
||||||
and_query.append({})
|
and_query.append({})
|
||||||
|
|
||||||
if since:
|
if since:
|
||||||
and_query.append({"created": {"$gte": since}})
|
and_query.append(
|
||||||
|
{"created": {"$gte": normalise_datetime(since).isoformat()}}
|
||||||
|
)
|
||||||
|
|
||||||
if until:
|
if until:
|
||||||
and_query.append({"created": {"$lte": until}})
|
and_query.append(
|
||||||
|
{"created": {"$lte": normalise_datetime(until).isoformat()}}
|
||||||
|
)
|
||||||
|
|
||||||
if has_feedback is not None:
|
if has_feedback is not None:
|
||||||
and_query.append(
|
and_query.append(
|
||||||
|
|
|
||||||
32
tests/test_mongodb_queries.py
Normal file
32
tests/test_mongodb_queries.py
Normal file
|
|
@ -0,0 +1,32 @@
|
||||||
|
from datetime import datetime, timedelta, timezone
|
||||||
|
from typing import Any, Iterator, Tuple
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from great_ai import MongoDbDriver
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def mongo() -> Iterator[Tuple[MongoDbDriver, Any]]:
|
||||||
|
driver = MongoDbDriver.__new__(MongoDbDriver)
|
||||||
|
driver.mongo_connection_string = "mongodb://test"
|
||||||
|
driver.mongo_database = "test"
|
||||||
|
with patch("great_ai.persistence.mongodb_driver.MongoClient") as factory:
|
||||||
|
client = factory.__getitem__.return_value.return_value.__enter__.return_value
|
||||||
|
collection = client.__getitem__.return_value.traces
|
||||||
|
collection.count_documents.return_value = 3
|
||||||
|
collection.find.return_value.__enter__.return_value = []
|
||||||
|
yield driver, collection
|
||||||
|
|
||||||
|
|
||||||
|
def test_mongo_date_bounds_match_stored_iso_strings(
|
||||||
|
mongo: Tuple[MongoDbDriver, Any],
|
||||||
|
) -> None:
|
||||||
|
driver, collection = mongo
|
||||||
|
since = datetime(2026, 1, 1, 1, tzinfo=timezone(timedelta(hours=1)))
|
||||||
|
until = datetime(2026, 1, 2)
|
||||||
|
driver.query(since=since, until=until)
|
||||||
|
predicates = collection.find.call_args.kwargs["filter"]["$and"]
|
||||||
|
assert {"created": {"$gte": "2026-01-01T00:00:00"}} in predicates
|
||||||
|
assert {"created": {"$lte": "2026-01-02T00:00:00"}} in predicates
|
||||||
Loading…
Reference in a new issue