You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ur...@apache.org on 2025/04/08 05:08:20 UTC
(airflow) branch main updated: Optimistically activate assets on addition (#48867)
This is an automated email from the ASF dual-hosted git repository.
uranusjr 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 86f76a1550b Optimistically activate assets on addition (#48867)
86f76a1550b is described below
commit 86f76a1550b4be7994e5c35f45c924e369b6a64d
Author: Tzu-ping Chung <ur...@gmail.com>
AuthorDate: Tue Apr 8 13:08:13 2025 +0800
Optimistically activate assets on addition (#48867)
This uses an SQL on-conflict clause to activate assets when they are
collected from DAG files in simple cases. This should cover most simple
cases. where the new asset does not conflict with anything that already
existed in the active table.
When this optimistic insertion fails, the scheduler would come in later
and check more comprehensively whether the existing record is actually
an orphan, do a cleanup, and activate the new asset.
This improves user ergonomics when adding a new asset. In most simple
cases, the new asset would be active more immediately. On edge cases
where an asset is deleted and added later and so on, the user is more
likely to expect the update might take some time.
---
.../src/airflow/dag_processing/collection.py | 44 +++++++++--
.../src/airflow/jobs/scheduler_job_runner.py | 2 +
airflow-core/src/airflow/models/dag.py | 1 +
.../airflow/serialization/serialized_objects.py | 19 ++---
.../tests/unit/dag_processing/test_collection.py | 86 ++++++++++++++++++++--
airflow-core/tests/unit/jobs/test_scheduler_job.py | 27 ++-----
airflow-core/tests/unit/models/test_dag.py | 43 -----------
.../tests/unit/models/test_taskinstance.py | 23 +++---
8 files changed, 150 insertions(+), 95 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py
index 2304e5612b1..4af608d7712 100644
--- a/airflow-core/src/airflow/dag_processing/collection.py
+++ b/airflow-core/src/airflow/dag_processing/collection.py
@@ -29,7 +29,7 @@ from __future__ import annotations
import logging
import traceback
-from typing import TYPE_CHECKING, NamedTuple, cast
+from typing import TYPE_CHECKING, NamedTuple
from sqlalchemy import delete, func, insert, select, tuple_
from sqlalchemy.exc import OperationalError
@@ -37,6 +37,7 @@ from sqlalchemy.orm import joinedload, load_only
from airflow.assets.manager import asset_manager
from airflow.models.asset import (
+ AssetActive,
AssetAliasModel,
AssetModel,
DagScheduleAssetAliasReference,
@@ -64,7 +65,7 @@ if TYPE_CHECKING:
from sqlalchemy.sql import Select
from airflow.models.dagwarning import DagWarning
- from airflow.serialization.serialized_objects import MaybeSerializedDAG, SerializedAssetWatcher
+ from airflow.serialization.serialized_objects import MaybeSerializedDAG
from airflow.typing_compat import Self
log = logging.getLogger(__name__)
@@ -665,6 +666,33 @@ class AssetModelOperation(NamedTuple):
)
return orm_aliases
+ def activate_assets_if_possible(self, models: Iterable[AssetModel], *, session: Session) -> None:
+ """
+ Try to activate assets eagerly.
+
+ This inserts a record to AssetActive for an asset so it is activated
+ on creation if its ``name`` and ``uri`` values do not conflict with
+ anything else. This is a best-effort operation; we simply give up if
+ there's a conflict. The scheduler makes a more comprehensive pass
+ through all assets in ``_update_asset_orphanage``.
+ """
+ if (dialect_name := session.bind.dialect.name) == "postgresql":
+ from sqlalchemy.dialects.postgresql import insert
+
+ stmt = insert(AssetActive).on_conflict_do_nothing()
+ elif dialect_name == "mysql":
+ from sqlalchemy.dialects.mysql import insert
+
+ # MySQL does not support "do nothing"; this updates the row in
+ # conflict with its own value to achieve the same idea.
+ stmt = insert(AssetActive).on_duplicate_key_update(name=AssetActive.name)
+ else:
+ from sqlalchemy.dialects.sqlite import insert
+
+ stmt = insert(AssetActive).on_conflict_do_nothing()
+ if values := [{"name": m.name, "uri": m.uri} for m in models]:
+ session.execute(stmt, values)
+
def add_dag_asset_references(
self,
dags: dict[str, DagModel],
@@ -787,6 +815,8 @@ class AssetModelOperation(NamedTuple):
def add_asset_trigger_references(
self, assets: dict[tuple[str, str], AssetModel], *, session: Session
) -> None:
+ from airflow.serialization.serialized_objects import _encode_trigger
+
# Update references from assets being used
refs_to_add: dict[tuple[str, str], set[int]] = {}
refs_to_remove: dict[tuple[str, str], set[int]] = {}
@@ -797,16 +827,14 @@ class AssetModelOperation(NamedTuple):
for name_uri, asset in self.assets.items():
# If the asset belong to a DAG not active or paused, consider there is no watcher associated to it
- asset_watchers: list[SerializedAssetWatcher] = (
- [cast("SerializedAssetWatcher", watcher) for watcher in asset.watchers]
+ asset_watcher_triggers = (
+ [_encode_trigger(watcher.trigger) for watcher in asset.watchers]
if name_uri in active_assets
else []
)
trigger_hash_to_trigger_dict: dict[int, dict] = {
- BaseEventTrigger.hash(
- watcher.trigger["classpath"], watcher.trigger["kwargs"]
- ): watcher.trigger
- for watcher in asset_watchers
+ BaseEventTrigger.hash(trigger["classpath"], trigger["kwargs"]): trigger
+ for trigger in asset_watcher_triggers
}
triggers.update(trigger_hash_to_trigger_dict)
trigger_hash_from_asset: set[int] = set(trigger_hash_to_trigger_dict.keys())
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 81f5f289fb2..43fc54386c5 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -2282,6 +2282,8 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
An orphaned asset is no longer referenced in any DAG schedule parameters
or task outlets. Active assets (non-orphaned) have entries in AssetActive
and must have unique names and URIs.
+
+ :seealso: :meth:`AssetModelOperation.activate_assets_if_possible`.
"""
# Group assets into orphaned=True and orphaned=False groups.
orphaned = (
diff --git a/airflow-core/src/airflow/models/dag.py b/airflow-core/src/airflow/models/dag.py
index 6995bb39fa9..35a6893ed13 100644
--- a/airflow-core/src/airflow/models/dag.py
+++ b/airflow-core/src/airflow/models/dag.py
@@ -1905,6 +1905,7 @@ class DAG(TaskSDKDag, LoggingMixin):
asset_op.add_dag_asset_name_uri_references(session=session)
asset_op.add_task_asset_references(orm_dags, orm_assets, session=session)
asset_op.add_asset_trigger_references(orm_assets, session=session)
+ asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
dag_op.update_dag_asset_expression(orm_dags=orm_dags, orm_assets=orm_assets)
diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py
index 4826cb01729..a032a453947 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -248,6 +248,16 @@ class _PriorityWeightStrategyNotRegistered(AirflowException):
)
+def _encode_trigger(trigger: BaseEventTrigger | dict):
+ if isinstance(trigger, dict):
+ return trigger
+ classpath, kwargs = trigger.serialize()
+ return {
+ "classpath": classpath,
+ "kwargs": kwargs,
+ }
+
+
def encode_asset_condition(var: BaseAsset) -> dict[str, Any]:
"""
Encode an asset condition.
@@ -262,15 +272,6 @@ def encode_asset_condition(var: BaseAsset) -> dict[str, Any]:
"trigger": _encode_trigger(watcher.trigger),
}
- def _encode_trigger(trigger: BaseEventTrigger | dict):
- if isinstance(trigger, dict):
- return trigger
- classpath, kwargs = trigger.serialize()
- return {
- "classpath": classpath,
- "kwargs": kwargs,
- }
-
asset = {
"__type": DAT.ASSET,
"name": var.name,
diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py
index f45d710692c..5a19747b759 100644
--- a/airflow-core/tests/unit/dag_processing/test_collection.py
+++ b/airflow-core/tests/unit/dag_processing/test_collection.py
@@ -42,6 +42,7 @@ from airflow.listeners.listener import get_listener_manager
from airflow.models import DagModel, DagRun, Trigger
from airflow.models.asset import (
AssetActive,
+ AssetModel,
DagScheduleAssetNameReference,
DagScheduleAssetUriReference,
asset_trigger_association_table,
@@ -150,14 +151,10 @@ class TestAssetModelOperation:
dag.is_paused = is_paused
orm_assets = asset_op.sync_assets(session=session)
- # Create AssetActive objects from assets. It is usually done in the scheduler
- for asset in orm_assets.values():
- session.add(AssetActive.for_asset(asset))
- session.commit()
-
+ session.flush()
+ asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
asset_op.add_asset_trigger_references(orm_assets, session=session)
-
- session.commit()
+ session.flush()
assert session.query(Trigger).count() == expected_num_triggers
assert session.query(asset_trigger_association_table).count() == expected_num_triggers
@@ -241,6 +238,81 @@ class TestAssetModelOperation:
assert next(iter(orm_aliases.values())).group == "new_group"
+@pytest.mark.db_test
+class TestAssetModelOperationSyncAssetActive:
+ @staticmethod
+ def clean_db():
+ clear_db_dags()
+ clear_db_assets()
+ clear_db_triggers()
+
+ @pytest.fixture(autouse=True)
+ def per_test(self) -> Generator:
+ self.clean_db()
+ yield
+ self.clean_db()
+
+ def test_add_asset_activate(self, dag_maker, session):
+ asset = Asset("myasset", "file://myasset/", group="old_group")
+ with dag_maker(schedule=[asset]) as dag:
+ EmptyOperator(task_id="mytask")
+
+ asset_op = AssetModelOperation.collect({dag.dag_id: dag})
+ orm_assets = asset_op.sync_assets(session=session)
+ session.flush()
+ assert len(orm_assets) == 1
+
+ asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
+ session.flush()
+ assert orm_assets["myasset", "file://myasset/"].active is not None
+
+ def test_add_asset_activate_already_exists(self, dag_maker, session):
+ asset = Asset("myasset", "file://myasset/", group="old_group")
+
+ session.add(AssetModel.from_public(asset))
+ session.flush()
+ session.add(AssetActive.for_asset(asset))
+ session.flush()
+
+ with dag_maker(schedule=[asset]) as dag:
+ EmptyOperator(task_id="mytask")
+
+ asset_op = AssetModelOperation.collect({dag.dag_id: dag})
+ orm_assets = asset_op.sync_assets(session=session)
+ session.flush()
+ assert len(orm_assets) == 1
+
+ asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
+ session.flush()
+ assert orm_assets["myasset", "file://myasset/"].active is not None, "should pick up existing active"
+
+ @pytest.mark.parametrize(
+ "existing_assets",
+ [
+ pytest.param([Asset("myasset", uri="file://different/asset")], id="name"),
+ pytest.param([Asset("another", uri="file://myasset/")], id="uri"),
+ ],
+ )
+ def test_add_asset_activate_conflict(self, dag_maker, session, existing_assets):
+ session.add_all(AssetModel.from_public(a) for a in existing_assets)
+ session.flush()
+ session.add_all(AssetActive.for_asset(a) for a in existing_assets)
+ session.flush()
+
+ asset = Asset(name="myasset", uri="file://myasset/", group="old_group")
+ with dag_maker(schedule=[asset]) as dag:
+ EmptyOperator(task_id="mytask")
+
+ asset_op = AssetModelOperation.collect({dag.dag_id: dag})
+ orm_assets = asset_op.sync_assets(session=session)
+ session.flush()
+ assert len(orm_assets) == 1
+
+ asset_op.activate_assets_if_possible(orm_assets.values(), session=session)
+ session.flush()
+ assert orm_assets["myasset", "file://myasset/"].active is None, "should not activate due to conflict"
+
+
@pytest.mark.db_test
class TestUpdateDagParsingResults:
"""Tests centred around the ``update_dag_parsing_results_in_db`` function."""
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index bcdbde271ca..1c6488cb936 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -6296,7 +6296,6 @@ class TestSchedulerJob:
).all()
return [a for a, v in assets if not v], [a for a, v in assets if v]
- @pytest.mark.want_activate_assets(False)
def test_asset_orphaning(self, dag_maker, session):
self.job_runner = SchedulerJobRunner(job=Job())
@@ -6309,27 +6308,22 @@ class TestSchedulerJob:
with dag_maker(dag_id="assets-1", schedule=[asset1, asset2], session=session):
BashOperator(task_id="task", bash_command="echo 1", outlets=[asset3, asset4])
- # Assets not activated yet; asset5 is not even registered (since it's not used anywhere).
+ # asset5 is not registered (since it's not used anywhere).
orphaned, active = self._find_assets_activation(session)
- assert active == []
- assert orphaned == [asset1, asset2, asset3, asset4]
+ assert active == [asset1, asset2, asset3, asset4]
+ assert orphaned == []
self.job_runner._update_asset_orphanage(session=session)
session.flush()
- # Assets are activated after scheduler loop.
- orphaned, active = self._find_assets_activation(session)
- assert active == [asset1, asset2, asset3, asset4]
- assert orphaned == []
-
# Now remove 2 asset references and add asset5.
with dag_maker(dag_id="assets-1", schedule=[asset1], session=session):
BashOperator(task_id="task", bash_command="echo 1", outlets=[asset3, asset5])
- # The DAG parser finds asset5, but it's not activated yet.
+ # The DAG parser finds asset5.
orphaned, active = self._find_assets_activation(session)
- assert active == [asset1, asset2, asset3, asset4]
- assert orphaned == [asset5]
+ assert active == [asset1, asset2, asset3, asset4, asset5]
+ assert orphaned == []
self.job_runner._update_asset_orphanage(session=session)
session.flush()
@@ -6339,7 +6333,6 @@ class TestSchedulerJob:
assert active == [asset1, asset3, asset5]
assert orphaned == [asset2, asset4]
- @pytest.mark.want_activate_assets(False)
def test_asset_orphaning_ignore_orphaned_assets(self, dag_maker, session):
self.job_runner = SchedulerJobRunner(job=Job())
@@ -6349,16 +6342,12 @@ class TestSchedulerJob:
BashOperator(task_id="task", bash_command="echo 1")
orphaned, active = self._find_assets_activation(session)
- assert active == []
- assert orphaned == [asset1]
+ assert active == [asset1]
+ assert orphaned == []
self.job_runner._update_asset_orphanage(session=session)
session.flush()
- orphaned, active = self._find_assets_activation(session)
- assert active == [asset1]
- assert orphaned == []
-
# now remove asset1 reference
with dag_maker(dag_id="assets-1", schedule=None, session=session):
BashOperator(task_id="task", bash_command="echo 1")
diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py
index 468b47a7565..e83cd485541 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -42,7 +42,6 @@ from airflow.exceptions import (
)
from airflow.models import DagBag
from airflow.models.asset import (
- AssetActive,
AssetAliasModel,
AssetDagRunQueue,
AssetEvent,
@@ -902,48 +901,6 @@ class TestDag:
.all()
) == {(task_id, dag_id1, asset2_orm.id)}
- @staticmethod
- def _find_assets_activation(session) -> tuple[list[AssetModel], list[AssetModel]]:
- assets = session.execute(
- select(AssetModel, AssetActive)
- .outerjoin(
- AssetActive,
- (AssetModel.name == AssetActive.name) & (AssetModel.uri == AssetActive.uri),
- )
- .order_by(AssetModel.uri)
- ).all()
- return [a for a, v in assets if not v], [a for a, v in assets if v]
-
- def test_bulk_write_to_db_does_not_activate(self, dag_maker, testing_dag_bundle, session):
- """
- Assets are not activated on write, but later in the scheduler by the SchedulerJob.
- """
- # Create four assets - two that have references and two that are unreferenced and marked as
- # orphans
- asset1 = Asset(uri="test://asset1", name="asset1", group="test-group")
- asset2 = Asset(uri="test://asset2", name="asset2", group="test-group")
- asset3 = Asset(uri="test://asset3", name="asset3", group="test-group")
- asset4 = Asset(uri="test://asset4", name="asset4", group="test-group")
-
- dag1 = DAG(dag_id="assets-1", start_date=DEFAULT_DATE, schedule=[asset1])
- BashOperator(dag=dag1, task_id="task", bash_command="echo 1", outlets=[asset3])
- DAG.bulk_write_to_db("testing", None, [dag1], session=session)
-
- assert session.scalars(select(AssetModel).order_by(AssetModel.uri)).all() == [asset1, asset3]
- assert session.scalars(select(AssetActive)).all() == []
-
- dag1 = DAG(dag_id="assets-1", start_date=DEFAULT_DATE, schedule=[asset1, asset2])
- BashOperator(dag=dag1, task_id="task", bash_command="echo 1", outlets=[asset3, asset4])
- DAG.bulk_write_to_db("testing", None, [dag1], session=session)
-
- assert session.scalars(select(AssetModel).order_by(AssetModel.uri)).all() == [
- asset1,
- asset2,
- asset3,
- asset4,
- ]
- assert session.scalars(select(AssetActive)).all() == []
-
def test_bulk_write_to_db_asset_aliases(self, testing_dag_bundle):
"""
Ensure that asset aliases referenced in a dag are correctly loaded into the database.
diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py
index 3df7f5f367e..c8a76913ff1 100644
--- a/airflow-core/tests/unit/models/test_taskinstance.py
+++ b/airflow-core/tests/unit/models/test_taskinstance.py
@@ -36,7 +36,7 @@ import pendulum
import pytest
import time_machine
import uuid6
-from sqlalchemy import select
+from sqlalchemy import delete, select
from airflow import settings
from airflow.exceptions import (
@@ -4222,29 +4222,34 @@ class TestTaskInstance:
"Asset(name='asset_first', uri='test://asset/')"
)
- @pytest.mark.want_activate_assets(True)
+ @pytest.mark.want_activate_assets(False)
def test_run_with_inactive_assets_in_inlets_within_the_same_dag(self, dag_maker, session):
- from airflow.sdk.definitions.asset import Asset
+ valid_asset = Asset("asset_first")
+ conflict_asset = Asset(name="asset_first", uri="test://asset/")
with dag_maker(schedule=None, serialized=True, session=session):
- @task(inlets=Asset("asset_first"))
+ @task(inlets=valid_asset)
def first_asset_task():
pass
- @task(inlets=Asset(name="asset_first", uri="test://asset"))
- def duplicate_asset_task():
+ @task(inlets=conflict_asset)
+ def conflict_asset_task():
pass
- first_asset_task() >> duplicate_asset_task()
+ first_asset_task() >> conflict_asset_task()
+
+ session.execute(delete(AssetActive))
+ session.add(AssetActive.for_asset(valid_asset))
tis = {ti.task_id: ti for ti in dag_maker.create_dagrun().task_instances}
+ tis["first_asset_task"].run(session=session)
with pytest.raises(AirflowInactiveAssetInInletOrOutletException) as exc:
- tis["first_asset_task"].run(session=session)
+ tis["conflict_asset_task"].run(session=session)
assert str(exc.value) == (
"Task has the following inactive assets in its inlets or outlets: "
- "Asset(name='asset_first', uri='asset_first')"
+ "Asset(name='asset_first', uri='test://asset/')"
)
@pytest.mark.want_activate_assets(True)