This is an automated email from the ASF dual-hosted git repository.
Lee-W pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4138878208c Use EXISTS instead of a fan-out join in asset orphanage
(#75404)
4138878208c is described below
commit 4138878208cc61e2fea420c4eda9a41a744a5ea3
Author: Aboulfoul <[email protected]>
AuthorDate: Mon Oct 26 10:43:34 2026 +0100
Use EXISTS instead of a fan-out join in asset orphanage (#75404)
Co-authored-by: Wei Lee <[email protected]>
---
.../src/airflow/jobs/scheduler_job_runner.py | 43 +++++----------
airflow-core/tests/unit/jobs/test_scheduler_job.py | 61 ++++++++++++++++++++++
2 files changed, 75 insertions(+), 29 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 82ea45fcb03..01c80d53379 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -35,7 +35,6 @@ from typing import TYPE_CHECKING, Any, cast
from uuid import UUID
from sqlalchemy import (
- CTE,
Text,
and_,
case,
@@ -2025,6 +2024,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
timers.call_regular_interval(
conf.getfloat("scheduler", "parsing_cleanup_interval"),
self._update_asset_orphanage,
+ non_fatal=True,
)
timers.call_regular_interval(
conf.getfloat("scheduler", "parsing_cleanup_interval"),
@@ -4219,37 +4219,24 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
:seealso: :meth:`AssetModelOperation.activate_assets_if_possible`.
"""
- # Group assets into orphaned=True and orphaned=False groups.
- orphaned = (
- (
- func.count(DagScheduleAssetReference.dag_id)
- + func.count(TaskOutletAssetReference.dag_id)
- + func.count(TaskInletAssetReference.dag_id)
- # An asset materialized from an ``AssetAlias`` outlet is
referenced only through
- # the alias association, not the schedule/outlet/inlet tables.
Count it so it is
- # activated (and shown in the UI) instead of being treated as
orphaned. See #58058.
- + func.count(alias_association_table.c.alias_id)
- )
- == 0
- ).label("orphaned")
- asset_reference_query = (
- select(AssetModel)
- .outerjoin(DagScheduleAssetReference)
- .outerjoin(TaskOutletAssetReference)
- .outerjoin(TaskInletAssetReference)
- .outerjoin(alias_association_table,
alias_association_table.c.asset_id == AssetModel.id)
- .group_by(AssetModel.id)
+ referenced = or_(
+ exists().where(DagScheduleAssetReference.asset_id ==
AssetModel.id),
+ exists().where(TaskOutletAssetReference.asset_id == AssetModel.id),
+ exists().where(TaskInletAssetReference.asset_id == AssetModel.id),
+ # An asset materialized from an ``AssetAlias`` outlet is
referenced only through
+ # the alias association, not the schedule/outlet/inlet tables.
Treat it as referenced so it
+ # is activated (and shown in the UI) instead of being orphaned.
See #58058.
+ exists().where(alias_association_table.c.asset_id ==
AssetModel.id),
)
-
- orphan_query = asset_reference_query.having(orphaned).cte()
- activate_query = asset_reference_query.having(~orphaned).cte()
+ orphan_query = select(AssetModel.name,
AssetModel.uri).where(~referenced).subquery()
+ activate_query = select(AssetModel.name,
AssetModel.uri).where(referenced).subquery()
self._orphan_unreferenced_assets(orphan_query, session=session)
self._activate_referenced_assets(activate_query, session=session)
self._cleanup_orphaned_asset_state_store(session=session)
@staticmethod
- def _orphan_unreferenced_assets(assets_query: CTE, *, session: Session) ->
None:
+ def _orphan_unreferenced_assets(assets_query: Subquery, *, session:
Session) -> None:
deleted_orphaned_assets = session.execute(
delete(AssetActive).where(
exists().where(
@@ -4261,7 +4248,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
stats.gauge("asset.orphaned", max(getattr(deleted_orphaned_assets,
"rowcount", 0), 0))
@staticmethod
- def _activate_referenced_assets(assets_query: CTE, *, session: Session) ->
None:
+ def _activate_referenced_assets(assets_query: Subquery, *, session:
Session) -> None:
active_assets_query = select(AssetActive.name, AssetActive.uri).join(
assets_query,
and_(AssetActive.name == assets_query.c.name, AssetActive.uri ==
assets_query.c.uri),
@@ -4304,9 +4291,7 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
),
)
.where(
- ~active_assets_query.where(
- and_(AssetActive.name == AssetModel.name,
AssetActive.uri == AssetModel.uri)
- ).exists()
+ ~exists().where(AssetActive.name == AssetModel.name,
AssetActive.uri == AssetModel.uri)
)
)
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index 13399201cf3..aa2a9c7c6d3 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -10262,6 +10262,67 @@ class TestSchedulerJob:
assert asset in active
assert asset not in orphaned
+ def test_asset_orphaning_asset_referenced_from_every_table(self,
dag_maker, session):
+ self.job_runner = SchedulerJobRunner(job=Job())
+
+ shared = Asset(uri="test://shared", name="shared", group="test_group")
+ dropped = Asset(uri="test://dropped", name="dropped",
group="test_group")
+ schedule_only = Asset(uri="test://schedule_only",
name="schedule_only", group="test_group")
+ outlet_only = Asset(uri="test://outlet_only", name="outlet_only",
group="test_group")
+ inlet_only = Asset(uri="test://inlet_only", name="inlet_only",
group="test_group")
+
+ for i in range(3):
+ with dag_maker(dag_id=f"consumer-{i}", schedule=[shared],
session=session):
+ BashOperator(task_id="task", bash_command="echo 1")
+ with dag_maker(dag_id="schedule-only", schedule=[schedule_only],
session=session):
+ BashOperator(task_id="task", bash_command="echo 1")
+ with dag_maker(dag_id="producer", schedule=None, session=session):
+ for i in range(3):
+ BashOperator(task_id=f"produce-{i}", bash_command="echo 1",
outlets=[shared, dropped])
+ BashOperator(task_id=f"read-{i}", bash_command="echo 1",
inlets=[shared])
+ BashOperator(task_id="outlet-only", bash_command="echo 1",
outlets=[outlet_only])
+ BashOperator(task_id="inlet-only", bash_command="echo 1",
inlets=[inlet_only])
+
+ # Attach aliases: one to ``shared`` and one to an asset with no other
reference.
+ alias_only = AssetModel(uri="test://alias_only", name="alias_only",
group="test_group")
+ alias_only.aliases.append(AssetAliasModel(name="alias_for_alias_only",
group="asset"))
+ shared_model = session.scalar(select(AssetModel).where(AssetModel.name
== "shared"))
+ shared_model.aliases.append(AssetAliasModel(name="alias_for_shared",
group="asset"))
+ session.add(alias_only)
+ session.flush()
+
+ self.job_runner._update_asset_orphanage(session=session)
+ session.flush()
+ orphaned, active = self._find_assets_activation(session)
+ assert {a.name for a in active} == {
+ "shared",
+ "dropped",
+ "schedule_only",
+ "outlet_only",
+ "inlet_only",
+ "alias_only",
+ }
+ assert orphaned == []
+
+ with dag_maker(dag_id="producer", schedule=None, session=session):
+ for i in range(3):
+ BashOperator(task_id=f"produce-{i}", bash_command="echo 1",
outlets=[shared])
+ BashOperator(task_id=f"read-{i}", bash_command="echo 1",
inlets=[shared])
+ BashOperator(task_id="outlet-only", bash_command="echo 1",
outlets=[outlet_only])
+ BashOperator(task_id="inlet-only", bash_command="echo 1",
inlets=[inlet_only])
+
+ self.job_runner._update_asset_orphanage(session=session)
+ session.flush()
+ orphaned, active = self._find_assets_activation(session)
+ assert {a.name for a in active} == {
+ "shared",
+ "schedule_only",
+ "outlet_only",
+ "inlet_only",
+ "alias_only",
+ }
+ assert [a.name for a in orphaned] == ["dropped"]
+
def test_asset_orphaning_ignore_orphaned_assets(self, dag_maker, session):
self.job_runner = SchedulerJobRunner(job=Job())
