Merge branch 'main' into kthota/249

Author: holtskinner

.github/actions/spelling/allow.txt

@@ -38,6 +38,7 @@ drivername
 dunders
 euo
 excinfo
+fernet
 fetchrow
 fetchval
 genai

.github/workflows/unit-tests.yml

@@ -15,11 +15,24 @@ jobs:
       postgres:
         image: postgres:15-alpine
         env:
-          POSTGRES_USER: postgres
-          POSTGRES_PASSWORD: postgres
+          POSTGRES_USER: a2a
+          POSTGRES_PASSWORD: a2a_password
           POSTGRES_DB: a2a_test
         ports:
           - 5432:5432
+        options: >-
+          --health-cmd pg_isready --health-interval 10s --health-timeout 5s --health-retries 5
+      mysql:
+        image: mysql:8.0
+        env:
+          MYSQL_ROOT_PASSWORD: root
+          MYSQL_DATABASE: a2a_test
+          MYSQL_USER: a2a
+          MYSQL_PASSWORD: a2a_password
+        ports:
+          - 3306:3306
+        options: >-
+          --health-cmd="mysqladmin ping -h localhost -u root -proot" --health-interval=10s --health-timeout=5s --health-retries=5
 
     strategy:
       matrix:
@@ -31,19 +44,18 @@ 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 }}/tests/docker/postgres/init.sql
-          export POSTGRES_TEST_DSN="postgresql+asyncpg://postgres:postgres@localhost:5432/a2a_test"
+      - name: Set up test environment variables
+        run: |          
+          echo "POSTGRES_TEST_DSN=postgresql+asyncpg://a2a:a2a_password@localhost:5432/a2a_test" >> $GITHUB_ENV
+          echo "MYSQL_TEST_DSN=mysql+aiomysql://a2a:a2a_password@localhost:3306/a2a_test" >> $GITHUB_ENV
 
       - 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 --extra sql
+        run: uv sync --dev --extra sql --extra encryption
       - name: Run tests and check coverage
         run: uv run pytest --cov=a2a --cov-report=xml --cov-fail-under=89
       - name: Show coverage summary in log

pyproject.toml

@@ -41,6 +41,7 @@ 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"]
+encryption = ["cryptography>=43.0.0"]
 
 [project.urls]
 homepage = "https://a2aproject.github.io/A2A/"

src/a2a/server/models.py

@@ -16,7 +16,7 @@ def override(func):  # noqa: ANN001, ANN201
 
 
 try:
-    from sqlalchemy import JSON, Dialect, String
+    from sqlalchemy import JSON, Dialect, LargeBinary, String
     from sqlalchemy.orm import (
         DeclarativeBase,
         Mapped,
@@ -208,3 +208,58 @@ class TaskModel(TaskMixin, Base):
     """Default task model with standard table name."""
 
     __tablename__ = 'tasks'
+
+
+# PushNotificationConfigMixin that can be used with any table name
+class PushNotificationConfigMixin:
+    """Mixin providing standard push notification config columns."""
+
+    task_id: Mapped[str] = mapped_column(String(36), primary_key=True)
+    config_id: Mapped[str] = mapped_column(String(255), primary_key=True)
+    config_data: Mapped[bytes] = mapped_column(LargeBinary, nullable=False)
+
+    @override
+    def __repr__(self) -> str:
+        """Return a string representation of the push notification config."""
+        repr_template = '<{CLS}(task_id="{TID}", config_id="{CID}")>'
+        return repr_template.format(
+            CLS=self.__class__.__name__,
+            TID=self.task_id,
+            CID=self.config_id,
+        )
+
+
+def create_push_notification_config_model(
+    table_name: str = 'push_notification_configs',
+    base: type[DeclarativeBase] = Base,
+) -> type:
+    """Create a PushNotificationConfigModel class with a configurable table name."""
+
+    class PushNotificationConfigModel(PushNotificationConfigMixin, base):
+        __tablename__ = table_name
+
+        @override
+        def __repr__(self) -> str:
+            """Return a string representation of the push notification config."""
+            repr_template = '<PushNotificationConfigModel[{TABLE}](task_id="{TID}", config_id="{CID}")>'
+            return repr_template.format(
+                TABLE=table_name,
+                TID=self.task_id,
+                CID=self.config_id,
+            )
+
+    PushNotificationConfigModel.__name__ = (
+        f'PushNotificationConfigModel_{table_name}'
+    )
+    PushNotificationConfigModel.__qualname__ = (
+        f'PushNotificationConfigModel_{table_name}'
+    )
+
+    return PushNotificationConfigModel
+
+
+# Default PushNotificationConfigModel for backward compatibility
+class PushNotificationConfigModel(PushNotificationConfigMixin, Base):
+    """Default push notification config model with standard table name."""
+
+    __tablename__ = 'push_notification_configs'

src/a2a/server/tasks/__init__.py

@@ -43,8 +43,31 @@ def __init__(self, *args, **kwargs):
             ) from _original_error
 
 
+try:
+    from a2a.server.tasks.database_push_notification_config_store import (
+        DatabasePushNotificationConfigStore,  # type: ignore
+    )
+except ImportError as e:
+    _original_error = e
+    # If the database push notification config store is not available, we can still use in-memory stores.
+    logger.debug(
+        'DatabasePushNotificationConfigStore not loaded. This is expected if database dependencies are not installed. Error: %s',
+        e,
+    )
+
+    class DatabasePushNotificationConfigStore:  # type: ignore
+        """Placeholder for DatabasePushNotificationConfigStore when dependencies are not installed."""
+
+        def __init__(self, *args, **kwargs):
+            raise ImportError(
+                'To use DatabasePushNotificationConfigStore, its dependencies must be installed. '
+                'You can install them with \'pip install "a2a-sdk[sql]"\''
+            ) from _original_error
+
+
 __all__ = [
     'BasePushNotificationSender',
+    'DatabasePushNotificationConfigStore',
     'DatabaseTaskStore',
     'InMemoryPushNotificationConfigStore',
     'InMemoryTaskStore',

src/a2a/server/tasks/base_push_notification_sender.py

@@ -58,7 +58,7 @@ async def _dispatch_notification(
             response = await self._client.post(
                 url,
                 json=task.model_dump(mode='json', exclude_none=True),
-                headers=headers
+                headers=headers,
             )
             response.raise_for_status()
             logger.info(