jason810496 commented on code in PR #63185:
URL: https://github.com/apache/airflow/pull/63185#discussion_r3510944307
##########
airflow-core/src/airflow/dag_processing/bundles/manager.py:
##########
@@ -366,6 +399,194 @@ def _extract_and_sign_template(bundle_name: str) ->
tuple[str | None, dict]:
session.execute(delete(ParseImportError).where(ParseImportError.bundle_name ==
name))
self.log.info("Deleted import errors for bundle %s which is no
longer configured", name)
+ # Airflow sets autoflush=False, so subsequent reads in the same session
+ # need an explicit flush to see ORM-side bundle state changes.
+ session.flush()
+
+ def reassign_dags_with_unconfigured_bundles(self) -> int:
+ """
+ Reassign Dags pointing at unconfigured bundles to a configured one.
+
+ Side effect of the ``0082_3_1_0_make_bundle_name_not_nullable``
+ migration (#63323): legacy rows get ``bundle_name='dags-folder'``
+ and NULL ``relative_fileloc``, which raises ``Requested bundle
+ '{name}' is not configured.`` at trigger time when the deployment
+ uses a custom bundle.
+
+ Each legacy-candidate row (NULL ``relative_fileloc`` and no
+ ``DagVersion``) is routed to the most-specific configured bundle
+ whose path contains its ``fileloc``, writing ``relative_fileloc``
+ atomically. Rows whose ``fileloc`` is under no configured bundle
+ are left untouched -- writing ``bundle_name`` without a verified
+ ``relative_fileloc`` would produce a row task workers cannot
+ execute -- and instead self-heal via the normal staleness
+ lifecycle: ``sync_bundles_to_db`` deactivates the old bundle, the
+ stale-scan marks the row stale, and the next successful parse
+ from any configured bundle resets everything via
+ ``update_dag_parsing_results_in_db``. No manual ``airflow dags
+ reserialize`` is required.
+
+ Multi-team-safe because a bundle path belongs to at most one team.
+ Each chunk runs in its own internally-owned transaction so the
+ row-lock window stays bounded and no caller-provided session is
+ committed.
+
+ :return: Number of Dags reassigned.
+ """
+ # Import here to avoid circular import
+ # (manager -> dag -> dagrun -> taskinstance -> dag_version -> manager)
+ from airflow.models.dag import DagModel
+ from airflow.models.dag_version import DagVersion
+
+ # Fast-skip once any 3.x parse cycle has run. DagVersion is written
+ # only by the parse path, and that path overwrites both bundle_name
+ # and relative_fileloc on every parse (see
+ # ``DagModelOperation.update_dags``), so any legacy row whose file
+ # is under a configured bundle self-heals at the next parse and any
+ # row whose file is under no configured bundle self-heals via the
+ # staleness lifecycle -- reassign has no work the parse path will
+ # not do itself. The probe is an index hit on dag_version's PK
+ # vs. a sequential scan of dag (no index on relative_fileloc).
+ with create_session() as session:
+ if session.scalar(select(DagVersion.id).limit(1)) is not None:
+ return 0
+
+ with create_session() as session:
+ if not (active_bundle_paths :=
self._resolve_active_bundle_paths(session=session)):
+ self.log.info(
+ "No active Dag bundles with resolvable paths; skipping
reassignment of Dags "
+ "with unconfigured bundles."
+ )
+ return 0
+
+ # Chunked UPDATEs ordered by dag_id, one transaction per chunk;
repaired
+ # rows drop out of the predicate because writing relative_fileloc
+ # makes the IS NULL clause false.
+ #
+ # Legacy-candidate predicate (rows never parsed in 3.x): NULL
+ # relative_fileloc (the 0082 migration leaves it NULL) AND NOT EXISTS
+ # DagVersion (the parse path writes DagModel.bundle_name before the
+ # DagVersion). Equivalent under that invariant; both stated as
+ # defense in depth and repeated on the UPDATE itself as a CAS guard
+ # so a concurrent parser write wins the race.
+ movements: dict[tuple[str | None, str], int] = defaultdict(int)
+ total_reassigned = 0
+ total_backfilled = 0
+ total_skipped = 0
+ last_dag_id: str | None = None
+
+ while True:
+ with create_session() as session:
+ query = (
+ select(DagModel.dag_id, DagModel.bundle_name,
DagModel.fileloc)
+ .where(
+ DagModel.relative_fileloc.is_(None),
+ ~exists().where(DagVersion.dag_id == DagModel.dag_id),
+ )
+ .order_by(DagModel.dag_id)
+ .limit(_REASSIGN_BATCH_SIZE)
+ )
+ if last_dag_id is not None:
+ query = query.where(DagModel.dag_id > last_dag_id)
+
+ if not (chunk := session.execute(query).all()):
+ break
+ last_dag_id = chunk[-1].dag_id
+
+ # Route every legacy row by fileloc, not just those on
+ # unconfigured bundles, so a migration-assigned dags-folder
+ # row whose file lives under a different configured bundle
+ # gets relocated instead of stranded. Classify as skip
+ # (no match), backfill (match == current bundle), or
+ # reassign (match != current bundle).
+ chunk_updates: list[tuple[str, str | None, str, str]] = []
+ for row in chunk:
+ match = _best_bundle_for_fileloc(row.fileloc,
active_bundle_paths) if row.fileloc else None
+ if match is None:
+ total_skipped += 1
+ continue
+ target, relative = match
+ chunk_updates.append((row.dag_id, row.bundle_name, target,
relative))
+
+ if not chunk_updates:
+ continue
+
+ with create_session() as session:
+ # create_session commits on context exit, bounding the
+ # row-lock window to one chunk.
+ for dag_id, prev_bundle, target, relative in chunk_updates:
+ result = cast(
+ "CursorResult",
+ session.execute(
+ update(DagModel)
+ .where(
+ DagModel.dag_id == dag_id,
+ DagModel.relative_fileloc.is_(None),
+ ~exists().where(DagVersion.dag_id ==
DagModel.dag_id),
+ )
+ .values(relative_fileloc=relative,
bundle_name=target)
+ .execution_options(synchronize_session=False)
+ ),
+ )
+
+ if result.rowcount:
+ # Rowcount is the source of truth for whether the CAS
actually fired
+ if target == prev_bundle:
+ total_backfilled += 1
+ else:
+ movements[(prev_bundle, target)] += 1
+ total_reassigned += 1
+ else:
+ self.log.debug("Skipping repair for Dag '%s': lost
race to parser.", dag_id)
+
+ for (prev, target), n in sorted(movements.items(), key=lambda item:
(str(item[0][0]), item[0][1])):
Review Comment:
Addressed in
https://github.com/apache/airflow/pull/63185/commits/5f5e118f627ffb6b33ec76ff0ed23bef2f9c4ce8
--
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]
