diff --git a/src/a2a/migrations/README.md b/src/a2a/migrations/README.md index 00b99f6fb..4767363ff 100644 --- a/src/a2a/migrations/README.md +++ b/src/a2a/migrations/README.md @@ -38,6 +38,8 @@ Or you can use the `--database-url` flag to specify the database URL for a singl ### 3. Apply Migrations Always run this command after installing or upgrading the SDK to ensure your database matches the required schema. This will upgrade the tables `tasks` and `push_notification_configs` in your database by adding columns `owner` and `last_updated` and an index `(owner, last_updated)` to the `tasks` table and a column `owner` to the `push_notification_configs` table. +Revision `b5e3d1c8a2f7` creates two additive tables, `task_versions` and `task_events`, used only by the clustered multi-server stores (`a2a.server.cluster`). It leaves the `tasks` table untouched, so deployments that use the default `DatabaseTaskStore` need not run it and are unaffected. Run it only when adopting a `VersionedTaskStore` for multi-server deployment. + ```bash uv run a2a-db ``` diff --git a/src/a2a/migrations/versions/b5e3d1c8a2f7_add_task_version_and_task_events.py b/src/a2a/migrations/versions/b5e3d1c8a2f7_add_task_version_and_task_events.py new file mode 100644 index 000000000..9d03ad033 --- /dev/null +++ b/src/a2a/migrations/versions/b5e3d1c8a2f7_add_task_version_and_task_events.py @@ -0,0 +1,100 @@ +"""add task_versions and task_events tables + +Revision ID: b5e3d1c8a2f7 +Revises: 38ce57e08137 +Create Date: 2026-09-14 10:00:00.000000 + +""" + +import logging + +from collections.abc import Sequence +from typing import Union + +import sqlalchemy as sa + + +try: + from alembic import context, op +except ImportError as e: + raise ImportError( + "A2A migrations require the 'db-cli' extra. Install with: 'pip install a2a-sdk[db-cli]'." + ) from e + +from a2a.migrations.migration_utils import table_exists + + +# revision identifiers, used by Alembic. +revision: str = 'b5e3d1c8a2f7' +down_revision: Union[str, Sequence[str], None] = '38ce57e08137' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def _seq_type() -> sa.types.TypeEngine: + """Primary-key type for the event log. + + BigInteger everywhere except SQLite, whose AUTOINCREMENT requires a plain + INTEGER column. Mirrors ``TaskEventMixin.seq`` in a2a.server.models. + """ + return sa.BigInteger().with_variant(sa.Integer(), 'sqlite') + + +def upgrade() -> None: + """Upgrade schema: create the task_versions and task_events tables.""" + versions_table = context.config.get_main_option( + 'task_versions_table', 'task_versions' + ) + events_table = context.config.get_main_option( + 'task_events_table', 'task_events' + ) + + if context.is_offline_mode() or not table_exists(versions_table): + op.create_table( + versions_table, + sa.Column('task_id', sa.String(36), primary_key=True), + sa.Column('owner', sa.String(255), nullable=True), + sa.Column('version', sa.BigInteger(), nullable=False), + ) + else: + logging.info( + "Table '%s' already exists. Skipping creation.", versions_table + ) + + if context.is_offline_mode() or not table_exists(events_table): + op.create_table( + events_table, + sa.Column('seq', _seq_type(), primary_key=True, autoincrement=True), + sa.Column('task_id', sa.String(36), nullable=False), + sa.Column('owner', sa.String(255), nullable=True), + sa.Column('task_version', sa.BigInteger(), nullable=False), + sa.Column('event_data', sa.LargeBinary(), nullable=False), + ) + op.create_index(f'ix_{events_table}_task_id', events_table, ['task_id']) + else: + logging.info( + "Table '%s' already exists. Skipping creation.", events_table + ) + + +def downgrade() -> None: + """Downgrade schema: drop the task_events and task_versions tables.""" + versions_table = context.config.get_main_option( + 'task_versions_table', 'task_versions' + ) + events_table = context.config.get_main_option( + 'task_events_table', 'task_events' + ) + + if context.is_offline_mode() or table_exists(events_table): + op.drop_index(f'ix_{events_table}_task_id', table_name=events_table) + op.drop_table(events_table) + else: + logging.info("Table '%s' does not exist. Skipping drop.", events_table) + + if context.is_offline_mode() or table_exists(versions_table): + op.drop_table(versions_table) + else: + logging.info( + "Table '%s' does not exist. Skipping drop.", versions_table + ) diff --git a/tests/migrations/versions/test_migration_b5e3d1c8a2f7.py b/tests/migrations/versions/test_migration_b5e3d1c8a2f7.py new file mode 100644 index 000000000..ad9d4aaf6 --- /dev/null +++ b/tests/migrations/versions/test_migration_b5e3d1c8a2f7.py @@ -0,0 +1,175 @@ +import importlib +import os +import sqlite3 +import tempfile + +from typing import Generator +from unittest.mock import patch + +import pytest + +from a2a.a2a_db_cli import run_migrations + + +# Explicitly import the migration module so it is tracked when Alembic loads it +# dynamically. Revision id starts with a letter, so this import is valid. +try: + importlib.import_module( + 'a2a.migrations.versions.b5e3d1c8a2f7_add_task_version_and_task_events' + ) +except (ImportError, AttributeError): + pass + + +REVISION = 'b5e3d1c8a2f7' +PREV_REVISION = '38ce57e08137' + + +@pytest.fixture(autouse=True) +def mock_logging_config(): + """Prevent tests from mutating global logging state.""" + with patch('logging.basicConfig'), patch('logging.config.fileConfig'): + yield + + +@pytest.fixture +def temp_db() -> Generator[str, None, None]: + """Create a temporary SQLite database for testing.""" + fd, path = tempfile.mkstemp(suffix='.db') + os.close(fd) + yield path + if os.path.exists(path): + os.remove(path) + + +def _setup_initial_schema(db_path: str) -> None: + """Create the base tasks/push tables the earlier migrations expect.""" + conn = sqlite3.connect(db_path) + cursor = conn.cursor() + cursor.execute(""" + CREATE TABLE tasks ( + id VARCHAR(36) PRIMARY KEY, + context_id VARCHAR(36) NOT NULL, + kind VARCHAR(16) NOT NULL, + status TEXT, + artifacts TEXT, + history TEXT, + metadata TEXT + ) + """) + cursor.execute(""" + CREATE TABLE push_notification_configs ( + task_id VARCHAR(36), + config_id VARCHAR(255), + config_data BLOB NOT NULL, + PRIMARY KEY (task_id, config_id) + ) + """) + conn.commit() + conn.close() + + +def _upgrade(db_url: str, revision: str = REVISION) -> None: + with patch( + 'sys.argv', + ['a2a-db', '--database-url', db_url, 'upgrade', revision], + ): + run_migrations() + + +def test_migration_b5e3d1c8a2f7_full_cycle(temp_db: str) -> None: + """Upgrade creates task_versions + task_events; downgrade reverses it.""" + db_url = f'sqlite+aiosqlite:///{temp_db}' + _setup_initial_schema(temp_db) + + # Upgrade through the whole chain up to this revision. + _upgrade(db_url) + + conn = sqlite3.connect(temp_db) + cursor = conn.cursor() + + # tasks is untouched: no version column. + cursor.execute('PRAGMA table_info(tasks)') + tasks_columns = {row[1] for row in cursor.fetchall()} + assert 'version' not in tasks_columns + + # task_versions table exists with the expected columns. + cursor.execute('PRAGMA table_info(task_versions)') + version_columns = {row[1] for row in cursor.fetchall()} + assert version_columns == {'task_id', 'owner', 'version'} + + # task_events table exists with the expected columns. + cursor.execute('PRAGMA table_info(task_events)') + event_columns = {row[1] for row in cursor.fetchall()} + assert event_columns == { + 'seq', + 'task_id', + 'owner', + 'task_version', + 'event_data', + } + + # Index on task_id exists. + cursor.execute('PRAGMA index_list(task_events)') + event_indexes = {row[1] for row in cursor.fetchall()} + assert 'ix_task_events_task_id' in event_indexes + conn.close() + + # Downgrade one step: both tables are gone. + with patch( + 'sys.argv', + ['a2a-db', '--database-url', db_url, 'downgrade', PREV_REVISION], + ): + run_migrations() + + conn = sqlite3.connect(temp_db) + cursor = conn.cursor() + cursor.execute("SELECT name FROM sqlite_schema WHERE type='table'") + tables = {row[0] for row in cursor.fetchall()} + assert 'task_events' not in tables + assert 'task_versions' not in tables + conn.close() + + +def test_migration_b5e3d1c8a2f7_idempotency(temp_db: str) -> None: + """Running the upgrade twice must not fail.""" + db_url = f'sqlite+aiosqlite:///{temp_db}' + _setup_initial_schema(temp_db) + _upgrade(db_url) + # Second run: both tables already exist. + _upgrade(db_url) + + conn = sqlite3.connect(temp_db) + cursor = conn.cursor() + cursor.execute("SELECT name FROM sqlite_schema WHERE type='table'") + tables = {row[0] for row in cursor.fetchall()} + assert {'task_versions', 'task_events'} <= tables + conn.close() + + +def test_migration_b5e3d1c8a2f7_offline( + temp_db: str, capsys: pytest.CaptureFixture[str] +) -> None: + """Offline (--sql) mode emits the DDL without touching the database.""" + db_url = f'sqlite+aiosqlite:///{temp_db}' + _setup_initial_schema(temp_db) + + with patch( + 'sys.argv', + ['a2a-db', '--database-url', db_url, '--sql', 'upgrade', REVISION], + ): + run_migrations() + + out = capsys.readouterr().out + assert 'CREATE TABLE task_versions' in out + assert 'CREATE TABLE task_events' in out + assert 'ix_task_events_task_id' in out + + # The database itself was not modified. + conn = sqlite3.connect(temp_db) + cursor = conn.cursor() + cursor.execute("SELECT name FROM sqlite_schema WHERE type='table'") + tables = {row[0] for row in cursor.fetchall()} + assert 'task_events' not in tables + assert 'task_versions' not in tables + conn.close()