deer-flow/backend/tests/test_persistence_migrations_env.py

"""Tests for the ``include_object`` filter used by ``migrations/env.py``.

LangGraph checkpointer tables (``checkpoints`` and friends) live alongside
DeerFlow's own tables in the same database. Alembic must NEVER emit DDL for
them or a future ``alembic revision --autogenerate`` would propose
``drop_table('checkpoints')`` whenever LangGraph's tables are reflected from
a live DB.

The filter is the only line of defence between an honest autogenerate run
and a destructive revision. It lives in ``_env_filters.py`` so it can be unit
tested without alembic's import-time machinery.
"""

from __future__ import annotations

import sqlalchemy as sa

from deerflow.persistence.migrations._env_filters import (
    LANGGRAPH_OWNED_TABLES,
    include_object,
)


def _table(name: str) -> sa.Table:
    return sa.Table(name, sa.MetaData())


def test_filter_excludes_langgraph_checkpoint_tables() -> None:
    for owned in (
        "checkpoints",
        "checkpoint_blobs",
        "checkpoint_writes",
        "checkpoint_migrations",
    ):
        assert include_object(_table(owned), owned, "table", True, None) is False


def test_filter_includes_deerflow_tables() -> None:
    for owned in ("runs", "threads_meta", "feedback", "users", "channel_connections"):
        assert include_object(_table(owned), owned, "table", True, None) is True


def test_filter_excludes_indexes_on_langgraph_tables() -> None:
    # An Index whose parent table is LangGraph-owned must also be filtered out;
    # otherwise autogenerate would emit drop_index against tables alembic does
    # not own.
    md = sa.MetaData()
    parent = sa.Table("checkpoints", md, sa.Column("id", sa.Integer, primary_key=True))
    idx = sa.Index("ix_checkpoints_anything", parent.c.id)
    assert include_object(idx, idx.name, "index", True, None) is False


def test_filter_includes_indexes_on_deerflow_tables() -> None:
    md = sa.MetaData()
    parent = sa.Table("runs", md, sa.Column("run_id", sa.String, primary_key=True))
    idx = sa.Index("ix_runs_something", parent.c.run_id)
    assert include_object(idx, idx.name, "index", True, None) is True


def test_langgraph_owned_tables_set_is_complete() -> None:
    # Pin the explicit set so an inadvertent removal -- e.g. someone simplifying
    # the filter -- requires a test diff that surfaces the change.
    assert LANGGRAPH_OWNED_TABLES == frozenset(
        {
            "checkpoints",
            "checkpoint_blobs",
            "checkpoint_writes",
            "checkpoint_migrations",
        }
    )


def test_env_module_wires_busy_timeout_for_sqlite() -> None:
    """Regression for the cross-process bootstrap pitfall: alembic spawns its
    own engine inside ``env.py::run_migrations_online`` and that engine does
    NOT inherit PRAGMAs from the production engine. Without an event listener
    here, its connections would use the default 5s busy_timeout and racy
    multi-process bootstrap would fail with ``database is locked`` instead of
    waiting for the file lock.

    We check the source rather than execute env.py (which would try to drive
    alembic on import) so this test stays a pure parity check.
    """
    from pathlib import Path  # noqa: PLC0415

    env_path = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/persistence/migrations/env.py"
    src = env_path.read_text(encoding="utf-8")
    assert "PRAGMA busy_timeout=30000" in src or "PRAGMA busy_timeout = 30000" in src, (
        "env.py must set busy_timeout on its alembic-spawned engine; without it, cross-process bootstrap on SQLite fails fast instead of waiting for the file lock"
    )
    assert 'listens_for(connectable.sync_engine, "connect")' in src, "busy_timeout must be wired via an event listener so EVERY connection alembic opens gets the PRAGMA, not just one initial probe"


class TestExtensionOwnedTables:
    """Extension tables share the database but not alembic's view of it.

    An extension owns its own MetaData and its own migration chain, so
    autogenerate reflecting them from a live database would find them absent
    from Base.metadata and propose dropping them.

    Note the scope: `make migrate-rev` is already safe, because
    `_autogen_revision.py` diffs against a throwaway SQLite built from the
    migration chain, where no extension table exists. The exposed path is a
    direct `alembic revision --autogenerate` from the migrations directory,
    whose `alembic.ini` points at a real `./data/deerflow.db` — the same path
    `LANGGRAPH_OWNED_TABLES` covers.
    """

    def setup_method(self):
        from deerflow.persistence.migrations import _env_filters

        self._saved = set(_env_filters.EXTENSION_TABLE_PREFIXES)

    def teardown_method(self):
        from deerflow.persistence.migrations import _env_filters

        _env_filters.EXTENSION_TABLE_PREFIXES.clear()
        _env_filters.EXTENSION_TABLE_PREFIXES.update(self._saved)

    def test_a_registered_prefix_is_excluded(self):
        from deerflow.persistence.migrations._env_filters import include_object, register_extension_table_prefix

        register_extension_table_prefix("ext_")
        assert include_object(None, "ext_events", "table", True, None) is False

    def test_an_unregistered_table_is_still_included(self):
        from deerflow.persistence.migrations._env_filters import include_object

        assert include_object(None, "runs", "table", True, None) is True

    def test_an_index_on_an_excluded_table_is_excluded_too(self):
        from types import SimpleNamespace

        from deerflow.persistence.migrations._env_filters import include_object, register_extension_table_prefix

        register_extension_table_prefix("ext_")
        index = SimpleNamespace(table=SimpleNamespace(name="ext_events"))
        assert include_object(index, "ix_ext_events_seq", "index", True, None) is False

    def test_a_constraint_on_an_excluded_table_is_excluded_too(self):
        # A filter that drops the table but keeps its constraints still emits
        # broken DDL (e.g. a dangling unique_constraint against a table
        # alembic no longer believes exists).
        from types import SimpleNamespace

        from deerflow.persistence.migrations._env_filters import include_object, register_extension_table_prefix

        register_extension_table_prefix("ext_")
        constraint = SimpleNamespace(table=SimpleNamespace(name="ext_events"))
        assert include_object(constraint, "uq_ext_events_seq", "unique_constraint", True, None) is False

    def test_registration_rejects_an_empty_prefix(self):
        import pytest

        from deerflow.persistence.migrations._env_filters import register_extension_table_prefix

        with pytest.raises(ValueError):
            register_extension_table_prefix("")

    def test_langgraph_exclusion_is_unaffected(self):
        from deerflow.persistence.migrations._env_filters import include_object

        assert include_object(None, "checkpoints", "table", True, None) is False

    def test_registration_rejects_a_prefix_that_would_hide_a_host_table(self):
        """A typo like table_prefix: "run" would silently stop alembic from
        managing the host's own `runs` / `run_events` tables. This must fail
        loudly at registration time rather than degrade autogenerate silently."""
        import pytest

        from deerflow.persistence.migrations._env_filters import register_extension_table_prefix

        with pytest.raises(ValueError, match="runs"):
            register_extension_table_prefix("run")

    def test_registration_accepts_a_prefix_that_matches_no_host_table(self):
        from deerflow.persistence.migrations._env_filters import EXTENSION_TABLE_PREFIXES, register_extension_table_prefix

        register_extension_table_prefix("acme_ext_")
        assert "acme_ext_" in EXTENSION_TABLE_PREFIXES


_PREFIX_PROBE = """
import json, sys

from alembic.autogenerate import compare_metadata
from alembic.migration import MigrationContext
from sqlalchemy import Column, Integer, MetaData, String, Table, create_engine

import deerflow.persistence.models  # noqa: F401 - populate Base.metadata
from deerflow.persistence.base import Base
from deerflow.persistence.migrations._env_filters import (
    include_object,
    register_configured_extension_table_prefixes,
)

db_path = sys.argv[1]

# An extension-owned table in the shape an extension's own chain leaves it:
# present in the database, absent from the host's Base.metadata.
foreign = MetaData()
Table("acme_events", foreign, Column("id", Integer, primary_key=True), Column("body", String(16)))
engine = create_engine("sqlite:///" + db_path)
foreign.create_all(engine)

registered = register_configured_extension_table_prefixes()

with engine.connect() as conn:
    ctx = MigrationContext.configure(conn, opts={"include_object": include_object, "compare_type": False})
    diffs = compare_metadata(ctx, Base.metadata)

dropped = [op[1].name for op in diffs if isinstance(op, tuple) and op and op[0] == "remove_table"]
print(json.dumps({"registered": list(registered), "dropped": dropped}))
"""


class TestPrefixesReachTheAlembicProcess:
    """The filter is only as good as whatever populates it.

    Every other test here registers a prefix in-process and then asserts
    ``include_object``. Those pass for a real reason — the filter works — and
    still missed that the sole caller of ``register_extension_table_prefix``
    was ``load_extensions()``, which runs in the Gateway and never in the
    alembic process, leaving the set empty exactly where alembic reads it.
    These cover the wiring rather than the filter.
    """

    def _probe(self, tmp_path, config_yaml: str) -> dict:
        import json
        import os
        import subprocess
        import sys
        from pathlib import Path

        (tmp_path / "config.yaml").write_text(config_yaml, encoding="utf-8")
        script = tmp_path / "probe.py"
        script.write_text(_PREFIX_PROBE, encoding="utf-8")

        backend = Path(__file__).resolve().parents[1]
        proc = subprocess.run(
            [sys.executable, str(script), str(tmp_path / "probe.db")],
            cwd=str(backend),
            env={**os.environ, "DEER_FLOW_CONFIG_PATH": str(tmp_path / "config.yaml"), "PYTHONPATH": str(backend)},
            capture_output=True,
            text=True,
            timeout=180,
        )
        assert proc.returncode == 0, f"probe failed:\n{proc.stdout}\n{proc.stderr}"
        return json.loads(proc.stdout.strip().splitlines()[-1])

    def test_a_declared_prefix_survives_into_a_fresh_process(self, tmp_path):
        """No Gateway has run here, so nothing called load_extensions()."""
        result = self._probe(tmp_path, "plugins:\n  - use: acme_ext:install\n    table_prefix: acme_\n")

        assert result["registered"] == ["acme_"]
        assert "acme_events" not in result["dropped"], "autogenerate proposed dropping an extension-owned table"

    def test_an_undeclared_table_is_still_proposed_for_dropping(self, tmp_path):
        """The control. Without it the test above could pass on a filter that excludes everything."""
        result = self._probe(tmp_path, "plugins: []\n")

        assert result["registered"] == []
        assert "acme_events" in result["dropped"], "expected an undeclared foreign table to be reflected and dropped"

    def test_an_empty_prefix_does_not_take_the_alembic_process_down(self, tmp_path):
        """One declaration must not mean two different things in two processes.

        This reader parses raw YAML so it never imports the extension, which
        also means ``ExtensionSpec`` never validates what it sees. Rejecting an
        empty prefix here would surface a config-schema error from the wrong
        process twice over: an operator does not expect to hear about a
        malformed ``config.yaml`` from alembic, and Gateway startup runs this
        same module through ``bootstrap_schema`` — so a raise would take the
        Gateway down with a message about migrations. ``ExtensionSpec`` owns
        that verdict (see the paired test below).
        """
        result = self._probe(tmp_path, 'plugins:\n  - use: acme_ext:install\n    table_prefix: ""\n')

        assert result["registered"] == []
        # And it degrades to the no-prefix behaviour rather than to
        # "" matching every table name.
        assert "acme_events" in result["dropped"]

    def test_a_non_string_prefix_does_not_take_the_alembic_process_down(self, tmp_path):
        result = self._probe(tmp_path, "plugins:\n  - use: acme_ext:install\n    table_prefix: 7\n")

        assert result["registered"] == []

    def test_env_module_registers_configured_prefixes(self):
        """``env.py`` cannot be imported outside alembic, so pin the call in its source."""
        import ast
        from pathlib import Path

        env_py = Path(__file__).resolve().parents[1] / "packages/harness/deerflow/persistence/migrations/env.py"
        called = {node.func.id for node in ast.walk(ast.parse(env_py.read_text(encoding="utf-8"))) if isinstance(node, ast.Call) and isinstance(node.func, ast.Name)}

        assert "register_configured_extension_table_prefixes" in called, "env.py must populate the prefix set; include_object reads it in that process"