diff --git a/.github/workflows/migration-validation.yml b/.github/workflows/migration-validation.yml index a129baf7df..066ee8ca86 100644 --- a/.github/workflows/migration-validation.yml +++ b/.github/workflows/migration-validation.yml @@ -3,9 +3,10 @@ name: Database Migration Validation on: pull_request: paths: - - 'src/backend/base/langflow/alembic/versions/*.py' - - 'src/backend/base/langflow/services/database/models/**/*.py' - - 'src/backend/tests/unit/alembic/test_migration_execution.py' + - 'src/backend/base/langflow/alembic/**' + - 'src/backend/base/langflow/services/database/models/**' + - 'src/backend/base/langflow/services/database/service.py' + - 'src/backend/tests/unit/alembic/**' - '.github/workflows/migration-validation.yml' jobs: @@ -13,6 +14,21 @@ jobs: name: Model/Migration Consistency runs-on: ubuntu-latest + services: + postgres: + image: postgres:16 + env: + POSTGRES_USER: langflow + POSTGRES_PASSWORD: langflow + POSTGRES_DB: langflow + ports: + - 5432:5432 + options: >- + --health-cmd="pg_isready -U langflow" + --health-interval=10s + --health-timeout=5s + --health-retries=5 + steps: - name: Checkout code uses: actions/checkout@v6 @@ -29,11 +45,12 @@ jobs: - name: Install dependencies run: | - uv sync + uv sync --extra postgresql - name: Check model/migration consistency env: MIGRATION_VALIDATION_CI: "true" + LANGFLOW_TEST_DATABASE_URI: "postgresql://langflow:langflow@localhost:5432/langflow" run: | uv run pytest src/backend/tests/unit/alembic/test_migration_execution.py -x -v diff --git a/src/backend/tests/unit/alembic/test_migration_execution.py b/src/backend/tests/unit/alembic/test_migration_execution.py index 3a9396a052..6bdfcef820 100644 --- a/src/backend/tests/unit/alembic/test_migration_execution.py +++ b/src/backend/tests/unit/alembic/test_migration_execution.py @@ -1,5 +1,6 @@ import errno import os +import re import shutil import subprocess import tempfile @@ -11,117 +12,196 @@ from alembic.autogenerate import compare_metadata from alembic.config import Config from alembic.migration import MigrationContext from langflow.services.database.service import SQLModel -from sqlalchemy import create_engine +from sqlalchemy import create_engine, text _WORKSPACE_ROOT = Path(__file__).resolve().parents[5] +_SCRIPT_LOCATION = _WORKSPACE_ROOT / "src/backend/base/langflow/alembic" -def _get_alembic_cfg(db_path: str) -> Config: +def _make_alembic_cfg(db_url: str) -> Config: """Create an Alembic Config pointing at the project's migration scripts.""" alembic_cfg = Config() - script_location = _WORKSPACE_ROOT / "src/backend/base/langflow/alembic" - if not script_location.exists(): - pytest.fail(f"Alembic script location not found at {script_location}") + if not _SCRIPT_LOCATION.exists(): + pytest.fail(f"Alembic script location not found at {_SCRIPT_LOCATION}") - alembic_cfg.set_main_option("script_location", str(script_location)) - alembic_cfg.set_main_option("sqlalchemy.url", f"sqlite+aiosqlite:///{db_path}") + alembic_cfg.set_main_option("script_location", str(_SCRIPT_LOCATION)) + alembic_cfg.set_main_option("sqlalchemy.url", db_url) return alembic_cfg +# --------------------------------------------------------------------------- +# Database fixtures +# --------------------------------------------------------------------------- + + +def _normalize_pg_url(url: str) -> str: + """Ensure a Postgres URL uses the psycopg (v3) async-capable driver. + + Alembic's env.py uses async_engine_from_config, which requires an + async-capable dialect. The psycopg driver supports both sync and async. + """ + if url.startswith("postgresql://"): + return url.replace("postgresql://", "postgresql+psycopg://", 1) + if url.startswith("postgres://"): + return url.replace("postgres://", "postgresql+psycopg://", 1) + return url + + +def _pg_url() -> str | None: + """Return a PostgreSQL URL from the environment, or None.""" + url = os.environ.get("LANGFLOW_TEST_DATABASE_URI") + if url is not None: + return _normalize_pg_url(url) + return None + + +def _create_pg_test_database(base_url: str, db_name: str) -> str: + """Create an isolated test database and return its URL.""" + engine = create_engine(base_url, isolation_level="AUTOCOMMIT") + try: + with engine.connect() as conn: + conn.execute( + text( + # db_name is generated internally from a hash, not user input + f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity " # noqa: S608 + f"WHERE datname = '{db_name}' AND pid <> pg_backend_pid()" + ) + ) + conn.execute(text(f"DROP DATABASE IF EXISTS {db_name}")) + conn.execute(text(f"CREATE DATABASE {db_name}")) + finally: + engine.dispose() + return base_url.rsplit("/", 1)[0] + f"/{db_name}" + + +def _drop_pg_test_database(base_url: str, db_name: str) -> None: + """Drop the test database.""" + engine = create_engine(base_url, isolation_level="AUTOCOMMIT") + try: + with engine.connect() as conn: + conn.execute( + text( + # db_name is generated internally from a hash, not user input + f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity " # noqa: S608 + f"WHERE datname = '{db_name}' AND pid <> pg_backend_pid()" + ) + ) + conn.execute(text(f"DROP DATABASE IF EXISTS {db_name}")) + finally: + engine.dispose() + + +@pytest.fixture(params=["sqlite", "postgres"]) +def db_url(request): + """Parametrized fixture that yields a database URL for each backend.""" + if request.param == "sqlite": + with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp: + db_path = tmp.name + yield f"sqlite+aiosqlite:///{db_path}" + for suffix in ("", "-wal", "-shm", "-journal"): + Path(db_path + suffix).unlink(missing_ok=True) + else: + base_url = _pg_url() + if base_url is None: + pytest.skip("LANGFLOW_TEST_DATABASE_URI not set") + # Use a unique DB name per test to allow parallel execution + import hashlib + + short_hash = hashlib.md5(request.node.name.encode()).hexdigest()[:10] # noqa: S324 + db_name = f"lf_mig_test_{short_hash}" + test_url = _create_pg_test_database(base_url, db_name) + yield test_url + _drop_pg_test_database(base_url, db_name) + + +def _parse_revision_values(line: str) -> list[str]: + """Extract revision ID(s) from a line like ``revision: str = "abc123"``. + + Handles single strings, tuples of strings, and None. Returns a list of + zero or more revision ID strings. + """ + if "=" not in line: + return [] + raw = line.split("=", 1)[1] + # Strip inline comments (e.g. "# pragma: allowlist secret") + if "#" in raw: + raw = raw[: raw.index("#")] + raw = raw.strip() + if raw == "None": + return [] + # Extract all quoted strings from the value (handles both single values + # and tuples like ("abc", "def")) + return re.findall(r"""["']([a-f0-9]+)["']""", raw) + + def _get_main_branch_head() -> str | None: """Get the alembic head revision that origin/main is at. - Finds migration files new on this branch (not on origin/main), then looks up - their down_revision to determine where main's DB would be. Uses the current - branch's alembic ScriptDirectory since it already contains all migrations. + Uses ``git grep`` to read the ``revision`` and ``down_revision`` variables + directly from migration files on origin/main, then walks the chain to find + the head revision. This avoids relying on filename conventions (which may + not match the actual revision IDs inside the files) and works regardless of + whether the branch adds, modifies, or deletes migration files. Returns None if git operations fail (e.g. shallow clone without origin/main). """ - from alembic.script import ScriptDirectory - git = shutil.which("git") if git is None: return None - # Find migration files that are new on this branch vs origin/main - try: - result = subprocess.run( # noqa: S603 - [ - git, - "diff", - "--name-only", - "--diff-filter=A", - "origin/main...HEAD", - "--", - "src/backend/base/langflow/alembic/versions/*.py", - ], - capture_output=True, - text=True, - check=True, - cwd=_WORKSPACE_ROOT, - ) - except subprocess.CalledProcessError as exc: - import warnings - - warnings.warn(f"git diff failed (rc={exc.returncode}): {exc.stderr.strip()}", stacklevel=2) - return None - except OSError as exc: - if exc.errno == errno.ENOENT: - return None # git binary not found at resolved path - raise # unexpected OS error (disk full, permissions, etc.) - - new_files = [f for f in result.stdout.strip().splitlines() if f.endswith(".py")] - - alembic_cfg = Config() - script_location = _WORKSPACE_ROOT / "src/backend/base/langflow/alembic" - alembic_cfg.set_main_option("script_location", str(script_location)) - script = ScriptDirectory.from_config(alembic_cfg) - - if not new_files: - # No new migrations on this branch — head is same as main - heads = script.get_heads() - if len(heads) == 1: - return heads[0] - if len(heads) > 1: - pytest.fail(f"Alembic has {len(heads)} head revisions — migration branches need merging: {heads}") - return None - - # Collect revision IDs of all new migrations - new_rev_ids = set() - for fpath in new_files: - filename = Path(fpath).name - new_rev_ids.add(filename.split("_", 1)[0]) - - # Find down_revisions that point outside the new migrations (i.e. into main) - main_revisions = set() - for rev_id in new_rev_ids: - rev_script = script.get_revision(rev_id) - if rev_script is None: - msg = ( - f"New migration file matched revision ID '{rev_id}' " - f"but Alembic has no such revision — check filename convention" + def _git_grep(pattern: str) -> str | None: + try: + result = subprocess.run( # noqa: S603 + [ + git, + "grep", + "-h", + pattern, + "origin/main", + "--", + "src/backend/base/langflow/alembic/versions/", + ], + capture_output=True, + text=True, + check=True, + cwd=_WORKSPACE_ROOT, ) - raise ValueError(msg) - if rev_script.down_revision is None: - msg = f"New migration {rev_id} has down_revision=None — it must chain from an existing migration" - raise ValueError(msg) - down = rev_script.down_revision - downs = set(down) if isinstance(down, (tuple, list)) else {down} - # Only keep down_revisions that are NOT themselves new migrations - main_revisions.update(downs - new_rev_ids) + except subprocess.CalledProcessError: + return None + except OSError as exc: + if exc.errno == errno.ENOENT: + return None + raise + return result.stdout - if len(main_revisions) > 1: - pytest.fail( - f"New migrations descend from {len(main_revisions)} different base revisions — " - f"they must share a single parent on main: {main_revisions}" - ) - if not main_revisions: - pytest.fail( - f"New migrations {new_rev_ids} form a disconnected chain — " - f"none of their down_revisions point to an existing migration on main" - ) - return main_revisions.pop() + # Extract all revision IDs from origin/main's migration files + rev_output = _git_grep("^revision:") + if not rev_output: + return None + + main_rev_ids: set[str] = set() + for line in rev_output.strip().splitlines(): + main_rev_ids.update(_parse_revision_values(line)) + + if not main_rev_ids: + return None + + # Extract all down_revision IDs to determine the chain + down_output = _git_grep("^down_revision:") + referenced: set[str] = set() + if down_output: + for line in down_output.strip().splitlines(): + referenced.update(_parse_revision_values(line)) + + # Head = revisions not referenced as down_revision by any other revision + heads = main_rev_ids - referenced + + if len(heads) == 1: + return heads.pop() + if len(heads) > 1: + pytest.fail(f"origin/main has {len(heads)} head revisions — migration branches need merging: {heads}") + return None def _filter_sqlite_noise(diffs: list) -> list: @@ -254,7 +334,21 @@ class TestFilterSqliteNoise: assert result == diffs -def test_no_phantom_migrations(): +def _engine_url(db_url: str) -> str: + """Convert an async DB URL to a sync one for SQLAlchemy create_engine.""" + if db_url.startswith("sqlite+aiosqlite"): + return db_url.replace("sqlite+aiosqlite", "sqlite", 1) + return db_url + + +def _filter_diffs(diffs: list, db_url: str) -> list: + """Apply backend-appropriate diff filtering.""" + if "sqlite" in db_url: + return _filter_sqlite_noise(diffs) + return list(diffs) + + +def test_no_phantom_migrations(db_url): """Verify that models and migrations are in sync. After migrating a fresh database to head, autogenerate should detect @@ -262,91 +356,93 @@ def test_no_phantom_migrations(): (e.g. pydantic, sqlmodel) change how column metadata is emitted, which would produce unintended migration diffs. """ - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp: - db_path = tmp.name + alembic_cfg = _make_alembic_cfg(db_url) + command.upgrade(alembic_cfg, "head") + engine = create_engine(_engine_url(db_url)) try: - alembic_cfg = _get_alembic_cfg(db_path) - command.upgrade(alembic_cfg, "head") - - engine = create_engine(f"sqlite:///{db_path}") - try: - with engine.connect() as connection: - migration_context = MigrationContext.configure(connection) - diffs = compare_metadata(migration_context, SQLModel.metadata) - finally: - engine.dispose() - - significant_diffs = _filter_sqlite_noise(diffs) - - if significant_diffs: - diff_descriptions = "\n".join(str(d) for d in significant_diffs) - pytest.fail( - f"Autogenerate detected {len(significant_diffs)} unexpected change(s) " - f"after migrating to head. This likely means a dependency upgrade changed " - f"how column metadata is generated.\n\nDiffs:\n{diff_descriptions}" - ) + with engine.connect() as connection: + migration_context = MigrationContext.configure(connection) + diffs = compare_metadata(migration_context, SQLModel.metadata) finally: - for suffix in ("", "-wal", "-shm", "-journal"): - Path(db_path + suffix).unlink(missing_ok=True) + engine.dispose() + + significant_diffs = _filter_diffs(diffs, db_url) + + if significant_diffs: + diff_descriptions = "\n".join(str(d) for d in significant_diffs) + pytest.fail( + f"Autogenerate detected {len(significant_diffs)} unexpected change(s) " + f"after migrating to head. This likely means a dependency upgrade changed " + f"how column metadata is generated.\n\nDiffs:\n{diff_descriptions}" + ) -def test_upgrade_from_main_branch(): +def test_upgrade_from_main_branch(db_url): """Verify that a DB at main's head can upgrade to current head and downgrade back. This catches the real-world scenario: a user running on main (or the latest release) upgrades to a branch with new migrations. The upgrade must succeed, the resulting schema must match the models, and downgrade back to main must also succeed. """ + from alembic.script import ScriptDirectory + main_head = _get_main_branch_head() if main_head is None: if os.environ.get("MIGRATION_VALIDATION_CI"): pytest.fail("Could not determine main branch head revision — ensure fetch-depth: 0 and origin/main exists") pytest.skip("Could not determine main branch head revision (shallow clone or no origin/main)") - with tempfile.NamedTemporaryFile(suffix=".db", delete=False) as tmp: - db_path = tmp.name + # Check if main and branch share the same alembic head (no new migrations). + # In that case this test is a no-op — alembic won't re-run already-applied + # migrations, so upgrade(main_head) -> upgrade(head) does nothing. + # Modified migrations are exercised by test_no_phantom_migrations instead. + branch_cfg = Config() + branch_cfg.set_main_option("script_location", str(_SCRIPT_LOCATION)) + branch_script = ScriptDirectory.from_config(branch_cfg) + branch_heads = branch_script.get_heads() + if len(branch_heads) == 1 and branch_heads[0] == main_head: + pytest.skip( + "No new migrations on this branch — main and branch share the same " + "alembic head. Modified migrations are tested by test_no_phantom_migrations." + ) + alembic_cfg = _make_alembic_cfg(db_url) + + # Step 1: Create DB at main's head revision (simulates existing user DB) + command.upgrade(alembic_cfg, main_head) + + # Step 2: Upgrade to the current branch head + command.upgrade(alembic_cfg, "head") + + # Step 3: Verify models match the migrated DB + engine = create_engine(_engine_url(db_url)) try: - alembic_cfg = _get_alembic_cfg(db_path) - - # Step 1: Create DB at main's head revision (simulates existing user DB) - command.upgrade(alembic_cfg, main_head) - - # Step 2: Upgrade to the current branch head - command.upgrade(alembic_cfg, "head") - - # Step 3: Verify models match the migrated DB - engine = create_engine(f"sqlite:///{db_path}") - try: - with engine.connect() as connection: - migration_context = MigrationContext.configure(connection) - diffs = compare_metadata(migration_context, SQLModel.metadata) - finally: - engine.dispose() - - significant_diffs = _filter_sqlite_noise(diffs) - - if significant_diffs: - diff_descriptions = "\n".join(str(d) for d in significant_diffs) - pytest.fail( - f"After upgrading from main ({main_head}) to head, " - f"autogenerate detected {len(significant_diffs)} schema mismatch(es).\n\n" - f"Diffs:\n{diff_descriptions}" - ) - - # Step 4: Downgrade back to main's head to verify rollback works - command.downgrade(alembic_cfg, main_head) - - # Step 5: Verify the DB is actually at main's revision after downgrade - engine = create_engine(f"sqlite:///{db_path}") - try: - with engine.connect() as connection: - ctx = MigrationContext.configure(connection) - current_rev = ctx.get_current_revision() - assert current_rev == main_head, f"After downgrade, expected revision {main_head} but got {current_rev}" - finally: - engine.dispose() + with engine.connect() as connection: + migration_context = MigrationContext.configure(connection) + diffs = compare_metadata(migration_context, SQLModel.metadata) finally: - for suffix in ("", "-wal", "-shm", "-journal"): - Path(db_path + suffix).unlink(missing_ok=True) + engine.dispose() + + significant_diffs = _filter_diffs(diffs, db_url) + + if significant_diffs: + diff_descriptions = "\n".join(str(d) for d in significant_diffs) + pytest.fail( + f"After upgrading from main ({main_head}) to head, " + f"autogenerate detected {len(significant_diffs)} schema mismatch(es).\n\n" + f"Diffs:\n{diff_descriptions}" + ) + + # Step 4: Downgrade back to main's head to verify rollback works + command.downgrade(alembic_cfg, main_head) + + # Step 5: Verify the DB is actually at main's revision after downgrade + engine = create_engine(_engine_url(db_url)) + try: + with engine.connect() as connection: + ctx = MigrationContext.configure(connection) + current_rev = ctx.get_current_revision() + assert current_rev == main_head, f"After downgrade, expected revision {main_head} but got {current_rev}" + finally: + engine.dispose()