diff --git a/great_ai/persistence/mongodb_driver.py b/great_ai/persistence/mongodb_driver.py index 6ea4d94..5c56a92 100644 --- a/great_ai/persistence/mongodb_driver.py +++ b/great_ai/persistence/mongodb_driver.py @@ -143,6 +143,8 @@ class MongoDbDriver(TracingDatabaseDriver): with MongoClient[Any](self.mongo_connection_string) as client: count = client[self.mongo_database].traces.count_documents(**query) + if take == 0: + return [], count if skip: query["skip"] = skip diff --git a/tests/test_mongodb_queries.py b/tests/test_mongodb_queries.py index 048fe0f..e2b984a 100644 --- a/tests/test_mongodb_queries.py +++ b/tests/test_mongodb_queries.py @@ -30,3 +30,11 @@ def test_mongo_date_bounds_match_stored_iso_strings( 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 + + +def test_zero_take_preserves_count_without_fetching( + mongo: Tuple[MongoDbDriver, Any], +) -> None: + driver, collection = mongo + assert driver.query(take=0) == ([], 3) + collection.find.assert_not_called()