-
Notifications
You must be signed in to change notification settings - Fork 28
feat(server): bundle migrations in wheel and add agent-control-migrate #209
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Merged
abhinav-galileo
merged 4 commits into
main
from
abhi/migrate-console-and-package-migrations
May 5, 2026
+352
−0
Merged
Changes from all commits
Commits
Show all changes
4 commits
Select commit
Hold shift + click to select a range
32a4839
feat(server): bundle migrations in wheel and add agent-control-migrate
abhinav-galileo beec8fb
test(server): cover migrate dispatch and force-include source paths
abhinav-galileo d891737
fix(server): require explicit revision for migrate downgrade
abhinav-galileo f21c959
fix(server): harden migrate CLI
abhinav-galileo File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,102 @@ | ||
| """Run bundled Alembic migrations for agent-control-server. | ||
|
|
||
| Exposed as the ``agent-control-migrate`` console script. The wheel ships | ||
| its Alembic config and migration scripts under the package so this | ||
| command works in any install location (Docker, venv, system Python). | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import argparse | ||
| import logging | ||
| import sys | ||
| from pathlib import Path | ||
|
|
||
| from alembic import command | ||
| from alembic.config import Config | ||
|
|
||
| import agent_control_server | ||
|
|
||
|
|
||
| def _bundled_config() -> Config: | ||
| pkg_dir = Path(agent_control_server.__file__).parent | ||
| ini_path = pkg_dir / "_alembic.ini" | ||
| alembic_dir = pkg_dir / "_alembic" | ||
| if not ini_path.exists() or not alembic_dir.exists(): | ||
| raise RuntimeError( | ||
| "Bundled Alembic resources not found. Expected " | ||
| f"{ini_path} and {alembic_dir}. The installed wheel is missing " | ||
| "migration assets." | ||
| ) | ||
| cfg = Config(str(ini_path)) | ||
| cfg.set_main_option("script_location", str(alembic_dir).replace("%", "%%")) | ||
| return cfg | ||
|
|
||
|
|
||
| def _build_parser() -> argparse.ArgumentParser: | ||
| parser = argparse.ArgumentParser( | ||
| prog="agent-control-migrate", | ||
| description="Run bundled Alembic migrations for agent-control-server.", | ||
| ) | ||
| subparsers = parser.add_subparsers(dest="command") | ||
|
|
||
| upgrade = subparsers.add_parser("upgrade", help="Upgrade to a revision.") | ||
| upgrade.add_argument("revision", nargs="?", default="head") | ||
| upgrade.add_argument("--sql", action="store_true", help="Emit SQL instead of executing.") | ||
|
|
||
| downgrade = subparsers.add_parser("downgrade", help="Downgrade to a revision.") | ||
| downgrade.add_argument("revision") | ||
| downgrade.add_argument("--sql", action="store_true", help="Emit SQL instead of executing.") | ||
|
|
||
| subparsers.add_parser("current", help="Show the current revision.") | ||
| subparsers.add_parser("history", help="List migration history.") | ||
| subparsers.add_parser("heads", help="Show current available heads.") | ||
| return parser | ||
|
|
||
|
|
||
| def _configure_logging() -> None: | ||
| logging.basicConfig( | ||
| level=logging.INFO, | ||
| format="%(levelname)s [%(name)s] %(message)s", | ||
| stream=sys.stderr, | ||
| ) | ||
|
|
||
|
|
||
| def main(argv: list[str] | None = None) -> int: | ||
| """Entry point for the ``agent-control-migrate`` console script. | ||
|
|
||
| With no arguments, runs ``upgrade head``. Supports a small subset of | ||
| Alembic commands sufficient for deploys and operational debugging: | ||
| ``upgrade``, ``downgrade``, ``current``, ``history``, ``heads``. | ||
| """ | ||
| args = list(argv) if argv is not None else sys.argv[1:] | ||
| if not args: | ||
| args = ["upgrade", "head"] | ||
|
abhinav-galileo marked this conversation as resolved.
|
||
|
|
||
| parser = _build_parser() | ||
| parsed = parser.parse_args(args) | ||
| _configure_logging() | ||
|
|
||
| try: | ||
| cfg = _bundled_config() | ||
| if parsed.command == "upgrade": | ||
| command.upgrade(cfg, parsed.revision, sql=parsed.sql) | ||
| elif parsed.command == "downgrade": | ||
| command.downgrade(cfg, parsed.revision, sql=parsed.sql) | ||
| elif parsed.command == "current": | ||
| command.current(cfg) | ||
| elif parsed.command == "history": | ||
| command.history(cfg) | ||
| elif parsed.command == "heads": | ||
| command.heads(cfg) | ||
| else: # pragma: no cover - argparse guarantees this cannot happen. | ||
| parser.error("missing command") | ||
| except Exception as exc: | ||
| print(f"agent-control-migrate: {exc}", file=sys.stderr) | ||
| return 1 | ||
|
|
||
| return 0 | ||
|
|
||
|
|
||
| if __name__ == "__main__": | ||
| raise SystemExit(main()) | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,242 @@ | ||
| """Unit tests for the bundled-migrations entry point. | ||
|
abhinav-galileo marked this conversation as resolved.
|
||
|
|
||
| These do not run migrations against a database. They verify the wheel-bundling | ||
| contract: the console script resolves to the right callable, dispatches | ||
| correctly to Alembic commands, and the bundled-config helper can load the | ||
| packaged migration layout. | ||
| """ | ||
|
|
||
| from __future__ import annotations | ||
|
|
||
| import tomllib | ||
| from pathlib import Path, PurePosixPath | ||
| from unittest.mock import MagicMock | ||
|
|
||
| import pytest | ||
| from alembic.script import ScriptDirectory | ||
|
|
||
| from agent_control_server import migrate | ||
|
|
||
|
|
||
| @pytest.fixture | ||
| def stub_config(monkeypatch: pytest.MonkeyPatch) -> object: | ||
| """Replace bundled-config building with a sentinel object. | ||
|
|
||
| Lets dispatch tests verify which Alembic command was called and | ||
| what config was passed without needing real migration assets. | ||
| """ | ||
| sentinel = object() | ||
| monkeypatch.setattr(migrate, "_bundled_config", lambda: sentinel) | ||
| return sentinel | ||
|
|
||
|
|
||
| def _patch_command(monkeypatch: pytest.MonkeyPatch, name: str) -> MagicMock: | ||
| mock = MagicMock() | ||
| monkeypatch.setattr(migrate.command, name, mock) | ||
| return mock | ||
|
|
||
|
|
||
| def test_main_default_runs_upgrade_head( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| upgrade = _patch_command(monkeypatch, "upgrade") | ||
| rc = migrate.main([]) | ||
| assert rc == 0 | ||
| upgrade.assert_called_once_with(stub_config, "head", sql=False) | ||
|
|
||
|
|
||
| def test_main_bare_upgrade_runs_upgrade_head( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| upgrade = _patch_command(monkeypatch, "upgrade") | ||
| rc = migrate.main(["upgrade"]) | ||
| assert rc == 0 | ||
| upgrade.assert_called_once_with(stub_config, "head", sql=False) | ||
|
|
||
|
|
||
| def test_main_explicit_upgrade_revision( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| upgrade = _patch_command(monkeypatch, "upgrade") | ||
| rc = migrate.main(["upgrade", "abc123"]) | ||
| assert rc == 0 | ||
| upgrade.assert_called_once_with(stub_config, "abc123", sql=False) | ||
|
|
||
|
|
||
| def test_main_upgrade_supports_sql( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| upgrade = _patch_command(monkeypatch, "upgrade") | ||
| rc = migrate.main(["upgrade", "head", "--sql"]) | ||
| assert rc == 0 | ||
| upgrade.assert_called_once_with(stub_config, "head", sql=True) | ||
|
|
||
|
|
||
| def test_main_bare_downgrade_requires_explicit_revision( | ||
| monkeypatch: pytest.MonkeyPatch, | ||
| ) -> None: | ||
| monkeypatch.setattr(migrate, "_bundled_config", pytest.fail) | ||
| with pytest.raises(SystemExit) as exc_info: | ||
| migrate.main(["downgrade"]) | ||
| assert exc_info.value.code == 2 | ||
|
|
||
|
|
||
| def test_main_explicit_downgrade_revision( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| downgrade = _patch_command(monkeypatch, "downgrade") | ||
| rc = migrate.main(["downgrade", "abc123"]) | ||
| assert rc == 0 | ||
| downgrade.assert_called_once_with(stub_config, "abc123", sql=False) | ||
|
|
||
|
|
||
| def test_main_downgrade_supports_sql( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| downgrade = _patch_command(monkeypatch, "downgrade") | ||
| rc = migrate.main(["downgrade", "-1", "--sql"]) | ||
| assert rc == 0 | ||
| downgrade.assert_called_once_with(stub_config, "-1", sql=True) | ||
|
|
||
|
|
||
| @pytest.mark.parametrize("op", ["current", "history", "heads"]) | ||
| def test_main_query_commands( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch, op: str | ||
| ) -> None: | ||
| cmd = _patch_command(monkeypatch, op) | ||
| rc = migrate.main([op]) | ||
| assert rc == 0 | ||
| cmd.assert_called_once_with(stub_config) | ||
|
|
||
|
|
||
| def test_main_unknown_command_returns_nonzero(monkeypatch: pytest.MonkeyPatch) -> None: | ||
| monkeypatch.setattr(migrate, "_bundled_config", pytest.fail) | ||
| with pytest.raises(SystemExit) as exc_info: | ||
| migrate.main(["does-not-exist"]) | ||
| assert exc_info.value.code == 2 | ||
|
|
||
|
|
||
| def test_main_unknown_command_prints_usage( | ||
| monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] | ||
| ) -> None: | ||
| monkeypatch.setattr(migrate, "_bundled_config", pytest.fail) | ||
| with pytest.raises(SystemExit): | ||
| migrate.main(["does-not-exist"]) | ||
| out = capsys.readouterr() | ||
| assert "invalid choice: 'does-not-exist'" in out.err | ||
| assert "usage:" in out.err | ||
|
|
||
|
|
||
| def test_main_help_prints_usage( | ||
| monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] | ||
| ) -> None: | ||
| monkeypatch.setattr(migrate, "_bundled_config", pytest.fail) | ||
| with pytest.raises(SystemExit) as exc_info: | ||
| migrate.main(["--help"]) | ||
| assert exc_info.value.code == 0 | ||
| out = capsys.readouterr() | ||
| assert "usage:" in out.out | ||
| assert "Run bundled Alembic migrations" in out.out | ||
|
|
||
|
|
||
| def test_main_rejects_extra_positional_args(monkeypatch: pytest.MonkeyPatch) -> None: | ||
| monkeypatch.setattr(migrate, "_bundled_config", pytest.fail) | ||
| with pytest.raises(SystemExit) as exc_info: | ||
| migrate.main(["upgrade", "head", "typo"]) | ||
| assert exc_info.value.code == 2 | ||
|
|
||
|
|
||
| def test_main_returns_nonzero_for_command_errors( | ||
| stub_config: object, monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] | ||
| ) -> None: | ||
| upgrade = _patch_command(monkeypatch, "upgrade") | ||
| upgrade.side_effect = RuntimeError("database unavailable") | ||
| rc = migrate.main(["upgrade", "head"]) | ||
| assert rc == 1 | ||
| out = capsys.readouterr() | ||
| assert "agent-control-migrate: database unavailable" in out.err | ||
|
|
||
|
|
||
| def test_bundled_config_raises_when_assets_missing( | ||
| tmp_path: Path, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| # Point the bundled-config lookup at a directory with no migration assets. | ||
| fake_pkg_init = tmp_path / "__init__.py" | ||
| fake_pkg_init.write_text("") | ||
| monkeypatch.setattr(migrate.agent_control_server, "__file__", str(fake_pkg_init)) | ||
|
|
||
| with pytest.raises(RuntimeError, match="Bundled Alembic resources not found"): | ||
| migrate._bundled_config() | ||
|
|
||
|
|
||
| def test_bundled_config_loads_real_bundled_layout( | ||
| tmp_path: Path, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| pkg_dir = tmp_path / "agent_control_server" | ||
| versions_dir = pkg_dir / "_alembic" / "versions" | ||
| versions_dir.mkdir(parents=True) | ||
| (pkg_dir / "__init__.py").write_text("") | ||
| (pkg_dir / "_alembic.ini").write_text("[alembic]\nscript_location = unused\n") | ||
| (pkg_dir / "_alembic" / "env.py").write_text("") | ||
| (pkg_dir / "_alembic" / "script.py.mako").write_text("") | ||
| (versions_dir / "abc123_initial.py").write_text( | ||
| '"""Initial revision."""\n' | ||
| 'revision = "abc123"\n' | ||
| "down_revision = None\n" | ||
| "branch_labels = None\n" | ||
| "depends_on = None\n" | ||
| ) | ||
| monkeypatch.setattr( | ||
| migrate.agent_control_server, | ||
| "__file__", | ||
| str(pkg_dir / "__init__.py"), | ||
| ) | ||
|
|
||
| cfg = migrate._bundled_config() | ||
| script_dir = ScriptDirectory.from_config(cfg) | ||
|
|
||
| assert script_dir.get_heads() == ["abc123"] | ||
|
|
||
|
|
||
| def test_bundled_config_escapes_percent_paths( | ||
| tmp_path: Path, monkeypatch: pytest.MonkeyPatch | ||
| ) -> None: | ||
| pkg_dir = tmp_path / "agent%control_server" | ||
| (pkg_dir / "_alembic").mkdir(parents=True) | ||
| (pkg_dir / "__init__.py").write_text("") | ||
| (pkg_dir / "_alembic.ini").write_text("[alembic]\nscript_location = unused\n") | ||
| monkeypatch.setattr( | ||
| migrate.agent_control_server, | ||
| "__file__", | ||
| str(pkg_dir / "__init__.py"), | ||
| ) | ||
|
|
||
| cfg = migrate._bundled_config() | ||
|
|
||
| assert cfg.get_main_option("script_location") == str(pkg_dir / "_alembic") | ||
|
|
||
|
|
||
| def test_force_include_source_paths_exist() -> None: | ||
| """Hatch force-include mappings must ship real migration assets under the package.""" | ||
| server_dir = Path(__file__).resolve().parent.parent | ||
| with (server_dir / "pyproject.toml").open("rb") as pyproject: | ||
| config = tomllib.load(pyproject) | ||
|
|
||
| scripts = config["project"]["scripts"] | ||
| assert scripts["agent-control-migrate"] == "agent_control_server.migrate:main" | ||
|
|
||
| wheel_config = config["tool"]["hatch"]["build"]["targets"]["wheel"] | ||
| force_include = wheel_config["force-include"] | ||
| assert force_include | ||
|
|
||
| for source, target in force_include.items(): | ||
| source_path = server_dir / source | ||
| assert source_path.exists(), f"missing force-include source: {source_path}" | ||
|
|
||
| target_path = PurePosixPath(target) | ||
| assert target_path.parts[0] == "agent_control_server" | ||
|
|
||
| alembic_target = force_include["alembic"] | ||
| versions = list((server_dir / "alembic" / "versions").glob("*.py")) | ||
| assert alembic_target == "agent_control_server/_alembic" | ||
| assert versions, f"no migration scripts under {server_dir / 'alembic' / 'versions'}" | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.