feat(db): support multiple SQL backends with optional extras (cont. #80)

.github/actions/spelling/expect.txt

@@ -1,5 +1,15 @@
 AUser
 excinfo
+fetchrow
+fetchval
 GVsb
+initdb
+isready
 notif
 otherurl
+POSTGRES
+postgres
+postgresql
+drivername
+aiomysql
+DSNs

.github/workflows/spelling.yaml

@@ -80,6 +80,7 @@ jobs:
             cspell:sql/src/tsql.txt
             cspell:terraform/dict/terraform.txt
             cspell:typescript/dict/typescript.txt
+
           check_extra_dictionaries: ""
           only_check_changed_files: true
           longest_word: "10"

.github/workflows/unit-tests.yml

@@ -15,6 +15,15 @@ jobs:
     runs-on: ubuntu-latest
 
     if: github.repository == 'google-a2a/a2a-python'
+    services:
+      postgres:
+        image: postgres:15-alpine
+        env:
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: a2a_test
+        ports:
+          - 5432:5432
 
     strategy:
       matrix:
@@ -28,6 +37,11 @@ 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://postgres:postgres@localhost:5432/a2a_test"
 
       - name: Install uv
         run: |
@@ -38,7 +52,7 @@ jobs:
           echo "$HOME/.cargo/bin" >> $GITHUB_PATH
 
       - name: Install dependencies
-        run: uv sync --dev
+        run: uv sync --dev --extra sql
 
       - name: Run tests
         run: uv run pytest

README.md

@@ -32,6 +32,21 @@ 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
@@ -40,6 +55,21 @@ 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/google-a2a/a2a-samples/tree/main/samples/python/agents/helloworld)

docker/postgres/docker-compose.yml

@@ -0,0 +1,28 @@
+version: "3.8"
+
+services:
+  postgres:
+    image: postgres:15-alpine
+    ports:
+      - "5432:5432"
+    environment:
+      - POSTGRES_USER=postgres
+      - POSTGRES_PASSWORD=postgres
+      - POSTGRES_DB=a2a_test
+    volumes:
+      - postgres_data:/var/lib/postgresql/data
+      - ./docker/postgres/init.sql:/docker-entrypoint-initdb.d/init.sql
+    networks:
+      - a2a-network
+    healthcheck:
+      test: ["CMD-SHELL", "pg_isready -U postgres"]
+      interval: 5s
+      timeout: 5s
+      retries: 5
+
+volumes:
+  postgres_data:
+
+networks:
+  a2a-network:
+    driver: bridge

docker/postgres/init.sql

@@ -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;
+

pyproject.toml

@@ -30,6 +30,26 @@ classifiers = [
   "License :: OSI Approved :: Apache Software License",
 ]
 
+[project.optional-dependencies]
+postgresql = [
+    "sqlalchemy>=2.0.0",
+    "asyncpg>=0.30.0",
+]
+mysql = [
+    "sqlalchemy>=2.0.0",
+    "aiomysql>=0.2.0",
+]
+sqlite = [
+    "sqlalchemy>=2.0.0",
+    "aiosqlite>=0.19.0",
+]
+sql = [
+    "sqlalchemy>=2.0.0",
+    "asyncpg>=0.30.0",
+    "aiomysql>=0.2.0",
+    "aiosqlite>=0.19.0",
+]
+
 [project.urls]
 homepage = "https://google.github.io/A2A/"
 repository = "https://github.com/google-a2a/a2a-python"
@@ -64,8 +84,11 @@ style = "pep440"
 
 [dependency-groups]
 dev = [
+    "asyncpg-stubs>=0.30.1",
     "datamodel-code-generator>=0.30.0",
+    "greenlet>=3.2.2",
     "mypy>=1.15.0",
+    "nox>=2025.5.1",
     "pytest>=8.3.5",
     "pytest-asyncio>=0.26.0",
     "pytest-cov>=6.1.1",

src/a2a/server/models.py

@@ -0,0 +1,194 @@
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+
+if TYPE_CHECKING:
+    from typing_extensions import override
+else:
+
+    def override(func):  # noqa: D103
+        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]):
+        self.pydantic_type = pydantic_type
+        super().__init__(**kwargs)
+
+    @override
+    def process_bind_param(
+        self, value: T | None, dialect: Dialect
+    ) -> dict[str, Any] | None:
+        if value is None:
+            return None
+        return (
+            value.model_dump(mode='json')
+            if isinstance(value, BaseModel)
+            else value
+        )
+
+    @override
+    def process_result_value(
+        self, value: dict[str, Any] | None, dialect: Dialect
+    ) -> T | None:
+        if value is None:
+            return None
+        return self.pydantic_type.model_validate(value)
+
+
+class PydanticListType(TypeDecorator[list[T]], 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]):
+        self.pydantic_type = pydantic_type
+        super().__init__(**kwargs)
+
+    @override
+    def process_bind_param(
+        self, value: list[T] | None, dialect: Dialect
+    ) -> list[dict[str, Any]] | None:
+        if value is None:
+            return None
+        return [
+            item.model_dump(mode='json')
+            if isinstance(item, BaseModel)
+            else item
+            for item in value
+        ]
+
+    @override
+    def process_result_value(
+        self, value: list[dict[str, Any]] | None, dialect: Dialect
+    ) -> list[T] | None:
+        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, primary_key=True, index=True)
+    contextId: Mapped[str] = mapped_column(String, nullable=False)  # noqa: N815
+    kind: Mapped[str] = mapped_column(String, 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]:
+        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'

src/a2a/server/tasks/__init__.py

@@ -1,5 +1,6 @@
 """Components for managing tasks within the A2A server."""
 
+from a2a.server.tasks.database_task_store import DatabaseTaskStore
 from a2a.server.tasks.inmemory_push_notifier import InMemoryPushNotifier
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
 from a2a.server.tasks.push_notifier import PushNotifier
@@ -10,6 +11,7 @@
 
 
 __all__ = [
+    'DatabaseTaskStore',
     'InMemoryPushNotifier',
     'InMemoryTaskStore',
     'PushNotifier',