Lee-W commented on code in PR #63185:
URL: https://github.com/apache/airflow/pull/63185#discussion_r3535240214
##########
airflow-core/src/airflow/dag_processing/bundles/manager.py:
##########
@@ -34,17 +36,48 @@
from airflow.models.team import Team
from airflow.providers_manager import ProvidersManager
from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.session import NEW_SESSION, create_session, provide_session
if TYPE_CHECKING:
from collections.abc import Iterable
+ from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session
log = logging.getLogger(__name__)
_example_dag_bundle_name = "example_dags"
+# Chunk size for the one-time startup repair of unconfigured bundles.
+_REASSIGN_BATCH_SIZE = 1000
+
+
+def _best_bundle_for_fileloc(
Review Comment:
```suggestion
def _guest_best_bundle_for_fileloc(
```
nit
##########
airflow-core/src/airflow/dag_processing/bundles/manager.py:
##########
@@ -34,17 +36,48 @@
from airflow.models.team import Team
from airflow.providers_manager import ProvidersManager
from airflow.utils.log.logging_mixin import LoggingMixin
-from airflow.utils.session import NEW_SESSION, provide_session
+from airflow.utils.session import NEW_SESSION, create_session, provide_session
if TYPE_CHECKING:
from collections.abc import Iterable
+ from sqlalchemy.engine import CursorResult
from sqlalchemy.orm import Session
log = logging.getLogger(__name__)
_example_dag_bundle_name = "example_dags"
+# Chunk size for the one-time startup repair of unconfigured bundles.
+_REASSIGN_BATCH_SIZE = 1000
+
+
+def _best_bundle_for_fileloc(
+ fileloc: str, descending_bundle_paths: dict[str, Path]
+) -> tuple[str, str] | None:
+ """
+ Return ``(bundle_name, relative_fileloc)`` for the first bundle whose path
contains ``fileloc``.
+
+ ``descending_bundle_paths`` must be sorted by path length descending so
+ the deepest bundle wins when paths overlap.
+
+ Returns ``None`` when ``fileloc`` is not under any bundle's path.
+
+ Uses the same plain ``Path.relative_to`` check as
+ ``BaseDagImporter.get_relative_path``, so the ``relative_fileloc`` written
+ here matches what the next parse computes for the same file. Filelocs are
+ produced by the Dag processor parsing admin-controlled bundle files, so
they
+ are trusted and need no path-traversal normalization.
+ """
+ file_path = Path(fileloc)
+ for name, path in descending_bundle_paths.items():
+ try:
+ relative = file_path.relative_to(path)
+ except ValueError:
+ continue
Review Comment:
```suggestion
with contextlib.suppress():
relative = file_path.relative_to(path)
```
##########
airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py:
##########
@@ -476,3 +489,1244 @@ 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 TestBestBundleForFileloc:
+ """Tests for ``_best_bundle_for_fileloc`` path matching."""
+
+ def test_returns_relative_path_for_match(self) -> None:
+ assert _best_bundle_for_fileloc("/dags/team_x/dag.py", {"team-x":
Path("/dags/team_x")}) == (
+ "team-x",
+ "dag.py",
+ )
+
+ def test_returns_none_for_no_match(self) -> None:
+ assert _best_bundle_for_fileloc("/elsewhere/dag.py", {"team-x":
Path("/dags/team_x")}) is None
+
+ def test_returns_none_for_empty_paths(self) -> None:
+ assert _best_bundle_for_fileloc("/dags/dag.py", {}) is None
+
+ def test_normalises_redundant_separators_and_dots(self) -> None:
+ # ``Path`` itself collapses ``//`` and ``.`` segments on construction,
+ # so the helper matches without any explicit normalization.
+ assert _best_bundle_for_fileloc("/dags//team_x/./dag.py", {"team-x":
Path("/dags/team_x")}) == (
+ "team-x",
+ "dag.py",
+ )
Review Comment:
I think we can parameterize these
##########
airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py:
##########
@@ -476,3 +489,1244 @@ 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 TestBestBundleForFileloc:
+ """Tests for ``_best_bundle_for_fileloc`` path matching."""
+
+ def test_returns_relative_path_for_match(self) -> None:
+ assert _best_bundle_for_fileloc("/dags/team_x/dag.py", {"team-x":
Path("/dags/team_x")}) == (
+ "team-x",
+ "dag.py",
+ )
+
+ def test_returns_none_for_no_match(self) -> None:
+ assert _best_bundle_for_fileloc("/elsewhere/dag.py", {"team-x":
Path("/dags/team_x")}) is None
+
+ def test_returns_none_for_empty_paths(self) -> None:
+ assert _best_bundle_for_fileloc("/dags/dag.py", {}) is None
+
+ def test_normalises_redundant_separators_and_dots(self) -> None:
+ # ``Path`` itself collapses ``//`` and ``.`` segments on construction,
+ # so the helper matches without any explicit normalization.
+ assert _best_bundle_for_fileloc("/dags//team_x/./dag.py", {"team-x":
Path("/dags/team_x")}) == (
+ "team-x",
+ "dag.py",
+ )
+
+
[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.
+ """
+ from airflow.models.dag_version import DagVersion
Review Comment:
do we need this inline import here
--
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]
