jason810496 commented on code in PR #63185:
URL: https://github.com/apache/airflow/pull/63185#discussion_r3658161524
##########
airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py:
##########
@@ -532,3 +546,1246 @@ def test_get_all_bundle_names():
# the naming suffix instead of pinning an exact list.
extra = [n for n in bundle_names if n not in {"dags-folder",
"example_dags"}]
assert all(n.endswith("-example-dags") for n in extra)
+
+
[email protected]
+def clear_dags_and_bundles():
+ clear_db_dags()
+ clear_db_dag_bundles()
+ yield
+ clear_db_dags()
+ clear_db_dag_bundles()
+
+
+def _add_dag(session, dag_id: str, bundle_name: str) -> DagModel:
+ dag = DagModel(dag_id=dag_id, bundle_name=bundle_name,
fileloc=f"/tmp/{dag_id}.py")
+ session.add(dag)
+ session.flush()
+ return dag
+
+
+class TestGuessBestBundleForFileloc:
+ """Tests for ``_guess_best_bundle_for_fileloc`` path matching."""
+
+ @pytest.mark.parametrize(
+ ("fileloc", "bundle_paths", "expected"),
+ [
+ pytest.param(
+ "/dags/team_x/dag.py",
+ {"team-x": Path("/dags/team_x")},
+ ("team-x", "dag.py"),
+ id="match",
+ ),
+ pytest.param("/elsewhere/dag.py", {"team-x":
Path("/dags/team_x")}, None, id="no_match"),
+ pytest.param("/dags/dag.py", {}, None, id="empty_paths"),
+ pytest.param(
+ # ``Path`` itself collapses ``//`` and ``.`` segments on
construction,
+ # so the helper matches without any explicit normalization.
+ "/dags//team_x/./dag.py",
+ {"team-x": Path("/dags/team_x")},
+ ("team-x", "dag.py"),
+ id="redundant_separators_and_dots",
+ ),
+ ],
+ )
+ def test_guess_best_bundle_for_fileloc(self, fileloc, bundle_paths,
expected) -> None:
+ assert _guess_best_bundle_for_fileloc(fileloc, bundle_paths) ==
expected
+
+
[email protected]_test
+class TestReassignDagsWithUnconfiguredBundles:
+ """Tests for DagBundlesManager.reassign_dags_with_unconfigured_bundles."""
+
+ def _manager_with_bundle_names(self, names: list[str]) ->
DagBundlesManager:
+ """Return a manager whose ``_resolve_active_bundle_paths`` reports
*names* with non-matching paths.
+
+ The fake paths are absolute roots that cannot contain any
+ ``/tmp/{dag_id}.py`` test fileloc, so rows are routed as unmatched
+ unless a test explicitly arranges otherwise.
+ """
+ with patch.dict(
+ os.environ,
+ {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST":
json.dumps(BASIC_BUNDLE_CONFIG)},
+ ):
+ manager = DagBundlesManager()
+ paths = {name: Path(f"/__unmatched_for_test__/{name}") for name in
names}
+ manager._resolve_active_bundle_paths = lambda *, session: paths #
type: ignore[method-assign]
+ return manager
+
+ def test_no_configured_bundles_is_noop(self, clear_dags_and_bundles,
session):
+ """Return 0 without raising when no bundles are configured."""
+ manager = self._manager_with_bundle_names([])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ def test_already_configured_is_noop(self, clear_dags_and_bundles, session)
-> None:
+ """No reassignment when every Dag already points at a configured
bundle."""
+ bundle = DagBundleModel(name="bundle-a")
+ bundle.active = True
+ session.add(bundle)
+ session.flush()
+ _add_dag(session, "dag-1", "bundle-a")
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["bundle-a"])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+ assert session.get(DagModel, "dag-1").bundle_name == "bundle-a"
+
+ def test_unmatched_fileloc_leaves_row_untouched(self,
clear_dags_and_bundles, session, caplog) -> None:
+ """Rows whose fileloc has no configured bundle path keep their
original bundle.
+
+ The fallback that wrote ``bundle_name`` without a verified
+ ``relative_fileloc`` produced active-but-un-runnable rows, so the
+ helper now leaves such rows on their unconfigured bundle and emits a
+ single warning naming the count.
+ """
+ active = DagBundleModel(name="active")
+ active.active = True
+ removed = DagBundleModel(name="removed-bundle")
+ removed.active = False
+ session.add(active)
+ session.add(removed)
+ session.flush()
+ _add_dag(session, "dag-1", "removed-bundle")
+ _add_dag(session, "dag-2", "removed-bundle")
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["active"])
+ with caplog.at_level("WARNING",
logger="airflow.dag_processing.bundles.manager.DagBundlesManager"):
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ session.expire_all()
+ for dag_id in ("dag-1", "dag-2"):
+ row = session.get(DagModel, dag_id)
+ assert row.bundle_name == "removed-bundle"
+ assert row.relative_fileloc is None
+ assert any("Skipped 2 legacy Dag(s)" in record.message for record in
caplog.records)
+
+ def test_row_with_populated_relative_fileloc_is_left_alone(self,
clear_dags_and_bundles, session) -> None:
+ """A 3.x row whose bundle is no longer configured must keep its bundle
assignment.
+
+ Only rows the 0082 migration touched (``relative_fileloc IS NULL``) are
+ candidates for reassignment. Rows that already carry a relative path
+ were written by 3.x serialization and must be left on their bundle so
+ the regular stale-Dag deactivation path can handle a removed bundle.
+ """
+ active_bundle = DagBundleModel(name="active")
+ active_bundle.active = True
+ removed_bundle = DagBundleModel(name="removed-bundle")
+ removed_bundle.active = False
+ session.add(active_bundle)
+ session.add(removed_bundle)
+ session.flush()
+
+ dag = DagModel(
+ dag_id="dag-1",
+ bundle_name="removed-bundle",
+ fileloc="/tmp/dag-1.py",
+ )
+ dag.relative_fileloc = "dag-1.py"
+ dag.bundle_version = "abc123"
+ session.add(dag)
+ session.flush()
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["active"])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ session.expire_all()
+ refreshed = session.get(DagModel, "dag-1")
+ assert refreshed.bundle_name == "removed-bundle"
+ assert refreshed.relative_fileloc == "dag-1.py"
+ assert refreshed.bundle_version == "abc123"
+
+ def test_row_with_existing_dag_version_is_left_alone(self,
clear_dags_and_bundles, session) -> None:
+ """A Dag with any DagVersion row must not be reassigned.
+
+ Defended at two layers: (1) the global DagVersion-existence fast-skip
+ short-circuits before any work, and (2) the per-row predicate
+ ``NOT EXISTS DagVersion`` excludes the row even if the fast-skip is
+ bypassed. The parse path is the source of truth for any Dag with a
+ DagVersion -- touching only the DagModel would leave the DagVersion
+ stale, and scheduler/executor paths prefer DagVersion.bundle_name
+ when building task workloads.
+ """
+ active = DagBundleModel(name="active")
+ active.active = True
+ removed = DagBundleModel(name="removed-bundle")
+ removed.active = False
+ session.add(active)
+ session.add(removed)
+ session.flush()
+
+ # NULL relative_fileloc would normally make this a repair candidate;
+ # the DagVersion row should exclude it from the predicate.
+ dag = DagModel(
+ dag_id="versioned",
+ bundle_name="removed-bundle",
+ fileloc="/tmp/versioned.py",
+ )
+ dag.relative_fileloc = None
+ session.add(dag)
+ session.flush()
+
+ version = DagVersion(
+ dag_id="versioned",
+ version_number=1,
+ bundle_name="removed-bundle",
+ bundle_version="v1",
+ )
+ session.add(version)
+ session.flush()
+ session.commit()
+
+ manager = self._manager_with_bundle_names(["active"])
+ assert manager.reassign_dags_with_unconfigured_bundles() == 0
+
+ session.expire_all()
+ refreshed = session.get(DagModel, "versioned")
+ assert refreshed.bundle_name == "removed-bundle"
+ assert refreshed.relative_fileloc is None
+ refreshed_version = session.get(DagVersion, version.id)
+ assert refreshed_version.bundle_name == "removed-bundle"
+ assert refreshed_version.bundle_version == "v1"
+
+ @conf_vars({("core", "multi_team"): "True"})
+ def test_runs_under_multi_team_mode(self, clear_dags_and_bundles, session,
tmp_path) -> None:
+ """Multi-team mode still repairs legacy rows.
+
+ Each team's bundle owns a distinct on-disk path, so routing a legacy
+ row to the most-specific bundle whose path contains its ``fileloc``
+ cannot cross a team boundary. The repair therefore runs unchanged
+ under ``core.multi_team`` and must not blanket-skip rows that have a
+ safe target.
+ """
+ bundle_dir = tmp_path / "team_x"
+ bundle_dir.mkdir()
+ legacy_file = bundle_dir / "legacy.py"
+ legacy_file.write_text("# legacy dag")
+
+ config = [
+ {
+ "name": "team-x-bundle",
+ "classpath":
"airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {"path": str(bundle_dir), "refresh_interval": 1},
+ }
+ ]
+ active_bundle = DagBundleModel(name="team-x-bundle")
+ active_bundle.active = True
+ removed_bundle = DagBundleModel(name="removed-bundle")
+ removed_bundle.active = False
+ session.add(active_bundle)
+ session.add(removed_bundle)
+ session.flush()
+
+ dag = DagModel(dag_id="legacy", bundle_name="removed-bundle",
fileloc=str(legacy_file))
+ dag.relative_fileloc = None
+ session.add(dag)
+ session.flush()
+ session.commit()
+
+ with patch.dict(
+ os.environ,
+ {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST":
json.dumps(config)},
+ ):
+ count =
DagBundlesManager().reassign_dags_with_unconfigured_bundles()
+
+ assert count == 1
+ session.expire_all()
+ refreshed = session.get(DagModel, "legacy")
+ assert refreshed.bundle_name == "team-x-bundle"
+ assert refreshed.relative_fileloc == "legacy.py"
+
+
[email protected]_test
+class TestBackfillRelativeFileloc:
+ """Tests for the legacy ``relative_fileloc`` backfill triggered by
reassignment."""
+
+ @pytest.mark.parametrize(
+ ("fileloc_under_bundle", "expected_bundle_name",
"expected_relative_fileloc"),
+ [
+ pytest.param(True, "my-bundle", "legacy.py",
id="fileloc_under_bundle_path_is_backfilled"),
+ pytest.param(False, "orphan-bundle", None,
id="fileloc_outside_bundle_path_is_left_alone"),
+ ],
+ )
+ def test_backfill_behavior(
+ self,
+ clear_dags_and_bundles,
+ session,
+ tmp_path,
+ fileloc_under_bundle: bool,
+ expected_bundle_name: str,
+ expected_relative_fileloc: str | None,
+ ) -> None:
+ """Reassignment only happens when ``fileloc`` lies under a configured
bundle path.
+
+ When the fileloc matches, both ``bundle_name`` and ``relative_fileloc``
+ are written atomically. When it does not match, the row is left on its
+ unconfigured bundle so it cannot become an active-but-un-runnable row.
+
+ :param fileloc_under_bundle: Whether the legacy Dag's absolute
``fileloc`` lies under
+ the configured bundle's path.
+ :param expected_bundle_name: Expected ``bundle_name`` after
reassignment.
+ :param expected_relative_fileloc: Expected ``relative_fileloc`` after
reassignment.
+ """
+ bundle_dir = tmp_path / "dags"
+ bundle_dir.mkdir()
+ legacy_file_under_bundle = bundle_dir / "legacy.py"
+ legacy_file_under_bundle.write_text("# legacy dag")
+ legacy_fileloc = str(legacy_file_under_bundle) if fileloc_under_bundle
else "/elsewhere/foo.py"
+
+ config = [
+ {
+ "name": "my-bundle",
+ "classpath":
"airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {"path": str(bundle_dir), "refresh_interval": 1},
+ }
+ ]
+ active_bundle = DagBundleModel(name="my-bundle")
+ active_bundle.active = True
+ orphan_bundle = DagBundleModel(name="orphan-bundle")
+ orphan_bundle.active = False
+ session.add(active_bundle)
+ session.add(orphan_bundle)
+ session.flush()
+
+ dag = DagModel(dag_id="legacy", bundle_name="orphan-bundle",
fileloc=legacy_fileloc)
+ dag.relative_fileloc = None
+ session.add(dag)
+ session.flush()
+ session.commit()
+
+ with patch.dict(
+ os.environ,
+ {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST":
json.dumps(config)},
+ ):
+ manager = DagBundlesManager()
+ manager.reassign_dags_with_unconfigured_bundles()
+
+ session.expire_all()
+ refreshed = session.get(DagModel, "legacy")
+ assert refreshed.bundle_name == expected_bundle_name
+ assert refreshed.relative_fileloc == expected_relative_fileloc
+
+ def test_post_2x_to_3x_migration_with_renamed_bundle(
+ self, clear_dags_and_bundles, session, tmp_path
+ ) -> None:
+ """End-to-end: 2.x→3.x upgrade where the operator's bundle is not
named ``dags-folder``.
+
+ Reproduces the original incident: the migration sets
``bundle_name='dags-folder'`` on
+ legacy Dag rows and leaves ``relative_fileloc`` NULL, but the operator
has configured a
+ single LocalDagBundle named ``custom-bundle`` pointing at the same
on-disk dags folder.
+ After ``reassign_dags_with_unconfigured_bundles`` runs at DFP startup,
the legacy Dags
+ should be:
+
+ 1. Reassigned to ``custom-bundle`` so triggering DagRuns no longer
fails with
+ "Requested bundle 'dags-folder' is not configured."
+ 2. Have ``relative_fileloc`` backfilled from ``fileloc`` so the
standard fileloc-based
+ stale-detection path can later detect real deletions.
+ """
+ dags_folder = tmp_path / "dags"
+ dags_folder.mkdir()
+ legacy_file = dags_folder / "my_dag.py"
+ legacy_file.write_text("# 2.x dag")
+
+ config = [
+ {
+ "name": "custom-bundle",
+ "classpath":
"airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {"path": str(dags_folder), "refresh_interval": 1},
+ }
+ ]
+ # State right after the 0082 migration: ``dags-folder`` row exists in
dag_bundle (added
+ # by migration backfill) but is not in the user's config; the user's
``custom-bundle``
+ # has not been registered yet (sync_bundles_to_db does that on
startup).
+ legacy_default_bundle = DagBundleModel(name="dags-folder")
+ legacy_default_bundle.active = True
+ session.add(legacy_default_bundle)
+ session.flush()
+
+ legacy_dag = DagModel(
+ dag_id="legacy_dag",
+ bundle_name="dags-folder",
+ fileloc=str(legacy_file),
+ )
+ legacy_dag.relative_fileloc = None
+ session.add(legacy_dag)
+ session.flush()
+ session.commit()
+
+ with patch.dict(
+ os.environ,
+ {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST":
json.dumps(config)},
+ ):
+ manager = DagBundlesManager()
+ # DFP startup order: register configured bundles, then reassign.
+ manager.sync_bundles_to_db(session=session)
+ session.commit()
+ count = manager.reassign_dags_with_unconfigured_bundles()
+
+ assert count == 1
+ session.expire_all()
+ refreshed = session.get(DagModel, "legacy_dag")
+ assert refreshed.bundle_name == "custom-bundle"
+ assert refreshed.relative_fileloc == "my_dag.py"
+
+ def test_backfill_commits_between_chunks(
+ self, clear_dags_and_bundles, session, tmp_path, monkeypatch
+ ) -> None:
+ """The legacy backfill chunks and commits like the reassignment loop.
+
+ Without chunked commits, a deployment where every legacy Dag is
+ already on a configured bundle would still UPDATE the whole set in
+ one transaction, holding row locks on the dag table for the full
+ DFP startup.
+ """
+ from airflow.dag_processing.bundles import manager as
bundles_manager_mod
+
+ bundle_dir = tmp_path / "dags"
+ bundle_dir.mkdir()
+ for i in range(5):
+ (bundle_dir / f"legacy_{i}.py").write_text(f"# legacy {i}")
+
+ config = [
+ {
+ "name": "my-bundle",
+ "classpath":
"airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {"path": str(bundle_dir), "refresh_interval": 1},
+ }
+ ]
+ active_bundle = DagBundleModel(name="my-bundle")
+ active_bundle.active = True
+ session.add(active_bundle)
+ session.flush()
+
+ for i in range(5):
+ dag = DagModel(
+ dag_id=f"legacy_{i}",
+ bundle_name="my-bundle",
+ fileloc=str(bundle_dir / f"legacy_{i}.py"),
+ )
+ dag.relative_fileloc = None
+ session.add(dag)
+ session.flush()
+ session.commit()
+
+ monkeypatch.setattr(bundles_manager_mod, "_REASSIGN_BATCH_SIZE", 2)
+ # Each chunk opens its own ``create_session`` (which commits on exit),
+ # so counting context-manager entries equals counting batch commits.
+ session_open_count = [0]
+ real_create_session = bundles_manager_mod.create_session
+
+ @contextlib.contextmanager
+ def _counting_create_session(*args, **kwargs):
+ session_open_count[0] += 1
+ with real_create_session(*args, **kwargs) as s:
+ yield s
+
+ monkeypatch.setattr(bundles_manager_mod, "create_session",
_counting_create_session)
+
+ with patch.dict(
+ os.environ,
+ {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST":
json.dumps(config)},
+ ):
+ manager = DagBundlesManager()
+ # No unconfigured rows exist, so the reassignment loop returns
+ # immediately and the helper drives the chunking we are testing.
+ manager.reassign_dags_with_unconfigured_bundles()
+
+ # 1 session for the active-bundle read + 5 backfill rows / batch size 2
+ # = chunks of 2, 2, 1, then an empty chunk that terminates the loop.
+ assert session_open_count[0] >= 4
Review Comment:
Addressed in
https://github.com/apache/airflow/commit/2472a3260fed8fe2869c91ca60920bcc622c51fc,
I added exact count with explaination.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
