feat: Zero-downtime support for Database migration (#831)

# Description
The old 0.3 version is not able to read the 1.0 entries from database
because of the inconsistencies between 1.0 and 0.3 data types. This
applies to both DatabaseTaskStore and
DatabasePushNotificationConfigStore. This PR fixes this issue by
allowing users to write 0.3 compatible entires during migration period.

## Changes
- adds new conversion methods to `compat/0_3/conversions.py`
- update `DatabaseTaskStore` and `DatabasePushNotificationConfigStore`
to accept new conversion methods
- utilize new conversion methods of `DatabaseTaskStore` and
`DatabasePushNotificationConfigStore`

## Tested
Created a database using `0.3` spec containing populated tables `task`
and `push_notification_configs`. Ran `uv run a2a-db` using `1.0` spec
against the database and added new entries using the new Zero-downtime
feature, `DatabaseTaskStore.core_to_model_conversion =
core_to_compat_task_model` and
`DatabasePushNotificationConfigStore.core_to_model_conversion =
core_to_compat_push_notification_config_model`. Succesfully read new
entries using `0.3` spec.


Fixes #811 🦕

Author: sokoliva

src/a2a/compat/v0_3/conversions.py

@@ -1,10 +1,15 @@
 import base64
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from cryptography.fernet import Fernet
 
 from google.protobuf.json_format import MessageToDict, ParseDict
 
 from a2a.compat.v0_3 import types as types_v03
+from a2a.server.models import PushNotificationConfigModel, TaskModel
 from a2a.types import a2a_pb2 as pb2_v10
 
 
@@ -1367,3 +1372,77 @@ def to_compat_get_extended_agent_card_request(
 ) -> types_v03.GetAuthenticatedExtendedCardRequest:
     """Convert get extended agent card request to v0.3 compat type."""
     return types_v03.GetAuthenticatedExtendedCardRequest(id=request_id)
+
+
+def core_to_compat_task_model(task: pb2_v10.Task, owner: str) -> TaskModel:
+    """Converts a 1.0 core Task to a TaskModel using v0.3 JSON structure."""
+    compat_task = to_compat_task(task)
+    data = compat_task.model_dump(mode='json')
+
+    return TaskModel(
+        id=task.id,
+        context_id=task.context_id,
+        owner=owner,
+        status=data.get('status'),
+        history=data.get('history'),
+        artifacts=data.get('artifacts'),
+        task_metadata=data.get('metadata'),
+        protocol_version='0.3',
+    )
+
+
+def compat_task_model_to_core(task_model: TaskModel) -> pb2_v10.Task:
+    """Converts a TaskModel with v0.3 structure to a 1.0 core Task."""
+    compat_task = types_v03.Task(
+        id=task_model.id,
+        context_id=task_model.context_id,
+        status=types_v03.TaskStatus.model_validate(task_model.status),
+        artifacts=(
+            [types_v03.Artifact.model_validate(a) for a in task_model.artifacts]
+            if task_model.artifacts
+            else []
+        ),
+        history=(
+            [types_v03.Message.model_validate(h) for h in task_model.history]
+            if task_model.history
+            else []
+        ),
+        metadata=task_model.task_metadata,
+    )
+    return to_core_task(compat_task)
+
+
+def core_to_compat_push_notification_config_model(
+    task_id: str,
+    config: pb2_v10.TaskPushNotificationConfig,
+    owner: str,
+    fernet: 'Fernet | None' = None,
+) -> PushNotificationConfigModel:
+    """Converts a 1.0 core TaskPushNotificationConfig to a PushNotificationConfigModel using v0.3 JSON structure."""
+    compat_config = to_compat_push_notification_config(config)
+
+    json_payload = compat_config.model_dump_json().encode('utf-8')
+    data_to_store = fernet.encrypt(json_payload) if fernet else json_payload
+
+    return PushNotificationConfigModel(
+        task_id=task_id,
+        config_id=config.id,
+        owner=owner,
+        config_data=data_to_store,
+        protocol_version='0.3',
+    )
+
+
+def compat_push_notification_config_model_to_core(
+    model_instance: str, task_id: str
+) -> pb2_v10.TaskPushNotificationConfig:
+    """Converts a PushNotificationConfigModel with v0.3 structure back to a 1.0 core TaskPushNotificationConfig."""
+    inner_config = types_v03.PushNotificationConfig.model_validate_json(
+        model_instance
+    )
+    return to_core_task_push_notification_config(
+        types_v03.TaskPushNotificationConfig(
+            task_id=task_id,
+            push_notification_config=inner_config,
+        )
+    )

src/a2a/server/tasks/database_push_notification_config_store.py

@@ -13,9 +13,7 @@
         AsyncSession,
         async_sessionmaker,
     )
-    from sqlalchemy.orm import (
-        class_mapper,
-    )
+    from sqlalchemy.orm import class_mapper
 except ImportError as e:
     raise ImportError(
         'DatabasePushNotificationConfigStore requires SQLAlchemy and a database driver. '
@@ -26,8 +24,11 @@
         "or 'pip install a2a-sdk[sql]'"
     ) from e
 
-from a2a.compat.v0_3 import conversions
-from a2a.compat.v0_3 import types as types_v03
+from collections.abc import Callable
+
+from a2a.compat.v0_3.conversions import (
+    compat_push_notification_config_model_to_core,
+)
 from a2a.server.context import ServerCallContext
 from a2a.server.models import (
     Base,
@@ -44,7 +45,6 @@
 if TYPE_CHECKING:
     from cryptography.fernet import Fernet
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -61,14 +61,34 @@ class DatabasePushNotificationConfigStore(PushNotificationConfigStore):
     config_model: type[PushNotificationConfigModel]
     _fernet: 'Fernet | None'
     owner_resolver: OwnerResolver
+    core_to_model_conversion: (
+        Callable[
+            [str, TaskPushNotificationConfig, str, 'Fernet | None'],
+            PushNotificationConfigModel,
+        ]
+        | None
+    )
+    model_to_core_conversion: (
+        Callable[[PushNotificationConfigModel], TaskPushNotificationConfig]
+        | None
+    )
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         engine: AsyncEngine,
         create_table: bool = True,
         table_name: str = 'push_notification_configs',
         encryption_key: str | bytes | None = None,
         owner_resolver: OwnerResolver = resolve_user_scope,
+        core_to_model_conversion: Callable[
+            [str, TaskPushNotificationConfig, str, 'Fernet | None'],
+            PushNotificationConfigModel,
+        ]
+        | None = None,
+        model_to_core_conversion: Callable[
+            [PushNotificationConfigModel], TaskPushNotificationConfig
+        ]
+        | None = None,
     ) -> None:
         """Initializes the DatabasePushNotificationConfigStore.
 
@@ -80,6 +100,8 @@ def __init__(
                 If provided, `config_data` will be encrypted in the database.
                 The key must be a URL-safe base64-encoded 32-byte key.
             owner_resolver: Function to resolve the owner from the context.
+            core_to_model_conversion: Optional function to convert a TaskPushNotificationConfig to a TaskPushNotificationConfigModel.
+            model_to_core_conversion: Optional function to convert a TaskPushNotificationConfigModel to a TaskPushNotificationConfig.
         """
         logger.debug(
             'Initializing DatabasePushNotificationConfigStore with existing engine, table: %s',
@@ -98,6 +120,8 @@ def __init__(
             else create_push_notification_config_model(table_name)
         )
         self._fernet = None
+        self.core_to_model_conversion = core_to_model_conversion
+        self.model_to_core_conversion = model_to_core_conversion
 
         if encryption_key:
             try:
@@ -152,6 +176,11 @@ def _to_orm(
 
         The config data is serialized to JSON bytes, and encrypted if a key is configured.
         """
+        if self.core_to_model_conversion:
+            return self.core_to_model_conversion(
+                task_id, config, owner, self._fernet
+            )
+
         json_payload = MessageToJson(config).encode('utf-8')
 
         if self._fernet:
@@ -174,6 +203,9 @@ def _from_orm(
 
         Handles decryption if a key is configured, with a fallback to plain JSON.
         """
+        if self.model_to_core_conversion:
+            return self.model_to_core_conversion(model_instance)
+
         payload = model_instance.config_data
 
         if self._fernet:
@@ -359,12 +391,7 @@ def _parse_config(
         """
         if protocol_version == '1.0':
             return Parse(json_payload, TaskPushNotificationConfig())
-        inner_config = types_v03.PushNotificationConfig.model_validate_json(
-            json_payload
-        )
-        return conversions.to_core_task_push_notification_config(
-            types_v03.TaskPushNotificationConfig(
-                task_id=task_id or '',
-                push_notification_config=inner_config,
-            )
+
+        return compat_push_notification_config_model_to_core(
+            json_payload, task_id or ''
         )

src/a2a/server/tasks/database_task_store.py

@@ -1,25 +1,17 @@
 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_, delete, func, or_, select
     from sqlalchemy.ext.asyncio import (
         AsyncEngine,
         AsyncSession,
         async_sessionmaker,
     )
-    from sqlalchemy.orm import (
-        class_mapper,
-    )
+    from sqlalchemy.orm import class_mapper
 except ImportError as e:
     raise ImportError(
         'DatabaseTaskStore requires SQLAlchemy and a database driver. '
@@ -29,11 +21,11 @@
         "'pip install a2a-sdk[sqlite]', "
         "or 'pip install a2a-sdk[sql]'"
     ) from e
-
 from google.protobuf.json_format import MessageToDict, ParseDict
 
-from a2a.compat.v0_3 import conversions
-from a2a.compat.v0_3 import types as types_v03
+from a2a.compat.v0_3.conversions import (
+    compat_task_model_to_core,
+)
 from a2a.server.context import ServerCallContext
 from a2a.server.models import Base, TaskModel, create_task_model
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
@@ -60,13 +52,18 @@ class DatabaseTaskStore(TaskStore):
     _initialized: bool
     task_model: type[TaskModel]
     owner_resolver: OwnerResolver
+    core_to_model_conversion: Callable[[Task, str], TaskModel] | None = None
+    model_to_core_conversion: Callable[[TaskModel], Task] | None = None
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         engine: AsyncEngine,
         create_table: bool = True,
         table_name: str = 'tasks',
         owner_resolver: OwnerResolver = resolve_user_scope,
+        core_to_model_conversion: Callable[[Task, str], TaskModel]
+        | None = None,
+        model_to_core_conversion: Callable[[TaskModel], Task] | None = None,
     ) -> None:
         """Initializes the DatabaseTaskStore.
 
@@ -75,6 +72,8 @@ def __init__(
             create_table: If true, create tasks table on initialization.
             table_name: Name of the database table. Defaults to 'tasks'.
             owner_resolver: Function to resolve the owner from the context.
+            core_to_model_conversion: Optional function to convert a Task to a TaskModel.
+            model_to_core_conversion: Optional function to convert a TaskModel to a Task.
         """
         logger.debug(
             'Initializing DatabaseTaskStore with existing engine, table: %s',
@@ -87,6 +86,8 @@ def __init__(
         self.create_table = create_table
         self._initialized = False
         self.owner_resolver = owner_resolver
+        self.core_to_model_conversion = core_to_model_conversion
+        self.model_to_core_conversion = model_to_core_conversion
 
         self.task_model = (
             TaskModel
@@ -119,6 +120,9 @@ async def _ensure_initialized(self) -> None:
 
     def _to_orm(self, task: Task, owner: str) -> TaskModel:
         """Maps a Proto Task to a SQLAlchemy TaskModel instance."""
+        if self.core_to_model_conversion:
+            return self.core_to_model_conversion(task, owner)
+
         return self.task_model(
             id=task.id,
             context_id=task.context_id,
@@ -140,6 +144,9 @@ def _to_orm(self, task: Task, owner: str) -> TaskModel:
 
     def _from_orm(self, task_model: TaskModel) -> Task:
         """Maps a SQLAlchemy TaskModel to a Proto Task instance."""
+        if self.model_to_core_conversion:
+            return self.model_to_core_conversion(task_model)
+
         if task_model.protocol_version == '1.0':
             task = Task(
                 id=task_model.id,
@@ -160,29 +167,7 @@ def _from_orm(self, task_model: TaskModel) -> Task:
             return task
 
         # Legacy conversion
-        legacy_task = types_v03.Task(
-            id=task_model.id,
-            context_id=task_model.context_id,
-            status=types_v03.TaskStatus.model_validate(task_model.status),
-            artifacts=(
-                [
-                    types_v03.Artifact.model_validate(a)
-                    for a in task_model.artifacts
-                ]
-                if task_model.artifacts
-                else []
-            ),
-            history=(
-                [
-                    types_v03.Message.model_validate(m)
-                    for m in task_model.history
-                ]
-                if task_model.history
-                else []
-            ),
-            metadata=task_model.task_metadata or {},
-        )
-        return conversions.to_core_task(legacy_task)
+        return compat_task_model_to_core(task_model)
 
     async def save(
         self, task: Task, context: ServerCallContext | None = None