Skip to content

Commit 2e2d431

Browse files
sokolivasokoliva
authored
feat: Add protocol_version column to Task and PushNotificationConfig models and create a migration (#789)
## Changes - Add `protocol_version` column to Task and PushNotificationConfig models - Add `add_column_protocol_version` migration - Refactor migration utilities ## Contributing Guide - [x] Follow the [`CONTRIBUTING` Guide](https://github.com/a2aproject/a2a-python/blob/main/CONTRIBUTING.md). - [x] Make your Pull Request title in the <https://www.conventionalcommits.org/> specification. - Important Prefixes for [release-please](https://github.com/googleapis/release-please): - `fix:` which represents bug fixes, and correlates to a [SemVer](https://semver.org/) patch. - `feat:` represents a new feature, and correlates to a SemVer minor. - `feat!:`, or `fix!:`, `refactor!:`, etc., which represent a breaking change (indicated by the `!`) and will result in a SemVer major. - [x] Ensure the tests and linter pass (Run `bash scripts/format.sh` from the repository root to format) - [x] Appropriate docs were updated (if necessary) Fixes #787 🦕
1 parent 13d0106 commit 2e2d431

5 files changed

Lines changed: 232 additions & 91 deletions

File tree

‎src/a2a/migrations/env.py‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,11 @@
3333

3434
# Interpret the config file for Python logging.
3535
# This line sets up loggers basically.
36-
if config.config_file_name is not None:
36+
if (
37+
config.config_file_name is not None
38+
and os.path.exists(config.config_file_name)
39+
and config.config_file_name.endswith('.ini')
40+
):
3741
fileConfig(config.config_file_name)
3842

3943
if config.get_main_option('verbose') == 'true':

‎src/a2a/migrations/migration_utils.py‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
"""Utility functions for Alembic migrations."""
2+
3+
import logging
4+
from typing import Any
5+
6+
import sqlalchemy as sa
7+
8+
try:
9+
from alembic import context, op
10+
except ImportError as e:
11+
raise ImportError(
12+
"A2A migrations require the 'db-cli' extra. Install with: 'pip install a2a-sdk[db-cli]'."
13+
) from e
14+
15+
16+
def _get_inspector() -> sa.engine.reflection.Inspector:
17+
"""Get the current database inspector."""
18+
bind = op.get_bind()
19+
inspector = sa.inspect(bind)
20+
return inspector
21+
22+
23+
def table_exists(table_name: str) -> bool:
24+
"""Check if a table exists in the database."""
25+
if context.is_offline_mode():
26+
return True
27+
inspector = _get_inspector()
28+
return table_name in inspector.get_table_names()
29+
30+
31+
def column_exists(
32+
table_name: str, column_name: str, downgrade_mode: bool = False
33+
) -> bool:
34+
"""Check if a column exists in a table."""
35+
if context.is_offline_mode():
36+
return downgrade_mode
37+
38+
inspector = _get_inspector()
39+
columns = [c['name'] for c in inspector.get_columns(table_name)]
40+
return column_name in columns
41+
42+
43+
def index_exists(
44+
table_name: str, index_name: str, downgrade_mode: bool = False
45+
) -> bool:
46+
"""Check if an index exists on a table."""
47+
if context.is_offline_mode():
48+
return downgrade_mode
49+
50+
inspector = _get_inspector()
51+
indexes = [i['name'] for i in inspector.get_indexes(table_name)]
52+
return index_name in indexes
53+
54+
55+
def add_column(
56+
table: str,
57+
column_name: str,
58+
nullable: bool,
59+
type_: sa.types.TypeEngine,
60+
default: Any | None = None,
61+
) -> None:
62+
"""Add a column to a table if it doesn't already exist."""
63+
if not column_exists(table, column_name):
64+
op.add_column(
65+
table,
66+
sa.Column(
67+
column_name,
68+
type_,
69+
nullable=nullable,
70+
server_default=default,
71+
),
72+
)
73+
else:
74+
logging.info(
75+
f"Column '{column_name}' already exists in table '{table}'. Skipping."
76+
)
77+
78+
79+
def drop_column(table: str, column_name: str) -> None:
80+
"""Drop a column from a table if it exists."""
81+
if column_exists(table, column_name, True):
82+
op.drop_column(table, column_name)
83+
else:
84+
logging.info(
85+
f"Column '{column_name}' does not exist in table '{table}'. Skipping."
86+
)
87+
88+
89+
def add_index(table: str, index_name: str, columns: list[str]) -> None:
90+
"""Create an index on a table if it doesn't already exist."""
91+
if not index_exists(table, index_name):
92+
op.create_index(
93+
index_name,
94+
table,
95+
columns,
96+
)
97+
else:
98+
logging.info(
99+
f"Index '{index_name}' already exists on table '{table}'. Skipping."
100+
)
101+
102+
103+
def drop_index(table: str, index_name: str) -> None:
104+
"""Drop an index from a table if it exists."""
105+
if index_exists(table, index_name, True):
106+
op.drop_index(index_name, table_name=table)
107+
else:
108+
logging.info(
109+
f"Index '{index_name}' does not exist on table '{table}'. Skipping."
110+
)

‎src/a2a/migrations/versions/38ce57e08137_add_column_protocol_version.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
"""add column protocol version
2+
3+
Revision ID: 38ce57e08137
4+
Revises: 6419d2d130f6
5+
Create Date: 2026-03-09 12:07:16.998955
6+
7+
"""
8+
9+
import logging
10+
from collections.abc import Sequence
11+
from typing import Union
12+
13+
import sqlalchemy as sa
14+
15+
try:
16+
from alembic import context
17+
except ImportError as e:
18+
raise ImportError(
19+
"A2A migrations require the 'db-cli' extra. Install with: 'pip install a2a-sdk[db-cli]'."
20+
) from e
21+
22+
from a2a.migrations.migration_utils import table_exists, add_column, drop_column
23+
24+
25+
# revision identifiers, used by Alembic.
26+
revision: str = '38ce57e08137'
27+
down_revision: Union[str, Sequence[str], None] = '6419d2d130f6'
28+
branch_labels: Union[str, Sequence[str], None] = None
29+
depends_on: Union[str, Sequence[str], None] = None
30+
31+
32+
def upgrade() -> None:
33+
"""Upgrade schema."""
34+
tasks_table = context.config.get_main_option('tasks_table', 'tasks')
35+
push_notification_configs_table = context.config.get_main_option(
36+
'push_notification_configs_table', 'push_notification_configs'
37+
)
38+
39+
if table_exists(tasks_table):
40+
add_column(tasks_table, 'protocol_version', True, sa.String(16))
41+
else:
42+
logging.warning(
43+
f"Table '{tasks_table}' does not exist. Skipping upgrade for this table."
44+
)
45+
46+
if table_exists(push_notification_configs_table):
47+
add_column(
48+
push_notification_configs_table,
49+
'protocol_version',
50+
True,
51+
sa.String(16),
52+
)
53+
else:
54+
logging.warning(
55+
f"Table '{push_notification_configs_table}' does not exist. Skipping upgrade for this table."
56+
)
57+
58+
59+
def downgrade() -> None:
60+
"""Downgrade schema."""
61+
tasks_table = context.config.get_main_option('tasks_table', 'tasks')
62+
push_notification_configs_table = context.config.get_main_option(
63+
'push_notification_configs_table', 'push_notification_configs'
64+
)
65+
66+
if table_exists(tasks_table):
67+
drop_column(tasks_table, 'protocol_version')
68+
else:
69+
logging.warning(
70+
f"Table '{tasks_table}' does not exist. Skipping downgrade for this table."
71+
)
72+
73+
if table_exists(push_notification_configs_table):
74+
drop_column(push_notification_configs_table, 'protocol_version')
75+
else:
76+
logging.warning(
77+
f"Table '{push_notification_configs_table}' does not exist. Skipping downgrade for this table."
78+
)

‎src/a2a/migrations/versions/6419d2d130f6_add_columns_owner_last_updated.py‎

Lines changed: 32 additions & 89 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,26 @@
66
77
"""
88

9+
import logging
910
from collections.abc import Sequence
1011

11-
import logging
1212
import sqlalchemy as sa
1313

1414
try:
15-
from alembic import context, op
15+
from alembic import context
1616
except ImportError as e:
1717
raise ImportError(
1818
"'Add columns owner and last_updated to database tables' migration requires Alembic. Install with: 'pip install a2a-sdk[db-cli]'."
1919
) from e
2020

21+
from a2a.migrations.migration_utils import (
22+
table_exists,
23+
add_column,
24+
add_index,
25+
drop_column,
26+
drop_index,
27+
)
28+
2129

2230
# revision identifiers, used by Alembic.
2331
revision: str = '6419d2d130f6'
@@ -26,80 +34,6 @@
2634
depends_on: str | Sequence[str] | None = None
2735

2836

29-
def _get_inspector() -> sa.engine.reflection.Inspector:
30-
bind = op.get_bind()
31-
inspector = sa.inspect(bind)
32-
return inspector
33-
34-
35-
def _add_column(
36-
table: str,
37-
column_name: str,
38-
nullable: bool,
39-
type_: sa.types.TypeEngine,
40-
value: str | None = None,
41-
) -> None:
42-
if not _column_exists(table, column_name):
43-
op.add_column(
44-
table,
45-
sa.Column(
46-
column_name,
47-
type_,
48-
nullable=nullable,
49-
server_default=value,
50-
),
51-
)
52-
53-
54-
def _add_index(table: str, index_name: str, columns: list[str]) -> None:
55-
if not _index_exists(table, index_name):
56-
op.create_index(
57-
index_name,
58-
table,
59-
columns,
60-
)
61-
62-
63-
def _drop_column(table: str, column_name: str) -> None:
64-
if _column_exists(table, column_name, True):
65-
op.drop_column(table, column_name)
66-
67-
68-
def _drop_index(table: str, index_name: str) -> None:
69-
if _index_exists(table, index_name, True):
70-
op.drop_index(index_name, table_name=table)
71-
72-
73-
def _table_exists(table_name: str) -> bool:
74-
if context.is_offline_mode():
75-
return True
76-
bind = op.get_bind()
77-
inspector = sa.inspect(bind)
78-
return table_name in inspector.get_table_names()
79-
80-
81-
def _column_exists(
82-
table_name: str, column_name: str, downgrade_mode: bool = False
83-
) -> bool:
84-
if context.is_offline_mode():
85-
return downgrade_mode
86-
87-
inspector = _get_inspector()
88-
columns = [c['name'] for c in inspector.get_columns(table_name)]
89-
return column_name in columns
90-
91-
92-
def _index_exists(
93-
table_name: str, index_name: str, downgrade_mode: bool = False
94-
) -> bool:
95-
if context.is_offline_mode():
96-
return downgrade_mode
97-
98-
inspector = _get_inspector()
99-
indexes = [i['name'] for i in inspector.get_indexes(table_name)]
100-
return index_name in indexes
101-
102-
10337
def upgrade() -> None:
10438
"""Upgrade schema."""
10539
# Get the default value from the config (passed via CLI)
@@ -112,10 +46,10 @@ def upgrade() -> None:
11246
'push_notification_configs_table', 'push_notification_configs'
11347
)
11448

115-
if _table_exists(tasks_table):
116-
_add_column(tasks_table, 'owner', False, sa.String(128), owner)
117-
_add_column(tasks_table, 'last_updated', True, sa.DateTime())
118-
_add_index(
49+
if table_exists(tasks_table):
50+
add_column(tasks_table, 'owner', False, sa.String(255), owner)
51+
add_column(tasks_table, 'last_updated', True, sa.DateTime())
52+
add_index(
11953
tasks_table,
12054
f'idx_{tasks_table}_owner_last_updated',
12155
['owner', 'last_updated'],
@@ -125,14 +59,19 @@ def upgrade() -> None:
12559
f"Table '{tasks_table}' does not exist. Skipping upgrade for this table."
12660
)
12761

128-
if _table_exists(push_notification_configs_table):
129-
_add_column(
62+
if table_exists(push_notification_configs_table):
63+
add_column(
13064
push_notification_configs_table,
13165
'owner',
13266
False,
133-
sa.String(128),
67+
sa.String(255),
13468
owner,
13569
)
70+
add_index(
71+
push_notification_configs_table,
72+
f'ix_{push_notification_configs_table}_owner',
73+
['owner'],
74+
)
13675
else:
13776
logging.warning(
13877
f"Table '{push_notification_configs_table}' does not exist. Skipping upgrade for this table."
@@ -146,20 +85,24 @@ def downgrade() -> None:
14685
'push_notification_configs_table', 'push_notification_configs'
14786
)
14887

149-
if _table_exists(tasks_table):
150-
_drop_index(
88+
if table_exists(tasks_table):
89+
drop_index(
15190
tasks_table,
15291
f'idx_{tasks_table}_owner_last_updated',
15392
)
154-
_drop_column(tasks_table, 'owner')
155-
_drop_column(tasks_table, 'last_updated')
93+
drop_column(tasks_table, 'owner')
94+
drop_column(tasks_table, 'last_updated')
15695
else:
15796
logging.warning(
15897
f"Table '{tasks_table}' does not exist. Skipping downgrade for this table."
15998
)
16099

161-
if _table_exists(push_notification_configs_table):
162-
_drop_column(push_notification_configs_table, 'owner')
100+
if table_exists(push_notification_configs_table):
101+
drop_index(
102+
push_notification_configs_table,
103+
f'ix_{push_notification_configs_table}_owner',
104+
)
105+
drop_column(push_notification_configs_table, 'owner')
163106
else:
164107
logging.warning(
165108
f"Table '{push_notification_configs_table}' does not exist. Skipping downgrade for this table."

0 commit comments

Comments
 (0)