diff --git a/src/hqt/db/migrations.py b/src/hqt/db/migrations.py index 7ee828c..d234f6c 100644 --- a/src/hqt/db/migrations.py +++ b/src/hqt/db/migrations.py @@ -7,12 +7,20 @@ from hqt.db.models import Base log = logging.getLogger(__name__) + # Ordered schema migrations. Each entry upgrades the schema from the previous # version to `version`. Version 1 (BASELINE_VERSION) is the schema produced by # Base.metadata.create_all, so entries here start at version 2, e.g.: # (2, lambda conn: conn.exec_driver_sql( # "ALTER TABLE sessions ADD COLUMN foo TEXT")), -MIGRATIONS: list[tuple[int, Callable[[Connection], None]]] = [] +def _migrate_v2(conn: Connection) -> None: + conn.exec_driver_sql("ALTER TABLE sessions ADD COLUMN worktree_path TEXT") + conn.exec_driver_sql("ALTER TABLE sessions ADD COLUMN worktree_branch TEXT") + + +MIGRATIONS: list[tuple[int, Callable[[Connection], None]]] = [ + (2, _migrate_v2), +] BASELINE_VERSION = 1 LATEST_VERSION = MIGRATIONS[-1][0] if MIGRATIONS else BASELINE_VERSION diff --git a/src/hqt/db/models.py b/src/hqt/db/models.py index db9c03a..d90f89e 100644 --- a/src/hqt/db/models.py +++ b/src/hqt/db/models.py @@ -38,6 +38,8 @@ class Session(Base): tmux_session_name: Mapped[str] = mapped_column(unique=True) harness_session_id: Mapped[str | None] = mapped_column(default=None) model: Mapped[str | None] = mapped_column(default=None) + worktree_path: Mapped[str | None] = mapped_column(default=None) + worktree_branch: Mapped[str | None] = mapped_column(default=None) created_at: Mapped[datetime] = mapped_column(insert_default=func.now()) last_activity_at: Mapped[datetime | None] = mapped_column(default=None) archived: Mapped[bool] = mapped_column(default=False) diff --git a/tests/test_db.py b/tests/test_db.py index f080593..d28d876 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -86,15 +86,19 @@ def test_fresh_db_stamped_latest_version(tmp_path): assert _user_version(get_engine(settings)) == migrations.LATEST_VERSION -def test_unversioned_existing_db_stamped_baseline(tmp_path): +def test_unversioned_existing_db_stamped_latest(tmp_path): # Simulates a user DB created before versioning existed: tables present, - # user_version still 0. + # user_version still 0. migrate() should apply all pending migrations and + # stamp the DB at LATEST_VERSION (not just BASELINE_VERSION). settings = _tmp_settings(tmp_path) settings.db_path.parent.mkdir(parents=True, exist_ok=True) engine = get_engine(settings) Base.metadata.create_all(engine) + with engine.begin() as conn: + conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_path") + conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_branch") migrate(engine) - assert _user_version(engine) == migrations.BASELINE_VERSION + assert _user_version(engine) == migrations.LATEST_VERSION def test_db_from_newer_hqt_raises(tmp_path): @@ -109,20 +113,95 @@ def test_db_from_newer_hqt_raises(tmp_path): def test_pending_migrations_applied_in_order_once(tmp_path, monkeypatch): settings = _tmp_settings(tmp_path) - ensure_db(settings) # stamps BASELINE_VERSION + ensure_db(settings) # stamps migrations.LATEST_VERSION (currently 2) applied: list[int] = [] + # Use versions above current LATEST_VERSION so none have been applied yet. monkeypatch.setattr( migrations, "MIGRATIONS", - [(2, lambda conn: applied.append(2)), (3, lambda conn: applied.append(3))], + [ + *migrations.MIGRATIONS, + (3, lambda conn: applied.append(3)), + (4, lambda conn: applied.append(4)), + ], ) - monkeypatch.setattr(migrations, "LATEST_VERSION", 3) + monkeypatch.setattr(migrations, "LATEST_VERSION", 4) engine = get_engine(settings) migrate(engine) - assert applied == [2, 3] - assert _user_version(engine) == 3 + assert applied == [3, 4] + assert _user_version(engine) == 4 migrate(engine) # re-run is a no-op - assert applied == [2, 3] + assert applied == [3, 4] + + +def _session_column_names(engine) -> list[str]: + with engine.connect() as conn: + rows = conn.exec_driver_sql("PRAGMA table_info(sessions)").fetchall() + return [row[1] for row in rows] + + +def test_fresh_db_has_worktree_columns(tmp_path): + settings = _tmp_settings(tmp_path) + ensure_db(settings) + engine = get_engine(settings) + cols = _session_column_names(engine) + assert "worktree_path" in cols + assert "worktree_branch" in cols + assert _user_version(engine) == 2 + + +def test_v1_db_upgraded_to_v2(tmp_path): + # Build a v1 database: create all tables, drop the two new columns, stamp v1. + settings = _tmp_settings(tmp_path) + settings.db_path.parent.mkdir(parents=True, exist_ok=True) + engine = get_engine(settings) + Base.metadata.create_all(engine) + with engine.begin() as conn: + conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_path") + conn.exec_driver_sql("ALTER TABLE sessions DROP COLUMN worktree_branch") + conn.exec_driver_sql("PRAGMA user_version = 1") + + migrate(engine) + + cols = _session_column_names(engine) + assert "worktree_path" in cols + assert "worktree_branch" in cols + assert _user_version(engine) == 2 + + +def test_session_worktree_fields_default_none(tmp_path): + settings = _tmp_settings(tmp_path) + ensure_db(settings) + engine = get_engine(settings) + factory = get_session_factory(engine) + + with factory() as session: + p = Project(name="wt-proj", path="/tmp/wt-proj") + h = Harness(name="claude-code-wt") + session.add_all([p, h]) + session.commit() + + # Session without worktree fields + s_plain = Session( + project_id=p.id, + harness_id=h.id, + tmux_session_name="hqt-plain", + ) + # Session with worktree fields + s_wt = Session( + project_id=p.id, + harness_id=h.id, + tmux_session_name="hqt-worktree", + worktree_path="/tmp/wt", + worktree_branch="feature/x", + ) + session.add_all([s_plain, s_wt]) + session.commit() + + assert s_plain.worktree_path is None + assert s_plain.worktree_branch is None + assert s_wt.worktree_path == "/tmp/wt" + assert s_wt.worktree_branch == "feature/x" def test_rows_stay_readable_after_session_closes(tmp_path):