fix: change last_updated type from String to DateTime

Author: sokoliva

src/a2a/migrations/versions/6419d2d130f6_add_columns_owner_last_updated.py

@@ -32,13 +32,18 @@ def _get_inspector() -> sa.engine.reflection.Inspector:
     return inspector
 
 
-def _add_column(table: str, value: str, column_name: str) -> None:
+def _add_column(
+    table: str,
+    column_name: str,
+    type_: sa.types.TypeEngine,
+    value: str | None = None,
+) -> None:
     if not _column_exists(table, column_name):
         op.add_column(
             table,
             sa.Column(
                 column_name,
-                sa.String(128),
+                type_,
                 nullable=False,
                 server_default=value,
             ),
@@ -107,8 +112,8 @@ def upgrade() -> None:
     )
 
     if _table_exists(tasks_table):
-        _add_column(tasks_table, owner, 'owner')
-        _add_column(tasks_table, '0', 'last_updated')
+        _add_column(tasks_table, 'owner', sa.String(128), owner)
+        _add_column(tasks_table, 'last_updated', sa.DateTime(timezone=True))
         _add_index(
             tasks_table,
             f'idx_{tasks_table}_owner_last_updated',
@@ -120,7 +125,9 @@ def upgrade() -> None:
         )
 
     if _table_exists(push_notification_configs_table):
-        _add_column(push_notification_configs_table, owner, 'owner')
+        _add_column(
+            push_notification_configs_table, 'owner', sa.String(128), owner
+        )
     else:
         logging.warning(
             f"Table '{push_notification_configs_table}' does not exist. Skipping upgrade for this table."

src/a2a/server/models.py

@@ -1,3 +1,4 @@
+from datetime import datetime
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 
@@ -18,7 +19,7 @@ def override(func):  # noqa: ANN001, ANN201
 
 
 try:
-    from sqlalchemy import JSON, Dialect, Index, LargeBinary, String
+    from sqlalchemy import JSON, DateTime, Dialect, Index, LargeBinary, String
     from sqlalchemy.orm import (
         DeclarativeBase,
         Mapped,
@@ -149,7 +150,9 @@ class TaskMixin:
         String(16), nullable=False, default='task'
     )
     owner: Mapped[str] = mapped_column(String(128), nullable=False)
-    last_updated: Mapped[str] = mapped_column(String(22), nullable=True)
+    last_updated: Mapped[datetime | None] = mapped_column(
+        DateTime(timezone=True), nullable=True
+    )
 
     # Properly typed Pydantic fields with automatic serialization
     status: Mapped[TaskStatus] = mapped_column(PydanticType(TaskStatus))

src/a2a/server/tasks/database_task_store.py

@@ -124,7 +124,7 @@ def _to_orm(self, task: Task, owner: str) -> TaskModel:
             kind='task',  # Default kind for tasks
             owner=owner,
             last_updated=(
-                task.status.timestamp.ToJsonString()
+                task.status.timestamp.ToDatetime()
                 if task.HasField('status') and task.status.HasField('timestamp')
                 else None
             ),
@@ -227,21 +227,16 @@ async def list(
                     == a2a_pb2.TaskState.Name(params.status)
                 )
             if params.HasField('status_timestamp_after'):
-                last_updated_after_iso = (
-                    params.status_timestamp_after.ToJsonString()
-                )
-                base_stmt = base_stmt.where(
-                    timestamp_col >= last_updated_after_iso
-                )
+                last_updated_after = params.status_timestamp_after.ToDatetime()
+                base_stmt = base_stmt.where(timestamp_col >= last_updated_after)
 
             # Get total count
             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 empty strings,
-            # which sort last in descending order
+            # Use nulls_last() to ensure NULL timestamps sort last in descending order
             stmt = base_stmt.order_by(
-                func.coalesce(timestamp_col, '').desc(),
+                timestamp_col.desc().nulls_last(),
                 self.task_model.id.desc(),
             )
 

tests/migrations/versions/test_migration_6419d2d130f6.py

@@ -102,6 +102,7 @@ def test_migration_6419d2d130f6_full_cycle(
     tasks_columns = {row[1]: row for row in cursor.fetchall()}
     assert 'owner' in tasks_columns
     assert 'last_updated' in tasks_columns
+    assert tasks_columns['last_updated'][2] == 'DATETIME'
 
     # Check default value for owner in tasks
     # row[4] is dflt_value in PRAGMA table_info