"""Migration 036: provenance, the operator switch, and membership intervals. The migration makes two irreversible-feeling judgement calls, and both need to be pinned: 1. **Existing rows become ``manual``.** Provenance cannot be recovered from ``created_at``, and manual is the fail-safe guess — manual rows are never pruned, so a wrong guess costs an un-pruned symbol rather than one swept out of the scanner. 2. **Existing "off" states are carried into ``scan_disabled_on``.** Whoever set ``untracked_on`` is ambiguous, and re-entry clears the prune flag. Copying it is what stops an automated sync reactivating something switched off by hand. """ from __future__ import annotations import importlib.util from pathlib import Path import sqlalchemy as sa from alembic.migration import MigrationContext from alembic.operations import Operations def _load_migration_module(): path = ( Path(__file__).resolve().parents[2] / "alembic" / "versions" / "036_universe_membership.py" ) spec = importlib.util.spec_from_file_location("migration_036", path) assert spec is not None and spec.loader is not None module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module def _tickers_table(connection: sa.Connection) -> None: """The schema as migration 035 leaves it.""" connection.execute(sa.text( """ CREATE TABLE tickers ( id INTEGER PRIMARY KEY, symbol VARCHAR(10) NOT NULL, delisted_on DATE, untracked_on DATE ) """ )) connection.execute(sa.text( "INSERT INTO tickers (id, symbol, delisted_on, untracked_on) VALUES " "(1, 'AAPL', NULL, NULL), " "(2, 'OFF', NULL, '2026-08-01')" )) def _columns(connection: sa.Connection, table: str) -> set[str]: return {row[1] for row in connection.execute(sa.text(f"PRAGMA table_info({table})"))} def test_upgrade_defaults_existing_rows_to_manual(): migration = _load_migration_module() engine = sa.create_engine("sqlite://") with engine.begin() as connection: _tickers_table(connection) with Operations.context(MigrationContext.configure(connection)): migration.upgrade() rows = dict(connection.execute( sa.text("SELECT symbol, tracking_source FROM tickers") ).all()) assert rows == {"AAPL": "manual", "OFF": "manual"} def test_upgrade_carries_an_existing_off_state_into_the_operator_switch(): migration = _load_migration_module() engine = sa.create_engine("sqlite://") with engine.begin() as connection: _tickers_table(connection) with Operations.context(MigrationContext.configure(connection)): migration.upgrade() rows = dict(connection.execute( sa.text("SELECT symbol, scan_disabled_on FROM tickers") ).all()) # The ambiguous existing "off" is preserved on the column no sync clears. assert rows["OFF"] == "2026-08-01" # A row that was never off stays on. assert rows["AAPL"] is None def test_upgrade_creates_the_membership_table_empty(): migration = _load_migration_module() engine = sa.create_engine("sqlite://") with engine.begin() as connection: _tickers_table(connection) with Operations.context(MigrationContext.configure(connection)): migration.upgrade() assert {"ticker_id", "universe", "first_seen_on", "absent_since"} <= _columns( connection, "universe_memberships" ) # History starts at first observation — the migration invents none. count = connection.execute( sa.text("SELECT COUNT(*) FROM universe_memberships") ).scalar_one() assert count == 0 def test_downgrade_removes_the_new_state_and_keeps_the_tickers(): migration = _load_migration_module() engine = sa.create_engine("sqlite://") with engine.begin() as connection: _tickers_table(connection) with Operations.context(MigrationContext.configure(connection)): migration.upgrade() with Operations.context(MigrationContext.configure(connection)): migration.downgrade() columns = _columns(connection, "tickers") assert "tracking_source" not in columns assert "scan_disabled_on" not in columns # 035's column and the rows themselves are untouched. assert "untracked_on" in columns assert connection.execute( sa.text("SELECT symbol FROM tickers ORDER BY symbol") ).all() == [("AAPL",), ("OFF",)]