Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 5 additions & 8 deletions src/a2a/server/tasks/database_task_store.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import logging

from collections.abc import Callable
from datetime import datetime, timezone


try:
from sqlalchemy import Table, and_, delete, func, or_, select
from sqlalchemy import Table, and_, case, delete, func, or_, select
from sqlalchemy.ext.asyncio import (
AsyncEngine,
AsyncSession,
Expand Down Expand Up @@ -245,13 +244,11 @@ async def list(
count_stmt = select(func.count()).select_from(base_stmt.alias())
total_count = (await session.execute(count_stmt)).scalar_one()

# Use coalesce to treat NULL timestamps as datetime.min,
# which sort last in descending order
# Sort NULL timestamps last without binding a sentinel value, which
# may fall outside a database's supported datetime range.
stmt = base_stmt.order_by(
func.coalesce(
timestamp_col,
datetime.min.replace(tzinfo=timezone.utc),
).desc(),
case((timestamp_col.is_(None), 1), else_=0).asc(),
timestamp_col.desc(),
self.task_model.id.desc(),
)

Expand Down
48 changes: 47 additions & 1 deletion tests/server/tasks/test_database_task_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

from collections.abc import AsyncGenerator
from datetime import datetime, timezone
from typing import Any
from unittest.mock import MagicMock

import pytest
Expand All @@ -10,7 +11,8 @@
from _pytest.mark.structures import ParameterSet
from a2a.compat.v0_3 import types as types_v03
from a2a.types.a2a_pb2 import ListTasksRequest
from sqlalchemy import insert
from sqlalchemy import event, insert
from sqlalchemy.engine import Connection, ExecutionContext


# Skip entire test module if SQLAlchemy is not installed
Expand Down Expand Up @@ -344,6 +346,50 @@ async def test_list_tasks(
await db_store_parameterized.delete(task.id, TEST_CONTEXT)


@pytest.mark.asyncio
async def test_list_tasks_does_not_bind_out_of_range_datetimes(
db_store_parameterized: DatabaseTaskStore,
) -> None:
"""Test list ordering without out-of-range datetime sentinels."""
bound_datetimes: list[datetime] = []
compiled_execution_count = 0

def capture_bound_datetimes(
_conn: Connection,
_cursor: Any,
_statement: str,
_parameters: Any,
context: ExecutionContext,
_executemany: bool,
) -> None:
nonlocal compiled_execution_count
if context.compiled is None:
return
compiled_execution_count += 1
bound_datetimes.extend(
value
for value in context.compiled.params.values()
if isinstance(value, datetime)
)

event.listen(
db_store_parameterized.engine.sync_engine,
'before_cursor_execute',
capture_bound_datetimes,
)
try:
await db_store_parameterized.list(ListTasksRequest(), TEST_CONTEXT)
finally:
event.remove(
db_store_parameterized.engine.sync_engine,
'before_cursor_execute',
capture_bound_datetimes,
)

assert compiled_execution_count > 0
assert all(value.year >= 1000 for value in bound_datetimes)


@pytest.mark.asyncio
@pytest.mark.parametrize(
'params, expected_error_message',
Expand Down
Loading