feat: Zero-Downtown migration support for DataBases

sokoliva · sokoliva · commit 123deeab8217 · 2026-03-13T17:32:42.000Z
diff --git a/src/a2a/compat/v0_3/conversions.py b/src/a2a/compat/v0_3/conversions.py
@@ -5,7 +5,11 @@
 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
+from cryptography.fernet import (
+    Fernet,
+)
 
 
 _COMPAT_TO_CORE_TASK_STATE: dict[types_v03.TaskState, Any] = {
@@ -1367,3 +1371,79 @@ 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)
+        if task_model.status
+        else None,
+        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,
+        )
+    )
diff --git a/src/a2a/server/tasks/database_push_notification_config_store.py b/src/a2a/server/tasks/database_push_notification_config_store.py
@@ -26,8 +26,12 @@
         "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
+if TYPE_CHECKING:
+    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,
@@ -40,11 +44,8 @@
 )
 from a2a.types.a2a_pb2 import TaskPushNotificationConfig
 
-
 if TYPE_CHECKING:
     from cryptography.fernet import Fernet
-
-
 logger = logging.getLogger(__name__)
 
 
@@ -62,6 +63,9 @@ class DatabasePushNotificationConfigStore(PushNotificationConfigStore):
     _fernet: 'Fernet | None'
     owner_resolver: OwnerResolver
 
+    core_to_model_conversion: 'Callable[[str, TaskPushNotificationConfig, str, Fernet | None], PushNotificationConfigModel] | None' = None
+    model_to_core_conversion: 'Callable[[PushNotificationConfigModel], TaskPushNotificationConfig] | None' = None
+
     def __init__(
         self,
         engine: AsyncEngine,
@@ -152,6 +156,14 @@ def _to_orm(
 
         The config data is serialized to JSON bytes, and encrypted if a key is configured.
         """
+        if self.core_to_model_conversion:
+            conversion = self.core_to_model_conversion
+            # bound method
+            if hasattr(conversion, '__func__'):
+                return conversion.__func__(task_id, config, owner, self._fernet)
+            # instance method
+            return conversion(task_id, config, owner, self._fernet)
+
         json_payload = MessageToJson(config).encode('utf-8')
 
         if self._fernet:
@@ -174,6 +186,14 @@ def _from_orm(
 
         Handles decryption if a key is configured, with a fallback to plain JSON.
         """
+        if self.model_to_core_conversion:
+            conversion = self.model_to_core_conversion
+            # bound method
+            if hasattr(conversion, '__func__'):
+                return conversion.__func__(model_instance)
+            # instance method
+            return conversion(model_instance)
+
         payload = model_instance.config_data
 
         if self._fernet:
@@ -359,12 +379,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 ''
         )
diff --git a/src/a2a/server/tasks/database_task_store.py b/src/a2a/server/tasks/database_task_store.py
@@ -1,6 +1,7 @@
 import logging
 
 from datetime import datetime, timezone
+from typing import TYPE_CHECKING
 
 
 try:
@@ -30,10 +31,12 @@
         "or 'pip install a2a-sdk[sql]'"
     ) from e
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
 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.server.context import ServerCallContext
 from a2a.server.models import Base, TaskModel, create_task_model
 from a2a.server.owner_resolver import OwnerResolver, resolve_user_scope
@@ -61,6 +64,9 @@ class DatabaseTaskStore(TaskStore):
     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__(
         self,
         engine: AsyncEngine,
@@ -119,6 +125,14 @@ 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:
+            conversion = self.core_to_model_conversion
+            # bound method
+            if hasattr(conversion, '__func__'):
+                return conversion.__func__(task, owner)
+            # instance method
+            return conversion(task, owner)
+
         return self.task_model(
             id=task.id,
             context_id=task.context_id,
@@ -140,6 +154,14 @@ 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:
+            conversion = self.model_to_core_conversion
+            # bound method
+            if hasattr(conversion, '__func__'):
+                return conversion.__func__(task_model)
+            # instance method
+            return conversion(task_model)
+
         if task_model.protocol_version == '1.0':
             task = Task(
                 id=task_model.id,
@@ -160,29 +182,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 conversions.compat_task_model_to_core(task_model)
 
     async def save(
         self, task: Task, context: ServerCallContext | None = None
diff --git a/tests/compat/v0_3/test_conversions.py b/tests/compat/v0_3/test_conversions.py
@@ -72,7 +72,13 @@
     to_core_task_push_notification_config,
     to_core_task_status,
     to_core_task_status_update_event,
+    core_to_compat_task_model,
+    compat_task_model_to_core,
+    core_to_compat_push_notification_config_model,
+    compat_push_notification_config_model_to_core,
 )
+from a2a.server.models import PushNotificationConfigModel, TaskModel
+from cryptography.fernet import Fernet
 from a2a.types import a2a_pb2 as pb2_v10
 
 
@@ -1911,3 +1917,100 @@ def test_to_core_part_unknown_part():
     assert not core_part.HasField('data')
     assert not core_part.HasField('raw')
     assert not core_part.HasField('url')
+
+
+def test_task_db_conversion():
+    v10_task = pb2_v10.Task(
+        id='task-123',
+        context_id='ctx-456',
+        status=pb2_v10.TaskStatus(
+            state=pb2_v10.TaskState.TASK_STATE_WORKING,
+        ),
+        metadata={'m1': 'v1'},
+    )
+    owner = 'owner-789'
+
+    # Test Core -> Model
+    model = core_to_compat_task_model(v10_task, owner)
+    assert model.id == 'task-123'
+    assert model.context_id == 'ctx-456'
+    assert model.owner == owner
+    assert model.protocol_version == '0.3'
+    assert model.status['state'] == 'working'
+    assert model.task_metadata == {'m1': 'v1'}
+
+    # Test Model -> Core
+    v10_restored = compat_task_model_to_core(model)
+    assert v10_restored.id == v10_task.id
+    assert v10_restored.context_id == v10_task.context_id
+    assert v10_restored.status.state == v10_task.status.state
+    assert v10_restored.metadata == v10_task.metadata
+
+
+def test_push_notification_config_db_conversion():
+    task_id = 'task-123'
+    v10_config = pb2_v10.TaskPushNotificationConfig(
+        id='pnc-1',
+        url='https://example.com/push',
+        token='secret-token',
+    )
+    owner = 'owner-789'
+
+    # Test Core -> Model (No encryption)
+    model = core_to_compat_push_notification_config_model(
+        task_id, v10_config, owner
+    )
+    assert model.task_id == task_id
+    assert model.config_id == 'pnc-1'
+    assert model.owner == owner
+    assert model.protocol_version == '0.3'
+
+    import json
+
+    data = json.loads(model.config_data.decode('utf-8'))
+    assert data['url'] == 'https://example.com/push'
+    assert data['token'] == 'secret-token'
+
+    # Test Model -> Core
+    v10_restored = compat_push_notification_config_model_to_core(
+        model.config_data.decode('utf-8'), task_id
+    )
+    assert v10_restored.id == v10_config.id
+    assert v10_restored.url == v10_config.url
+    assert v10_restored.token == v10_config.token
+
+
+def test_push_notification_config_persistence_conversion_with_encryption():
+    task_id = 'task-123'
+    v10_config = pb2_v10.TaskPushNotificationConfig(
+        id='pnc-1',
+        url='https://example.com/push',
+        token='secret-token',
+    )
+    owner = 'owner-789'
+    key = Fernet.generate_key()
+    fernet = Fernet(key)
+
+    # Test Core -> Model (With encryption)
+    model = core_to_compat_push_notification_config_model(
+        task_id, v10_config, owner, fernet=fernet
+    )
+    assert (
+        model.config_data != v10_config.SerializeToString()
+    )  # Should be encrypted
+
+    # Decrypt and verify
+    decrypted_data = fernet.decrypt(model.config_data)
+    import json
+
+    data = json.loads(decrypted_data.decode('utf-8'))
+    assert data['url'] == 'https://example.com/push'
+    assert data['token'] == 'secret-token'
+
+    # Test Model -> Core
+    v10_restored = compat_push_notification_config_model_to_core(
+        decrypted_data.decode('utf-8'), task_id
+    )
+    assert v10_restored.id == v10_config.id
+    assert v10_restored.url == v10_config.url
+    assert v10_restored.token == v10_config.token
diff --git a/tests/server/tasks/test_database_push_notification_config_store.py b/tests/server/tasks/test_database_push_notification_config_store.py
diff --git a/tests/server/tasks/test_database_task_store.py b/tests/server/tasks/test_database_task_store.py
