feat: rename a2a-db-cli to db-cli in installation instructions and update related CLI arguments

sokoliva · sokoliva · commit bf7961d718e8 · 2026-02-27T14:36:29.000Z
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,7 +39,7 @@ postgresql = ["sqlalchemy[asyncio,postgresql-asyncpg]>=2.0.0"]
 mysql = ["sqlalchemy[asyncio,aiomysql]>=2.0.0"]
 signing = ["PyJWT>=2.0.0"]
 sqlite = ["sqlalchemy[asyncio,aiosqlite]>=2.0.0"]
-a2a-db-cli = ["alembic>=1.14.0"]
+db-cli = ["alembic>=1.14.0"]
 
 sql = ["a2a-sdk[postgresql,mysql,sqlite]"]
 
@@ -50,7 +50,7 @@ all = [
   "a2a-sdk[grpc]",
   "a2a-sdk[telemetry]",
   "a2a-sdk[signing]",
-  "a2a-sdk[a2a-db-cli]",
+  "a2a-sdk[db-cli]",
 ]
 
 [project.urls]
diff --git a/src/a2a/a2a_db_cli.py b/src/a2a/a2a_db_cli.py
@@ -11,7 +11,7 @@
 
 except ImportError as e:
     raise ImportError(
-        "CLI requires Alembic. Install with: 'pip install a2a-sdk[a2a-db-cli]'."
+        "CLI requires Alembic. Install with: 'pip install a2a-sdk[db-cli]'."
     ) from e
 
 
@@ -28,14 +28,12 @@ def _add_shared_args(
     parser.add_argument(
         '--tasks-table',
         dest=f'{prefix}tasks_table',
-        help='Custom tasks table to update.',
-        action='append',
+        help='Custom tasks table to update. If not set, the default is "tasks".',
     )
     parser.add_argument(
-        '--push-notification-table',
-        dest=f'{prefix}push_notification_table',
-        help='Custom push notification configs table to update.',
-        action='append',
+        '--push-notification-configs-table',
+        dest=f'{prefix}push_notification_configs_table',
+        help='Custom push notification configs table to update. If not set, the default is "push_notification_configs".',
     )
     parser.add_argument(
         '-v',
@@ -122,29 +120,27 @@ def run_migrations() -> None:
     # Consolidate owner, db_url, tables, verbose and sql values
     owner = args.owner or getattr(args, 'sub_owner', None)
     db_url = args.database_url or getattr(args, 'sub_database_url', None)
-    task_tables = args.tasks_table or getattr(args, 'sub_tasks_table', None)
-    push_notification_tables = args.push_notification_table or getattr(
-        args, 'sub_push_notification_table', None
+    task_table = args.tasks_table or getattr(args, 'sub_tasks_table', None)
+    push_notification_configs_table = (
+        args.push_notification_configs_table
+        or getattr(args, 'sub_push_notification_configs_table', None)
     )
+
     verbose = args.verbose or getattr(args, 'sub_verbose', False)
     sql = args.sql or getattr(args, 'sub_sql', False)
 
     # Pass custom arguments to the migration context
     if owner:
-        if args.cmd == 'downgrade':
-            parser.error(
-                "The --add_columns_owner_last_updated-default-owner option is not supported for the 'downgrade' command."
-            )
         cfg.set_main_option(
             'add_columns_owner_last_updated_default_owner', owner
         )
     if db_url:
         os.environ['DATABASE_URL'] = db_url
-    if task_tables:
-        cfg.set_main_option('tasks_tables', ','.join(task_tables))
-    if push_notification_tables:
+    if task_table:
+        cfg.set_main_option('tasks_table', task_table)
+    if push_notification_configs_table:
         cfg.set_main_option(
-            'push_notification_tables', ','.join(push_notification_tables)
+            'push_notification_configs_table', push_notification_configs_table
         )
     if verbose:
         cfg.set_main_option('verbose', 'true')
diff --git a/src/a2a/migrations/README.md b/src/a2a/migrations/README.md
@@ -4,11 +4,11 @@ This directory handles the database schema updates for the A2A SDK. It uses a bu
 
 ## Installation
 
-To use the `a2a-db` migration tool, install the `a2a-sdk` with the `a2a-db-cli` extra.
+To use the `a2a-db` migration tool, install the `a2a-sdk` with the `db-cli` extra.
 
 | Extra | `uv` Command | `pip` Command |
 | :--- | :--- | :--- |
-| **CLI Only** | `uv add "a2a-sdk[a2a-db-cli]"` | `pip install "a2a-sdk[a2a-db-cli]"` |
+| **CLI Only** | `uv add "a2a-sdk[db-cli]"` | `pip install "a2a-sdk[db-cli]"` |
 | **All Extras** | `uv add "a2a-sdk[all]"` | `pip install "a2a-sdk[all]"` |
 
 
@@ -55,11 +55,11 @@ You can use the `--database-url` flag to specify the database URL for a single c
 ```bash
 uv run a2a-db --database-url "sqlite+aiosqlite:///my_database.db"
 ```
-#### --tasks-table / --push-notification-table
-The tables `tasks` and `push_notification_configs` are always updated by default. If your application uses additional tables that require the same migration, you can specify them using these flags.
+#### --tasks-table / --push-notification-configs-table
+Custom tasks and push notification configs tables to update. If not set, the default are `tasks` and `push_notification_configs`.
 
 ```bash
-uv run a2a-db --tasks-table "my_tasks" --push-notification-table "my_configs"
+uv run a2a-db --tasks-table "my_tasks" --push-notification-configs-table "my_configs"
 ```
 #### -v / --verbose
 Enables verbose output by setting `sqlalchemy.engine` logging to `INFO`.
diff --git a/src/a2a/migrations/env.py b/src/a2a/migrations/env.py
@@ -13,7 +13,7 @@
     from alembic import context
 except ImportError as e:
     raise ImportError(
-        "Migrations require Alembic. Install with: 'pip install a2a-sdk[a2a-db-cli]'."
+        "Migrations require Alembic. Install with: 'pip install a2a-sdk[db-cli]'."
     ) from e
 
 
diff --git a/src/a2a/migrations/script.py.mako b/src/a2a/migrations/script.py.mako
@@ -13,7 +13,7 @@ try:
     from alembic import op
 except ImportError as e:
     raise ImportError(
-        "A2A migrations require the 'a2a-db-cli' extra. Install with: 'pip install a2a-sdk[a2a-db-cli]'."
+        "A2A migrations require the 'db-cli' extra. Install with: 'pip install a2a-sdk[db-cli]'."
     ) from e
 
 ${imports if imports else ""}
diff --git a/src/a2a/migrations/versions/6419d2d130f6_add_columns_owner_last_updated.py b/src/a2a/migrations/versions/6419d2d130f6_add_columns_owner_last_updated.py
@@ -15,7 +15,7 @@
     from alembic import context, op
 except ImportError as e:
     raise ImportError(
-        "'Add columns owner and last_updated to database tables' migration requires Alembic. Install with: 'pip install a2a-sdk[a2a-db-cli]'."
+        "'Add columns owner and last_updated to database tables' migration requires Alembic. Install with: 'pip install a2a-sdk[db-cli]'."
     ) from e
 
 
@@ -32,6 +32,38 @@ def _get_inspector() -> sa.engine.reflection.Inspector:
     return inspector
 
 
+def _add_column(table: str, value: str, column_name: str) -> None:
+    if not _column_exists(table, column_name):
+        op.add_column(
+            table,
+            sa.Column(
+                column_name,
+                sa.String(128),
+                nullable=False,
+                server_default=value,
+            ),
+        )
+
+
+def _add_index(table: str, index_name: str, columns: list[str]) -> None:
+    if not _index_exists(table, index_name):
+        op.create_index(
+            index_name,
+            table,
+            columns,
+        )
+
+
+def _drop_column(table: str, column_name: str) -> None:
+    if _column_exists(table, column_name, True):
+        op.drop_column(table, column_name)
+
+
+def _drop_index(table: str, index_name: str) -> None:
+    if _index_exists(table, index_name, True):
+        op.drop_index(index_name, table_name=table)
+
+
 def _table_exists(table_name: str) -> bool:
     if context.is_offline_mode():
         return True
@@ -69,89 +101,54 @@ def upgrade() -> None:
         'add_columns_owner_last_updated_default_owner',
         'legacy_v03_no_user_info',
     )
-    tasks_tables = ['tasks']
-    push_notification_tables = ['push_notification_configs']
-
-    if tasks_tables_str := context.config.get_main_option('tasks_tables', None):
-        tasks_tables.extend([t.strip() for t in tasks_tables_str.split(',')])
-    if push_notification_tables_str := context.config.get_main_option(
-        'push_notification_tables', None
-    ):
-        push_notification_tables.extend(
-            [t.strip() for t in push_notification_tables_str.split(',')]
+    tasks_table = context.config.get_main_option('tasks_table', 'tasks')
+    push_notification_configs_table = context.config.get_main_option(
+        'push_notification_configs_table', 'push_notification_configs'
+    )
+
+    if _table_exists(tasks_table):
+        _add_column(tasks_table, owner, 'owner')
+        _add_column(tasks_table, '0', 'last_updated')
+        _add_index(
+            tasks_table,
+            f'idx_{tasks_table}_owner_last_updated',
+            ['owner', 'last_updated'],
+        )
+    else:
+        logging.warning(
+            f"Table '{tasks_table}' does not exist. Skipping upgrade for this table."
         )
 
-    for table in tasks_tables + push_notification_tables:
-        if _table_exists(table):
-            if not _column_exists(table, 'owner'):
-                op.add_column(
-                    table,
-                    sa.Column(
-                        'owner',
-                        sa.String(128),
-                        nullable=False,
-                        server_default=owner,
-                    ),
-                )
-        else:
-            if table in tasks_tables:
-                tasks_tables.remove(table)
-            logging.warning(
-                f"Table '{table}' does not exist. Skipping upgrade for this table."
-            )
-
-    for table in tasks_tables:
-        if _table_exists(table):
-            if not _column_exists(table, 'last_updated'):
-                op.add_column(
-                    table,
-                    sa.Column('last_updated', sa.String(22), nullable=True),
-                )
-            if not _index_exists(table, f'idx_{table}_owner_last_updated'):
-                op.create_index(
-                    f'idx_{table}_owner_last_updated',
-                    table,
-                    ['owner', 'last_updated'],
-                )
-        else:
-            logging.warning(
-                f"Table '{table}' does not exist. Skipping upgrade for this table."
-            )
+    if _table_exists(push_notification_configs_table):
+        _add_column(push_notification_configs_table, owner, 'owner')
+    else:
+        logging.warning(
+            f"Table '{push_notification_configs_table}' does not exist. Skipping upgrade for this table."
+        )
 
 
 def downgrade() -> None:
     """Downgrade schema."""
-    tasks_tables = ['tasks']
-    push_notification_tables = ['push_notification_configs']
-
-    if tasks_tables_str := context.config.get_main_option('tasks_tables', None):
-        tasks_tables.extend([t.strip() for t in tasks_tables_str.split(',')])
-    if push_notification_tables_str := context.config.get_main_option(
-        'push_notification_tables', None
-    ):
-        push_notification_tables.extend(
-            [t.strip() for t in push_notification_tables_str.split(',')]
+    tasks_table = context.config.get_main_option('tasks_table', 'tasks')
+    push_notification_configs_table = context.config.get_main_option(
+        'push_notification_configs_table', 'push_notification_configs'
+    )
+
+    if _table_exists(tasks_table):
+        _drop_index(
+            tasks_table,
+            f'idx_{tasks_table}_owner_last_updated',
+        )
+        _drop_column(tasks_table, 'owner')
+        _drop_column(tasks_table, 'last_updated')
+    else:
+        logging.warning(
+            f"Table '{tasks_table}' does not exist. Skipping downgrade for this table."
         )
 
-    for table in tasks_tables:
-        if _table_exists(table):
-            if _index_exists(table, f'idx_{table}_owner_last_updated', True):
-                op.drop_index(
-                    f'idx_{table}_owner_last_updated', table_name=table
-                )
-            if _column_exists(table, 'last_updated', True):
-                op.drop_column(table, 'last_updated')
-        else:
-            tasks_tables.remove(table)
-            logging.warning(
-                f"Table '{table}' does not exist. Skipping downgrade for this table."
-            )
-
-    for table in tasks_tables + push_notification_tables:
-        if _table_exists(table):
-            if _column_exists(table, 'owner', True):
-                op.drop_column(table, 'owner')
-        else:
-            logging.warning(
-                f"Table '{table}' does not exist. Skipping downgrade for this table."
-            )
+    if _table_exists(push_notification_configs_table):
+        _drop_column(push_notification_configs_table, 'owner')
+    else:
+        logging.warning(
+            f"Table '{push_notification_configs_table}' does not exist. Skipping downgrade for this table."
+        )
diff --git a/tests/migrations/test_a2a_db_cli.py b/tests/migrations/test_a2a_db_cli.py
@@ -26,7 +26,7 @@ def test_cli_upgrade_offline(mock_alembic_command, mock_alembic_config):
     tasks_table = 'my_tasks'
     push_table = 'my_push'
 
-    # Simulate: a2a-db upgrade head --sql --add_columns_owner_last_updated-default-ownertest-owner --tasks-table my_tasks --push-notification-table my_push -v
+    # Simulate: a2a-db upgrade head --sql --add_columns_owner_last_updated-default-ownertest-owner --tasks-table my_tasks --push-notification-configs-table my_push -v
     test_args = [
         'a2a-db',
         'upgrade',
@@ -36,7 +36,7 @@ def test_cli_upgrade_offline(mock_alembic_command, mock_alembic_config):
         custom_owner,
         '--tasks-table',
         tasks_table,
-        '--push-notification-table',
+        '--push-notification-configs-table',
         push_table,
         '-v',
     ]
@@ -55,10 +55,10 @@ def test_cli_upgrade_offline(mock_alembic_command, mock_alembic_config):
         'add_columns_owner_last_updated_default_owner', custom_owner
     )
     mock_alembic_config.return_value.set_main_option.assert_any_call(
-        'tasks_tables', tasks_table
+        'tasks_table', tasks_table
     )
     mock_alembic_config.return_value.set_main_option.assert_any_call(
-        'push_notification_tables', push_table
+        'push_notification_configs_table', push_table
     )
     mock_alembic_config.return_value.set_main_option.assert_any_call(
         'verbose', 'true'
@@ -88,7 +88,7 @@ def test_cli_downgrade_offline(mock_alembic_command, mock_alembic_config):
 
     # Verify tables option
     mock_alembic_config.return_value.set_main_option.assert_any_call(
-        'tasks_tables', tasks_table
+        'tasks_table', tasks_table
     )
 
 
diff --git a/tests/migrations/versions/test_migration_6419d2d130f6.py b/tests/migrations/versions/test_migration_6419d2d130f6.py
@@ -177,7 +177,7 @@ def test_migration_6419d2d130f6_custom_tables(
         db_url,
         '--tasks-table',
         custom_tasks,
-        '--push-notification-table',
+        '--push-notification-configs-table',
         custom_push,
         'upgrade',
         '6419d2d130f6',
diff --git a/uv.lock b/uv.lock
