Lee-W commented on code in PR #75404:
URL: https://github.com/apache/airflow/pull/75404#discussion_r4375236070
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -4217,30 +4217,17 @@ def _update_asset_orphanage(self, *, session: Session =
NEW_SESSION) -> None:
: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.
Count it so it is
+ # activated (and shown in the UI) instead of being treated as
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).where(~referenced).cte()
Review Comment:
Now using `subquery()` insetad of CTEs and both
`_orphan_unreferenced_assets` and `_activate_referenced_assets` are typed as
`Subquery`
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -4217,30 +4217,17 @@ def _update_asset_orphanage(self, *, session: Session =
NEW_SESSION) -> None:
: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.
Count it so it is
Review Comment:
I'm actually ok with count 🤔 but agree it could be ambiguous. updated it as
suggested
##########
airflow-core/tests/unit/jobs/test_scheduler_job.py:
##########
@@ -10166,6 +10166,37 @@ def
test_asset_orphaning_keeps_alias_materialized_asset_active(self, session):
assert asset in active
assert asset not in orphaned
+ def test_asset_orphaning_asset_referenced_from_every_table(self,
dag_maker, session):
Review Comment:
The test now has an inlet-only asset and an alias attached to `shared`, and
re-fetches `_find_assets_activation` after each `_update_asset_orphanage` call,
so each of the four `EXISTS` branches is pinned on its own.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -4217,30 +4217,17 @@ def _update_asset_orphanage(self, *, session: Session =
NEW_SESSION) -> None:
: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.
Count it so it is
+ # activated (and shown in the UI) instead of being treated as
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).where(~referenced).cte()
+ activate_query = select(AssetModel).where(referenced).cte()
Review Comment:
Both queries now select only `AssetModel.name` and `AssetModel.uri`.
##########
airflow-core/src/airflow/jobs/scheduler_job_runner.py:
##########
@@ -4217,30 +4217,17 @@ def _update_asset_orphanage(self, *, session: Session =
NEW_SESSION) -> None:
: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_(
Review Comment:
`_update_asset_orphanage` is now registered with `non_fatal=True`, so a
`statement_timeout` skips the cycle instead of exiting the SchedulerJob.
--
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]
