Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/actions/spelling/allow.txt
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ drivername
dunders
euo
excinfo
fernet
fetchrow
fetchval
genai
Expand Down
28 changes: 20 additions & 8 deletions .github/workflows/unit-tests.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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/"
Expand Down
57 changes: 56 additions & 1 deletion src/a2a/server/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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'
23 changes: 23 additions & 0 deletions src/a2a/server/tasks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
2 changes: 1 addition & 1 deletion src/a2a/server/tasks/base_push_notification_sender.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
Loading
Loading