This is an automated email from the ASF dual-hosted git repository.
vincbeck 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 98c4b8306f3 Store `allow_producer_teams` in
`dag_schedule_asset_reference` instead of `asset` (#66487)
98c4b8306f3 is described below
commit 98c4b8306f3af0829d6fbfba9571cb9f37975df4
Author: Vincent <[email protected]>
AuthorDate: Fri May 8 08:19:03 2026 -0400
Store `allow_producer_teams` in `dag_schedule_asset_reference` instead of
`asset` (#66487)
---
airflow-core/docs/migrations-ref.rst | 5 +-
airflow-core/src/airflow/assets/manager.py | 9 +++-
.../src/airflow/dag_processing/collection.py | 16 ++++--
...oducer_teams_to_dag_schedule_asset_reference.py | 52 ++++++++++++++++++
airflow-core/src/airflow/models/asset.py | 15 +++---
airflow-core/src/airflow/serialization/decoders.py | 1 +
airflow-core/src/airflow/serialization/encoders.py | 2 +
airflow-core/src/airflow/utils/db.py | 2 +-
airflow-core/tests/unit/assets/test_manager.py | 22 +++++---
.../tests/unit/dag_processing/test_collection.py | 39 ++++++++++++++
airflow-core/tests/unit/models/test_asset.py | 63 ----------------------
11 files changed, 140 insertions(+), 86 deletions(-)
diff --git a/airflow-core/docs/migrations-ref.rst
b/airflow-core/docs/migrations-ref.rst
index 71687ef384b..a9502d28fe9 100644
--- a/airflow-core/docs/migrations-ref.rst
+++ b/airflow-core/docs/migrations-ref.rst
@@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are
executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description
|
+=========================+==================+===================+==============================================================+
-| ``b8f3e4a1d2c9`` (head) | ``fde9ed84d07b`` | ``3.3.0`` | Add
retry_delay_override and retry_reason to task_instance. |
+| ``a7f3b2c1d4e5`` (head) | ``b8f3e4a1d2c9`` | ``3.3.0`` | Add
allow_producer_teams column to |
+| | | |
dag_schedule_asset_reference table. |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``b8f3e4a1d2c9`` | ``fde9ed84d07b`` | ``3.3.0`` | Add
retry_delay_override and retry_reason to task_instance. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``fde9ed84d07b`` | ``9fabad868fdb`` | ``3.3.0`` | Add
task_state and asset_state tables. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
diff --git a/airflow-core/src/airflow/assets/manager.py
b/airflow-core/src/airflow/assets/manager.py
index bcfd7cfd061..9f333efe851 100644
--- a/airflow-core/src/airflow/assets/manager.py
+++ b/airflow-core/src/airflow/assets/manager.py
@@ -187,7 +187,7 @@ class AssetManager(LoggingMixin):
:param dags_to_queue: set of DagModel instances to potentially queue.
:param source_teams: set of team names the source belongs to. Empty
set means teamless.
- :param asset_model: the AssetModel whose extra may contain
allow_producer_teams.
+ :param asset_model: the AssetModel whose scheduled_dags carry
allow_producer_teams.
:param source_is_api: True if the event was triggered via the REST API
(not a DAG task).
:param session: SQLAlchemy session.
"""
@@ -199,12 +199,16 @@ class AssetManager(LoggingMixin):
from airflow.models.dag import DagModel
- allow_producer_teams: list[str] =
asset_model.extra.get("allow_producer_teams", [])
is_teamless_source = len(source_teams) == 0
dag_ids = [dag.dag_id for dag in dags_to_queue]
dag_id_to_team = DagModel.get_dag_id_to_team_name_mapping(dag_ids,
session=session)
+ # Build per-consumer allow_producer_teams from the schedule reference
rows.
+ dag_id_to_allow_teams: dict[str, list[str]] = {
+ ref.dag_id: ref.allow_producer_teams or [] for ref in
asset_model.scheduled_dags
+ }
+
filtered = set()
for dag in dags_to_queue:
consumer_team = dag_id_to_team.get(dag.dag_id)
@@ -227,6 +231,7 @@ class AssetManager(LoggingMixin):
filtered.add(dag)
continue
+ allow_producer_teams = dag_id_to_allow_teams.get(dag.dag_id, [])
if source_teams & set(allow_producer_teams):
# Cross-team via allow_producer_teams
filtered.add(dag)
diff --git a/airflow-core/src/airflow/dag_processing/collection.py
b/airflow-core/src/airflow/dag_processing/collection.py
index 4df0b589c1f..6f8bc752bbc 100644
--- a/airflow-core/src/airflow/dag_processing/collection.py
+++ b/airflow-core/src/airflow/dag_processing/collection.py
@@ -891,15 +891,23 @@ class AssetModelOperation(NamedTuple):
if not references:
dags[dag_id].schedule_asset_references = []
continue
- referenced_asset_ids = {asset.id for asset in (assets[r.name,
r.uri] for r in references)}
+ referenced_assets = {assets[r.name, r.uri]: r.allow_producer_teams
for r in references}
+ referenced_asset_ids = {a.id for a in referenced_assets}
orm_refs = {r.asset_id: r for r in
dags[dag_id].schedule_asset_references}
for asset_id, ref in orm_refs.items():
if asset_id not in referenced_asset_ids:
session.delete(ref)
+ for asset_model, teams in referenced_assets.items():
+ if asset_model.id in orm_refs:
+ orm_refs[asset_model.id].allow_producer_teams = teams
session.bulk_save_objects(
- DagScheduleAssetReference(asset_id=asset_id, dag_id=dag_id)
- for asset_id in referenced_asset_ids
- if asset_id not in orm_refs
+ DagScheduleAssetReference(
+ asset_id=asset_model.id,
+ dag_id=dag_id,
+ allow_producer_teams=teams,
+ )
+ for asset_model, teams in referenced_assets.items()
+ if asset_model.id not in orm_refs
)
def add_dag_asset_alias_references(
diff --git
a/airflow-core/src/airflow/migrations/versions/0114_3_3_0_add_allow_producer_teams_to_dag_schedule_asset_reference.py
b/airflow-core/src/airflow/migrations/versions/0114_3_3_0_add_allow_producer_teams_to_dag_schedule_asset_reference.py
new file mode 100644
index 00000000000..d354607a84f
--- /dev/null
+++
b/airflow-core/src/airflow/migrations/versions/0114_3_3_0_add_allow_producer_teams_to_dag_schedule_asset_reference.py
@@ -0,0 +1,52 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add allow_producer_teams column to dag_schedule_asset_reference table.
+
+Revision ID: a7f3b2c1d4e5
+Revises: b8f3e4a1d2c9
+Create Date: 2026-05-06 12:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = "a7f3b2c1d4e5"
+down_revision = "b8f3e4a1d2c9"
+branch_labels = None
+depends_on = None
+airflow_version = "3.3.0"
+
+
+def upgrade():
+ """Add allow_producer_teams column to dag_schedule_asset_reference."""
+ with op.batch_alter_table("dag_schedule_asset_reference", schema=None) as
batch_op:
+ batch_op.add_column(sa.Column("allow_producer_teams", sa.JSON(),
nullable=True))
+
+
+def downgrade():
+ """Remove allow_producer_teams column from dag_schedule_asset_reference."""
+ from airflow.migrations.utils import disable_sqlite_fkeys
+
+ with disable_sqlite_fkeys(op):
+ with op.batch_alter_table("dag_schedule_asset_reference", schema=None)
as batch_op:
+ batch_op.drop_column("allow_producer_teams")
diff --git a/airflow-core/src/airflow/models/asset.py
b/airflow-core/src/airflow/models/asset.py
index 7a43e056d0d..3e8390fdfaf 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -337,10 +337,12 @@ class AssetModel(Base):
@classmethod
def from_serialized(cls, obj: SerializedAsset) -> AssetModel:
- extra = dict(obj.extra)
- if obj.allow_producer_teams:
- extra["allow_producer_teams"] = obj.allow_producer_teams
- return cls(name=obj.name, uri=obj.uri, group=obj.group, extra=extra)
+ return cls(
+ name=obj.name,
+ uri=obj.uri,
+ group=obj.group,
+ extra=dict(obj.extra),
+ )
def __init__(self, name: str = "", uri: str = "", **kwargs):
if not name and not uri:
@@ -379,14 +381,12 @@ class AssetModel(Base):
def to_serialized(self) -> SerializedAsset:
from airflow.serialization.definitions.assets import SerializedAsset
- extra = {k: v for k, v in self.extra.items() if k !=
"allow_producer_teams"}
return SerializedAsset(
name=self.name,
uri=self.uri,
group=self.group,
- extra=extra,
+ extra=self.extra,
watchers=[],
- allow_producer_teams=self.extra.get("allow_producer_teams", []),
)
def add_trigger(self, trigger: Trigger, watcher_name: str):
@@ -591,6 +591,7 @@ class DagScheduleAssetReference(Base):
asset_id: Mapped[int] = mapped_column(Integer, primary_key=True,
nullable=False)
dag_id: Mapped[str] = mapped_column(StringID(), primary_key=True,
nullable=False)
+ allow_producer_teams: Mapped[list | None] = mapped_column(sa.JSON(),
nullable=True)
created_at: Mapped[datetime] = mapped_column(UtcDateTime,
default=timezone.utcnow, nullable=False)
updated_at: Mapped[datetime] = mapped_column(
UtcDateTime, default=timezone.utcnow, onupdate=timezone.utcnow,
nullable=False
diff --git a/airflow-core/src/airflow/serialization/decoders.py
b/airflow-core/src/airflow/serialization/decoders.py
index 27db6134b13..683efdc87e6 100644
--- a/airflow-core/src/airflow/serialization/decoders.py
+++ b/airflow-core/src/airflow/serialization/decoders.py
@@ -106,6 +106,7 @@ def _decode_asset(var: dict[str, Any]):
)
for watcher in watchers
],
+ allow_producer_teams=var.get("allow_producer_teams", []),
)
diff --git a/airflow-core/src/airflow/serialization/encoders.py
b/airflow-core/src/airflow/serialization/encoders.py
index 87d7ef1f137..9e341cbe783 100644
--- a/airflow-core/src/airflow/serialization/encoders.py
+++ b/airflow-core/src/airflow/serialization/encoders.py
@@ -190,6 +190,8 @@ def encode_asset_like(a: BaseAsset | SerializedAssetBase)
-> dict[str, Any]:
d = {"__type": DAT.ASSET, "name": a.name, "uri": a.uri, "group":
a.group, "extra": a.extra}
if a.watchers:
d["watchers"] = [{"name": w.name, "trigger":
encode_trigger(w.trigger)} for w in a.watchers]
+ if a.allow_producer_teams:
+ d["allow_producer_teams"] = a.allow_producer_teams
return d
case AssetAlias() | SerializedAssetAlias():
return {"__type": DAT.ASSET_ALIAS, "name": a.name, "group":
a.group}
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index 59b16228e37..4caa9901bfb 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -116,7 +116,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.1.0": "cc92b33c6709",
"3.1.8": "509b94a1042d",
"3.2.0": "1d6611b6ab7c",
- "3.3.0": "b8f3e4a1d2c9",
+ "3.3.0": "a7f3b2c1d4e5",
}
# Prefix used to identify tables holding data moved during migration.
diff --git a/airflow-core/tests/unit/assets/test_manager.py
b/airflow-core/tests/unit/assets/test_manager.py
index 03ec82139cb..1f3eead3dc2 100644
--- a/airflow-core/tests/unit/assets/test_manager.py
+++ b/airflow-core/tests/unit/assets/test_manager.py
@@ -326,12 +326,18 @@ def _make_dag(dag_id: str) -> DagModel:
return dag
-def _make_asset_model(allow_producer_teams: list[str] | None = None) ->
AssetModel:
+def _make_asset_model(
+ scheduled_dags: dict[str, list[str]] | None = None,
+) -> AssetModel:
+ """Create a mock AssetModel.
+
+ :param scheduled_dags: mapping of dag_id -> allow_producer_teams for each
consumer reference.
+ """
model = mock.Mock(spec=AssetModel)
- extra = {}
- if allow_producer_teams:
- extra["allow_producer_teams"] = allow_producer_teams
- model.extra = extra
+ model.scheduled_dags = [
+ mock.Mock(dag_id=dag_id, allow_producer_teams=teams)
+ for dag_id, teams in (scheduled_dags or {}).items()
+ ]
return model
@@ -390,7 +396,7 @@ class TestFilterDagsByTeam:
result = AssetManager._filter_dags_by_team(
dags_to_queue={dag},
source_teams={"team_a"},
- asset_model=_make_asset_model(allow_producer_teams=[]),
+ asset_model=_make_asset_model(scheduled_dags={"dag1": []}),
source_is_api=False,
session=mock.Mock(),
)
@@ -406,7 +412,7 @@ class TestFilterDagsByTeam:
result = AssetManager._filter_dags_by_team(
dags_to_queue={dag},
source_teams={"team_a"},
- asset_model=_make_asset_model(allow_producer_teams=["team_a"]),
+ asset_model=_make_asset_model(scheduled_dags={"dag1":
["team_a"]}),
source_is_api=False,
session=mock.Mock(),
)
@@ -490,7 +496,7 @@ class TestFilterDagsByTeam:
result = AssetManager._filter_dags_by_team(
dags_to_queue={dag},
source_teams={"team_a"},
- asset_model=_make_asset_model(allow_producer_teams=["team_a"]),
+ asset_model=_make_asset_model(scheduled_dags={"dag1":
["team_a"]}),
source_is_api=True,
session=mock.Mock(),
)
diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py
b/airflow-core/tests/unit/dag_processing/test_collection.py
index 8dd02fd5a88..42be7792381 100644
--- a/airflow-core/tests/unit/dag_processing/test_collection.py
+++ b/airflow-core/tests/unit/dag_processing/test_collection.py
@@ -143,6 +143,45 @@ class TestAssetModelOperation:
yield
self.clean_db()
+ @pytest.mark.usefixtures("testing_dag_bundle")
+ def
test_sync_assets_preserves_allow_producer_teams_from_other_bundle(self,
dag_maker, session):
+ """When a producer bundle (without allow_producer_teams) is synced
after a consumer bundle
+ (with allow_producer_teams), the stored allow_producer_teams must not
be wiped out."""
+ from airflow.models.asset import DagScheduleAssetReference
+
+ # First sync: consumer bundle sets allow_producer_teams on the asset.
+ consumer_asset = Asset("shared_asset", allow_producer_teams=["team1",
"team2"])
+ with dag_maker(dag_id="consumer_dag", schedule=[consumer_asset]) as
consumer_dag:
+ EmptyOperator(task_id="mytask")
+
+ consumer_dags = {consumer_dag.dag_id:
LazyDeserializedDAG.from_dag(consumer_dag)}
+ orm_dags = DagModelOperation(consumer_dags, "testing",
None).add_dags(session=session)
+ asset_op = AssetModelOperation.collect(consumer_dags)
+ orm_assets = asset_op.sync_assets(session=session)
+ session.flush()
+ asset_op.add_dag_asset_references(orm_dags, orm_assets,
session=session)
+ session.flush()
+
+ ref = session.scalar(
+
select(DagScheduleAssetReference).where(DagScheduleAssetReference.dag_id ==
"consumer_dag")
+ )
+ assert ref.allow_producer_teams == ["team1", "team2"]
+
+ # Second sync: producer bundle references the same asset WITHOUT
allow_producer_teams.
+ producer_asset = Asset("shared_asset")
+ with dag_maker(dag_id="producer_dag", schedule="@once") as
producer_dag:
+ EmptyOperator(task_id="produce", outlets=[producer_asset])
+
+ producer_dags = {producer_dag.dag_id:
LazyDeserializedDAG.from_dag(producer_dag)}
+ DagModelOperation(producer_dags, "testing",
None).add_dags(session=session)
+ asset_op = AssetModelOperation.collect(producer_dags)
+ asset_op.sync_assets(session=session)
+ session.flush()
+
+ # Consumer's allow_producer_teams must still be preserved.
+ session.expire(ref)
+ assert ref.allow_producer_teams == ["team1", "team2"]
+
@pytest.mark.parametrize(
("is_active", "is_paused", "expected_num_triggers"),
[
diff --git a/airflow-core/tests/unit/models/test_asset.py
b/airflow-core/tests/unit/models/test_asset.py
index 1f4ffa579f5..f4c9e5cc866 100644
--- a/airflow-core/tests/unit/models/test_asset.py
+++ b/airflow-core/tests/unit/models/test_asset.py
@@ -163,66 +163,3 @@ def test_remove_reference_for_inactive_dag(
_simulate_soft_dag_deletion("test2")
remove_references_to_deleted_dags(session=session)
assert set(session.execute(select_stmt)) == set()
-
-
-class TestAssetModelAllowProducerTeams:
- @pytest.mark.parametrize(
- ("extra", "expected_teams"),
- [
- pytest.param({"allow_producer_teams": ["team_a"]}, ["team_a"],
id="with_teams"),
- pytest.param({}, [], id="without_teams"),
- ],
- )
- def test_to_serialized_allow_producer_teams(self, extra, expected_teams):
- model = AssetModel(name="test_asset", uri="test://asset", extra=extra)
- serialized = model.to_serialized()
- assert serialized.allow_producer_teams == expected_teams
- assert "allow_producer_teams" not in serialized.extra
-
- @pytest.mark.parametrize(
- ("teams", "expected_in_extra"),
- [
- pytest.param(["team_a", "team_b"], True, id="with_teams"),
- pytest.param([], False, id="empty_teams"),
- ],
- )
- def test_from_serialized_allow_producer_teams(self, teams,
expected_in_extra):
- from airflow.serialization.definitions.assets import SerializedAsset
-
- serialized = SerializedAsset(
- name="test_asset",
- uri="test://asset",
- group="asset",
- extra={},
- watchers=[],
- allow_producer_teams=teams,
- )
- model = AssetModel.from_serialized(serialized)
- if expected_in_extra:
- assert model.extra["allow_producer_teams"] == teams
- else:
- assert "allow_producer_teams" not in model.extra
-
- @pytest.mark.parametrize(
- ("teams", "extra"),
- [
- pytest.param(["team_analytics", "team_ml"], {"some_key":
"some_value"}, id="with_teams"),
- pytest.param([], {}, id="empty_teams"),
- ],
- )
- def test_round_trip_preserves_allow_producer_teams(self, teams, extra):
- from airflow.serialization.definitions.assets import SerializedAsset
-
- serialized = SerializedAsset(
- name="test_asset",
- uri="test://asset",
- group="asset",
- extra=extra,
- watchers=[],
- allow_producer_teams=teams,
- )
- model = AssetModel.from_serialized(serialized)
- result = model.to_serialized()
- assert result.allow_producer_teams == teams
- assert "allow_producer_teams" not in result.extra
- assert result.extra == extra
