feat: Support for database backend Task Store (#259)

```
Release-As: 0.2.11
```

Add support for database backed TaskStore

### Installation

```bash
uv add a2a-sdk[postgresql] #postgres driver
or 
uv add a2a-sdk[mysql] #mysql driver
or
uv add a2a-sdk[sqlite] #sqlite driver
or
uv add a2a-sdk[sql]. #install all three sql drivers
```

### Usage

```py
from a2a.server.tasks import DatabaseTaskStore
from sqlalchemy.ext.asyncio import (
        create_async_engine,
    )
...
...
# Postgres - "postgresql+asyncpg://postgres:postgres@localhost:5432/a2a_test"
# Sqlite - "sqlite+aiosqlite:///file::memory:?cache=shared"
# Mysql - "mysql+aiomysql://admin:saysoyeah@localhost:3306/testDB"
engine = create_async_engine(
    "postgresql+asyncpg://postgres:postgres@localhost:5432/a2a_test", echo=False
)
request_handler = DefaultRequestHandler(
    agent_executor=CurrencyAgentExecutor(),
    task_store=DatabaseTaskStore(engine=engine),
    push_notifier=InMemoryPushNotifier(httpx_client),
)
```

---------

Co-authored-by: Zac <[email protected]>
Co-authored-by: MEUNIER Laurent <[email protected]>
Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
Co-authored-by: Zac <[email protected]>
Co-authored-by: Holt Skinner <[email protected]>
Co-authored-by: Holt Skinner <[email protected]>

Author: 5 people

.github/actions/spelling/allow.txt

@@ -9,17 +9,22 @@ AServer
 AServers
 AService
 AStarlette
+AUser
+DSNs
 EUR
 GBP
+GVsb
 INR
 JPY
 JSONRPCt
 Llm
+POSTGRES
 RUF
 aconnect
 adk
 agentic
 aio
+aiomysql
 aproject
 autouse
 backticks
@@ -29,23 +34,34 @@ coc
 codegen
 coro
 datamodel
+drivername
 dunders
 euo
+excinfo
+fetchrow
+fetchval
 genai
 getkwargs
 gle
+initdb
 inmemory
+isready
 kwarg
 langgraph
 lifecycles
 linting
 lstrips
 mockurl
+notif
 oauthoidc
 oidc
 opensource
+otherurl
+postgres
+postgresql
 protoc
 pyi
+pypistats
 pyversions
 respx
 resub

.github/actions/spelling/expect.txt

@@ -9,7 +9,18 @@ jobs:
   test:
     name: Test with Python ${{ matrix.python-version }}
     runs-on: ubuntu-latest
+
     if: github.repository == 'a2aproject/a2a-python'
+    services:
+      postgres:
+        image: postgres:15-alpine
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: a2a_test
+        ports:
+          - 5432:5432
+
     strategy:
       matrix:
         python-version: ['3.10', '3.13']
@@ -20,13 +31,19 @@ jobs:
         uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
+      - name: Set postgres for tests
+        run: |
+          sudo apt-get update && sudo apt-get install -y postgresql-client
+          PGPASSWORD=postgres psql -h localhost -p 5432 -U postgres -d a2a_test -f ${{ github.workspace }}/docker/postgres/init.sql
+          export POSTGRES_TEST_DSN="postgresql+asyncpg://postgres:postgres@localhost:5432/a2a_test"
+
       - name: Install uv
         uses: astral-sh/setup-uv@v6
       - name: Add uv to PATH
         run: |
           echo "$HOME/.cargo/bin" >> $GITHUB_PATH
       - name: Install dependencies
-        run: uv sync --dev
+        run: uv sync --dev --extra sql
       - name: Run tests and check coverage
         run: uv run pytest --cov=a2a --cov-report=xml --cov-fail-under=90
       - name: Show coverage summary in log

.github/workflows/unit-tests.yml

@@ -33,6 +33,22 @@ When you're working within a uv project or a virtual environment managed by uv,
 uv add a2a-sdk
 ```
 
+To install with database support:
+
+```bash
+# PostgreSQL support
+uv add "a2a-sdk[postgresql]"
+
+# MySQL support
+uv add "a2a-sdk[mysql]"
+
+# SQLite support
+uv add "a2a-sdk[sqlite]"
+
+# All database drivers
+uv add "a2a-sdk[sql]"
+```
+
 ### Using `pip`
 
 If you prefer to use pip, the standard Python package installer, you can install `a2a-sdk` as follows
@@ -41,6 +57,22 @@ If you prefer to use pip, the standard Python package installer, you can install
 pip install a2a-sdk
 ```
 
+To install with database support:
+
+```bash
+# PostgreSQL support
+pip install "a2a-sdk[postgresql]"
+
+# MySQL support
+pip install "a2a-sdk[mysql]"
+
+# SQLite support
+pip install "a2a-sdk[sqlite]"
+
+# All database drivers
+pip install "a2a-sdk[sql]"
+```
+
 ## Examples
 
 ### [Helloworld Example](https://github.com/a2aproject/a2a-samples/tree/main/samples/python/agents/helloworld)
@@ -60,7 +92,7 @@ pip install a2a-sdk
    uv run test_client.py
    ```
 
-3. You can validate your agent using the agent inspector. Follow the instructions at the [a2a-inspector](https://github.com/a2aproject/a2a-inspector) repo. 
+3. You can validate your agent using the agent inspector. Follow the instructions at the [a2a-inspector](https://github.com/a2aproject/a2a-inspector) repo.
 
 You can also find more Python samples [here](https://github.com/a2aproject/a2a-samples/tree/main/samples/python) and JavaScript samples [here](https://github.com/a2aproject/a2a-samples/tree/main/samples/js).
 

README.md

@@ -0,0 +1,8 @@
+-- Create a dedicated user for the application
+CREATE USER a2a WITH PASSWORD 'a2a_password';
+
+-- Create the tasks database
+CREATE DATABASE a2a_tasks;
+
+GRANT ALL PRIVILEGES ON DATABASE a2a_test TO a2a;
+

docker/postgres/init.sql

@@ -36,6 +36,12 @@ classifiers = [
   "License :: OSI Approved :: Apache Software License",
 ]
 
+[project.optional-dependencies]
+postgresql = ["sqlalchemy[asyncio,postgresql-asyncpg]>=2.0.0"]
+mysql = ["sqlalchemy[asyncio,aiomysql]>=2.0.0"]
+sqlite = ["sqlalchemy[asyncio,aiosqlite]>=2.0.0"]
+sql = ["sqlalchemy[asyncio,postgresql-asyncpg,aiomysql,aiosqlite]>=2.0.0"]
+
 [project.urls]
 homepage = "https://a2aproject.github.io/A2A/"
 repository = "https://github.com/a2aproject/a2a-python"

pyproject.toml

@@ -0,0 +1,210 @@
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+
+if TYPE_CHECKING:
+    from typing_extensions import override
+else:
+
+    def override(func):  # noqa: ANN001, ANN201
+        """Override decorator."""
+        return func
+
+
+from pydantic import BaseModel
+
+from a2a.types import Artifact, Message, TaskStatus
+
+
+try:
+    from sqlalchemy import JSON, Dialect, String
+    from sqlalchemy.orm import (
+        DeclarativeBase,
+        Mapped,
+        declared_attr,
+        mapped_column,
+    )
+    from sqlalchemy.types import TypeDecorator
+except ImportError as e:
+    raise ImportError(
+        'Database models require SQLAlchemy. '
+        'Install with one of: '
+        "'pip install a2a-sdk[postgresql]', "
+        "'pip install a2a-sdk[mysql]', "
+        "'pip install a2a-sdk[sqlite]', "
+        "or 'pip install a2a-sdk[sql]'"
+    ) from e
+
+
+T = TypeVar('T', bound=BaseModel)
+
+
+class PydanticType(TypeDecorator[T], Generic[T]):
+    """SQLAlchemy type that handles Pydantic model serialization."""
+
+    impl = JSON
+    cache_ok = True
+
+    def __init__(self, pydantic_type: type[T], **kwargs: dict[str, Any]):
+        """Initialize the PydanticType.
+
+        Args:
+            pydantic_type: The Pydantic model type to handle.
+            **kwargs: Additional arguments for TypeDecorator.
+        """
+        self.pydantic_type = pydantic_type
+        super().__init__(**kwargs)
+
+    def process_bind_param(
+        self, value: T | None, dialect: Dialect
+    ) -> dict[str, Any] | None:
+        """Convert Pydantic model to a JSON-serializable dictionary for the database."""
+        if value is None:
+            return None
+        return (
+            value.model_dump(mode='json')
+            if isinstance(value, BaseModel)
+            else value
+        )
+
+    def process_result_value(
+        self, value: dict[str, Any] | None, dialect: Dialect
+    ) -> T | None:
+        """Convert a JSON-like dictionary from the database back to a Pydantic model."""
+        if value is None:
+            return None
+        return self.pydantic_type.model_validate(value)
+
+
+class PydanticListType(TypeDecorator, Generic[T]):
+    """SQLAlchemy type that handles lists of Pydantic models."""
+
+    impl = JSON
+    cache_ok = True
+
+    def __init__(self, pydantic_type: type[T], **kwargs: dict[str, Any]):
+        """Initialize the PydanticListType.
+
+        Args:
+            pydantic_type: The Pydantic model type for items in the list.
+            **kwargs: Additional arguments for TypeDecorator.
+        """
+        self.pydantic_type = pydantic_type
+        super().__init__(**kwargs)
+
+    def process_bind_param(
+        self, value: list[T] | None, dialect: Dialect
+    ) -> list[dict[str, Any]] | None:
+        """Convert a list of Pydantic models to a JSON-serializable list for the DB."""
+        if value is None:
+            return None
+        return [
+            item.model_dump(mode='json')
+            if isinstance(item, BaseModel)
+            else item
+            for item in value
+        ]
+
+    def process_result_value(
+        self, value: list[dict[str, Any]] | None, dialect: Dialect
+    ) -> list[T] | None:
+        """Convert a JSON-like list from the DB back to a list of Pydantic models."""
+        if value is None:
+            return None
+        return [self.pydantic_type.model_validate(item) for item in value]
+
+
+# Base class for all database models
+class Base(DeclarativeBase):
+    """Base class for declarative models in A2A SDK."""
+
+
+# TaskMixin that can be used with any table name
+class TaskMixin:
+    """Mixin providing standard task columns with proper type handling."""
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, index=True)
+    contextId: Mapped[str] = mapped_column(String(36), nullable=False)  # noqa: N815
+    kind: Mapped[str] = mapped_column(
+        String(16), nullable=False, default='task'
+    )
+
+    # Properly typed Pydantic fields with automatic serialization
+    status: Mapped[TaskStatus] = mapped_column(PydanticType(TaskStatus))
+    artifacts: Mapped[list[Artifact] | None] = mapped_column(
+        PydanticListType(Artifact), nullable=True
+    )
+    history: Mapped[list[Message] | None] = mapped_column(
+        PydanticListType(Message), nullable=True
+    )
+
+    # Using declared_attr to avoid conflict with Pydantic's metadata
+    @declared_attr
+    @classmethod
+    def task_metadata(cls) -> Mapped[dict[str, Any] | None]:
+        """Define the 'metadata' column, avoiding name conflicts with Pydantic."""
+        return mapped_column(JSON, nullable=True, name='metadata')
+
+    @override
+    def __repr__(self) -> str:
+        """Return a string representation of the task."""
+        repr_template = (
+            '<{CLS}(id="{ID}", contextId="{CTX_ID}", status="{STATUS}")>'
+        )
+        return repr_template.format(
+            CLS=self.__class__.__name__,
+            ID=self.id,
+            CTX_ID=self.contextId,
+            STATUS=self.status,
+        )
+
+
+def create_task_model(
+    table_name: str = 'tasks', base: type[DeclarativeBase] = Base
+) -> type:
+    """Create a TaskModel class with a configurable table name.
+
+    Args:
+        table_name: Name of the database table. Defaults to 'tasks'.
+        base: Base declarative class to use. Defaults to the SDK's Base class.
+
+    Returns:
+        TaskModel class with the specified table name.
+
+    Example:
+        # Create a task model with default table name
+        TaskModel = create_task_model()
+
+        # Create a task model with custom table name
+        CustomTaskModel = create_task_model('my_tasks')
+
+        # Use with a custom base
+        from myapp.database import Base as MyBase
+        TaskModel = create_task_model('tasks', MyBase)
+    """
+
+    class TaskModel(TaskMixin, base):
+        __tablename__ = table_name
+
+        @override
+        def __repr__(self) -> str:
+            """Return a string representation of the task."""
+            repr_template = '<TaskModel[{TABLE}](id="{ID}", contextId="{CTX_ID}", status="{STATUS}")>'
+            return repr_template.format(
+                TABLE=table_name,
+                ID=self.id,
+                CTX_ID=self.contextId,
+                STATUS=self.status,
+            )
+
+    # Set a dynamic name for better debugging
+    TaskModel.__name__ = f'TaskModel_{table_name}'
+    TaskModel.__qualname__ = f'TaskModel_{table_name}'
+
+    return TaskModel
+
+
+# Default TaskModel for backward compatibility
+class TaskModel(TaskMixin, Base):
+    """Default task model with standard table name."""
+
+    __tablename__ = 'tasks'