diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index bc62e79318c3c..2a78fd76397c0 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -220,8 +220,23 @@ Another way to achieve the same is by accessing ``outlet_events`` in a task's ex .. code-block:: python @asset(schedule=None) - def write_to_s3(self, context): - context["outlet_events"][self].extra = {"row_count": len(df)} + def write_to_s3(self, *, outlet_events): + outlet_events[self].extra = {"row_count": len(df)} + +These two APIs are equivalent for extra: Airflow writes yielded ``Metadata.extra`` onto the same per-task-instance accessor (merging with ``update``). Extra never implies a partition key; see :ref:`asset-partitions` for attaching keys with ``add_partitions`` or ``Metadata.partition_key``. + +From a ``@task``, ``outlet_events`` is a context key. A default of ``None`` is fine; a non-``None`` default is not. Prefer a keyword-only parameter. A positional ``outlet_events=None`` plus a colliding positional argument is not supported: + +.. code-block:: python + + from airflow.sdk import Asset, task + + example_asset = Asset("s3://asset/example.csv") + + + @task(outlets=[example_asset]) + def produce(*, outlet_events): + outlet_events[example_asset].extra = {"row_count": 1} There's minimal magic here---Airflow simply writes the yielded values to the exact same accessor. This also works in classic operators, including ``execute``, ``pre_execute``, and ``post_execute``. @@ -399,6 +414,12 @@ The following example creates an asset event against the S3 URI ``f"s3://bucket/ s3_asset = Asset(uri="s3://bucket/my-task", name="example_s3") yield Metadata(s3_asset, extra={"k": "v"}, alias=AssetAlias("my-task-outputs")) +``Metadata.partition_key`` is recorded on the concrete asset accessor only +(``outlet_events[asset].add_partitions``). Alias-resolved events keep using the +Dag-run partition key. Declare the asset as an outlet, or call +``add_partitions`` on that asset, if the key must appear on the event. Do not +expect ``alias=`` to fan out per-emission partition keys. + Only one asset event is emitted for an added asset, even if it is added to the alias multiple times, or added to multiple aliases. However, if different ``extra`` values are passed, it can emit multiple asset events. In the following example, two asset events will be emitted. .. code-block:: python @@ -533,6 +554,8 @@ depend on whether the producer and consumer have a team association: When Multi-Team mode is disabled, ``access_control`` is ignored and all asset events are delivered to all consuming Dags, preserving backward compatibility. +.. _asset-partitions: + Asset partitions ---------------- @@ -958,26 +981,36 @@ When the partition key is not known ahead of time (for example, a watermark discovered from the source data, a late-arriving file, or a backfill request), let the producing task decide it while it runs. Schedule the producer with ``PartitionedAtRuntime()`` and record the key(s) on the emitted event with -``outlet_events[self].add_partitions(...)``: +``outlet_events[self].add_partitions(...)``, or by yielding ``Metadata`` with +``partition_key``. Extra still only annotates the event; it does not select a +partition: .. code-block:: python - from airflow.sdk import PartitionedAtRuntime, asset + from airflow.sdk import Metadata, PartitionedAtRuntime, asset @asset( uri="file://incoming/player-stats/live-region.csv", schedule=PartitionedAtRuntime(), ) - def live_region_player_stats(self, outlet_events): + def live_region_player_stats(self, *, outlet_events): # The key is only known once the task runs. + outlet_events[self].extra = {"row_count": 1} outlet_events[self].add_partitions("us") + # Same as: + # yield Metadata(self, extra={"row_count": 1}, partition_key="us") + +``Metadata.partition_key`` is recorded on the concrete asset accessor. Events +emitted only through an alias still use the producing Dag run's partition key. Inside an ``@asset`` function, ``self`` (the emitted ``Asset``) and ``outlet_events`` (the outlet event accessor) are reserved parameter names that -Airflow populates at runtime. Pass a single key, or a list to fan out to several -partitions in one run. Each key produces its own asset event, and duplicate -keys collapse to a single event: +Airflow populates at runtime. Prefer keyword-only ``outlet_events``. Pass a +single key, or a list to fan out to several partitions in one run. Each key +produces its own asset event, and duplicate keys collapse to a single event. +Two yields with different ``partition_key`` values on one task instance also +fan out, and share the same merged extra: .. code-block:: python @@ -985,8 +1018,22 @@ keys collapse to a single event: uri="file://incoming/player-stats/multi-region.csv", schedule=PartitionedAtRuntime(), ) - def multi_region_player_stats(self, outlet_events): + def multi_region_player_stats(self, *, outlet_events): outlet_events[self].add_partitions(["us", "eu", "apac"]) + # Same as: + # yield Metadata(self, partition_key="us") + # yield Metadata(self, partition_key="eu") + # yield Metadata(self, partition_key="apac") + +A partitioned consumer (``PartitionedAssetTimetable``) requires a partition +key. An extra-only emit (no ``add_partitions`` / no ``Metadata.partition_key``) +logs a missing-key warning and does not create an ``AssetPartitionDagRun``. +Airflow does not invent a key from extra. + +.. note:: + + Airflow 3 batches asset events that share a partition key into one downstream + Dag run. Distinct keys are required for one run per mapped task instance. When a runtime run emits exactly one partition key, the producing ``dag_run.partition_key`` is back-filled to that key. Downstream Dags consume diff --git a/airflow-core/newsfragments/71993.bugfix.rst b/airflow-core/newsfragments/71993.bugfix.rst new file mode 100644 index 0000000000000..2e67df90f732b --- /dev/null +++ b/airflow-core/newsfragments/71993.bugfix.rst @@ -0,0 +1 @@ +Add optional ``Metadata.partition_key`` so yielding ``Metadata`` records partitions the same way as ``outlet_events.add_partitions``. diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 9f14dee0b786c..2bdc9b01a3e87 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -127,6 +127,7 @@ FixedKeyMapper, HourWindow, IdentityMapper, + Metadata, MinimumCount, RollupMapper, SegmentWindow, @@ -12832,6 +12833,325 @@ def test_consumer_dag_listen_to_two_partitioned_asset_with_key_1_mapper( assert asset_event.source_run_id == "test" +def _run_mapped_partition_producer( + *, + dag_id: str, + asset: Asset, + items: list[str], + use_metadata: bool, + extra_only: bool, + dag_maker: DagMaker, + session: Session, +) -> None: + with dag_maker(dag_id=dag_id, schedule=None, session=session): + if extra_only and use_metadata: + + @task(outlets=[asset]) + def produce(item): + yield Metadata(asset, extra={"section": item}) + + elif extra_only: + + @task(outlets=[asset]) + def produce(item, *, outlet_events): + outlet_events[asset].extra = {"section": item} + + elif use_metadata: + + @task(outlets=[asset]) + def produce(item): + yield Metadata(asset, extra={"section": item}, partition_key=item) + + else: + + @task(outlets=[asset]) + def produce(item, *, outlet_events): + outlet_events[asset].extra = {"section": item} + outlet_events[asset].add_partitions(item) + + produce.expand(item=items) + + dr = dag_maker.create_dagrun(session=session) + for map_index in range(len(items)): + dag_maker.run_ti("produce", dag_run=dr, map_index=map_index, session=session) + + +@pytest.mark.need_serialized_dag +@pytest.mark.usefixtures("clear_asset_partition_rows") +@pytest.mark.parametrize( + "use_metadata", + [True, False], + ids=["metadata", "outlet_events"], +) +def test_mapped_producer_partition_keys_match_between_apis( + dag_maker: DagMaker, session: Session, use_metadata: bool +): + items = ["us", "eu", "apac"] + api = "meta" if use_metadata else "oe" + asset = Asset(name=f"parity-{api}") + consumer_id = f"consumer-{api}" + producer_id = f"producer-{api}" + + with dag_maker( + dag_id=consumer_id, + schedule=PartitionedAssetTimetable( + assets=asset, + default_partition_mapper=IdentityMapper(), + ), + session=session, + ): + EmptyOperator(task_id="hi") + session.commit() + + _run_mapped_partition_producer( + dag_id=producer_id, + asset=asset, + items=items, + use_metadata=use_metadata, + extra_only=False, + dag_maker=dag_maker, + session=session, + ) + + events = session.scalars(select(AssetEvent).where(AssetEvent.source_dag_id == producer_id)).all() + assert {event.partition_key for event in events} == set(items) + assert {event.extra.get("section") for event in events} == set(items) + + apdrs = session.scalars( + select(AssetPartitionDagRun).where(AssetPartitionDagRun.target_dag_id == consumer_id) + ).all() + assert {apdr.partition_key for apdr in apdrs} == set(items) + + runner = SchedulerJobRunner( + job=Job(job_type=SchedulerJobRunner.job_type), executors=[MockExecutor(do_update=False)] + ) + runner._create_dagruns_for_partitioned_asset_dags(session=session) + + consumer_runs = session.scalars(select(DagRun).where(DagRun.dag_id == consumer_id)).all() + assert {run.partition_key for run in consumer_runs} == set(items) + for run in consumer_runs: + consumed = run.consumed_asset_events + assert len(consumed) == 1 + assert consumed[0].partition_key == run.partition_key + assert consumed[0].extra == {"section": run.partition_key} + + +@pytest.mark.need_serialized_dag +@pytest.mark.usefixtures("clear_asset_partition_rows") +@pytest.mark.parametrize( + "use_metadata", + [True, False], + ids=["metadata", "outlet_events"], +) +def test_extra_only_mapped_emit_does_not_create_apdr( + dag_maker: DagMaker, session: Session, use_metadata: bool +): + items = ["us", "eu", "apac"] + api = "meta" if use_metadata else "oe" + asset = Asset(name=f"extra-only-{api}") + consumer_id = f"extra-only-consumer-{api}" + producer_id = f"extra-only-producer-{api}" + + session.execute(delete(Log)) + session.commit() + + with dag_maker( + dag_id=consumer_id, + schedule=PartitionedAssetTimetable( + assets=asset, + default_partition_mapper=IdentityMapper(), + ), + session=session, + ): + EmptyOperator(task_id="hi") + session.commit() + + _run_mapped_partition_producer( + dag_id=producer_id, + asset=asset, + items=items, + use_metadata=use_metadata, + extra_only=True, + dag_maker=dag_maker, + session=session, + ) + + events = session.scalars(select(AssetEvent).where(AssetEvent.source_dag_id == producer_id)).all() + assert len(events) == 3 + assert {event.extra.get("section") for event in events} == set(items) + assert all(event.partition_key is None for event in events) + + assert session.scalar(select(AssetPartitionDagRun)) is None + logs = session.scalars(select(Log).where(Log.event == "missing partition key")).all() + assert logs + + runner = SchedulerJobRunner( + job=Job(job_type=SchedulerJobRunner.job_type), executors=[MockExecutor(do_update=False)] + ) + runner._create_dagruns_for_partitioned_asset_dags(session=session) + + assert session.scalars(select(DagRun).where(DagRun.dag_id == consumer_id)).all() == [] + + +@pytest.mark.need_serialized_dag +@pytest.mark.usefixtures("clear_asset_partition_rows") +def test_runtime_partition_key_without_partition_date_still_queues(dag_maker: DagMaker, session: Session): + """IdentityMapper + Metadata-style key (no producer date) still creates an APDR.""" + asset = Asset(name="rt-no-date") + with dag_maker( + dag_id="rt-no-date-consumer", + schedule=PartitionedAssetTimetable( + assets=asset, + default_partition_mapper=IdentityMapper(), + ), + session=session, + ): + EmptyOperator(task_id="hi") + session.commit() + + with dag_maker(dag_id="rt-no-date-producer", schedule=None, session=session) as dag: + EmptyOperator(task_id="hi", outlets=[asset]) + + dr = dag_maker.create_dagrun(session=session) + [ti] = dr.get_task_instances(session=session) + session.commit() + + TaskInstance.register_asset_changes_in_db( + ti=ti, + task_outlets=[o.asprofile() for o in dag.get_task("hi").outlets], + outlet_events=[ + { + "dest_asset_key": {"name": "rt-no-date", "uri": "rt-no-date"}, + "extra": {}, + "partition_key": "us", + } + ], + session=session, + ) + session.commit() + + apdr = session.scalar(select(AssetPartitionDagRun)) + assert apdr is not None + assert apdr.partition_key == "us" + assert apdr.partition_date is None + + runner = SchedulerJobRunner( + job=Job(job_type=SchedulerJobRunner.job_type), executors=[MockExecutor(do_update=False)] + ) + runner._create_dagruns_for_partitioned_asset_dags(session=session) + session.refresh(apdr) + assert apdr.created_dag_run_id is not None + + +@pytest.mark.need_serialized_dag +@pytest.mark.usefixtures("clear_asset_partition_rows") +def test_partitioned_event_does_not_queue_non_partitioned_consumer(dag_maker: DagMaker, session: Session): + asset = Asset(name="non-part-consumer-asset") + with dag_maker(dag_id="non-part-consumer", schedule=asset, session=session): + EmptyOperator(task_id="hi") + session.commit() + + with dag_maker(dag_id="non-part-producer", schedule=None, session=session) as dag: + EmptyOperator(task_id="hi", outlets=[asset]) + + dr = dag_maker.create_dagrun(session=session) + [ti] = dr.get_task_instances(session=session) + session.commit() + + TaskInstance.register_asset_changes_in_db( + ti=ti, + task_outlets=[o.asprofile() for o in dag.get_task("hi").outlets], + outlet_events=[ + { + "dest_asset_key": { + "name": "non-part-consumer-asset", + "uri": "non-part-consumer-asset", + }, + "extra": {}, + "partition_key": "us", + } + ], + session=session, + ) + session.commit() + + assert ( + session.scalar( + select(func.count()) + .select_from(AssetDagRunQueue) + .where(AssetDagRunQueue.target_dag_id == "non-part-consumer") + ) + == 0 + ) + assert session.scalar(select(AssetPartitionDagRun)) is None + + +@pytest.mark.need_serialized_dag +@pytest.mark.usefixtures("clear_asset_partition_rows") +def test_multi_partition_payloads_share_merged_extra(dag_maker: DagMaker, session: Session): + """add_partitions(['us','eu']) and two Metadata yields are equivalent if extras match.""" + asset = Asset(name="multi-part-extra") + with dag_maker( + dag_id="multi-part-consumer", + schedule=PartitionedAssetTimetable( + assets=asset, + default_partition_mapper=IdentityMapper(), + ), + session=session, + ): + EmptyOperator(task_id="hi") + session.commit() + + with dag_maker(dag_id="multi-part-producer", schedule=None, session=session) as dag: + EmptyOperator(task_id="hi", outlets=[asset]) + + dr = dag_maker.create_dagrun(session=session) + [ti] = dr.get_task_instances(session=session) + session.commit() + + extra = {"row_count": 1} + TaskInstance.register_asset_changes_in_db( + ti=ti, + task_outlets=[o.asprofile() for o in dag.get_task("hi").outlets], + outlet_events=[ + { + "dest_asset_key": {"name": "multi-part-extra", "uri": "multi-part-extra"}, + "extra": extra, + "partition_key": "us", + }, + { + "dest_asset_key": {"name": "multi-part-extra", "uri": "multi-part-extra"}, + "extra": extra, + "partition_key": "eu", + }, + ], + session=session, + ) + session.commit() + + events = session.scalars( + select(AssetEvent).where(AssetEvent.source_dag_id == "multi-part-producer") + ).all() + assert {event.partition_key for event in events} == {"us", "eu"} + assert all(event.extra == extra for event in events) + + apdrs = session.scalars(select(AssetPartitionDagRun)).all() + assert {apdr.partition_key for apdr in apdrs} == {"us", "eu"} + + runner = SchedulerJobRunner( + job=Job(job_type=SchedulerJobRunner.job_type), executors=[MockExecutor(do_update=False)] + ) + runner._create_dagruns_for_partitioned_asset_dags(session=session) + + consumer_runs = session.scalars(select(DagRun).where(DagRun.dag_id == "multi-part-consumer")).all() + assert {run.partition_key for run in consumer_runs} == {"us", "eu"} + for run in consumer_runs: + consumed = run.consumed_asset_events + assert len(consumed) == 1 + assert consumed[0].extra == extra + assert consumed[0].partition_key == run.partition_key + + def _make_n_satisfied_apdrs( *, consumer_dag_id: str, diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 6a348a3f2ca26..ddc9c57b81c59 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -2065,6 +2065,73 @@ def producer(*, outlet_events): assert len(asset_alias_obj.assets) == 1 assert asset_alias_obj.assets[0].uri == asset_uri + @pytest.mark.parametrize( + "keyword_only", + [False, True], + ids=["positional_or_keyword", "keyword_only"], + ) + def test_mapped_outlet_events_extra_injection(self, dag_maker, session, keyword_only): + asset = Asset("test_mapped_outlet_events_extra") + with dag_maker(schedule=None, serialized=True, session=session): + if keyword_only: + + @task(outlets=[asset]) + def write(x, *, outlet_events=None): + outlet_events[asset].extra = {"n": x} + + else: + + @task(outlets=[asset]) + def write(x, outlet_events=None): + outlet_events[asset].extra = {"n": x} + + write.expand(x=[1, 2, 3]) + + dr = dag_maker.create_dagrun() + for map_index in range(3): + dag_maker.run_ti("write", dag_run=dr, map_index=map_index, session=session) + + events = session.scalars(select(AssetEvent).where(AssetEvent.source_task_id == "write")).all() + assert len(events) == 3 + assert {event.extra["n"] for event in events} == {1, 2, 3} + assert all(event.partition_key is None for event in events) + + def test_mapped_metadata_partition_key_and_extra_per_ti(self, dag_maker, session): + asset = Asset("test_mapped_metadata_partition") + items = ["us", "eu", "apac"] + with dag_maker(schedule=None, serialized=True, session=session): + + @task(outlets=[asset]) + def write(item): + yield Metadata(asset, {"section": item}, partition_key=item) + + write.expand(item=items) + + dr = dag_maker.create_dagrun() + for map_index in range(len(items)): + dag_maker.run_ti("write", dag_run=dr, map_index=map_index, session=session) + + events = session.scalars(select(AssetEvent).where(AssetEvent.source_task_id == "write")).all() + assert len(events) == 3 + assert {event.partition_key for event in events} == set(items) + for event in events: + assert event.extra == {"section": event.partition_key} + + def test_metadata_invalid_partition_key_fails_task_without_event(self, dag_maker, session): + asset = Asset("test_metadata_invalid_partition_key") + with dag_maker(schedule=None, serialized=True, session=session): + + @task(outlets=[asset]) + def write(): + yield Metadata(asset, extra={"n": 1}, partition_key="") + + write() + + with pytest.raises(ValueError, match="must not be empty"): + dag_maker.run_ti("write") + + assert session.scalar(select(AssetEvent)) is None + @pytest.mark.need_serialized_dag def test_outlet_asset_alias_asset_not_exists(self, dag_maker, session): asset_alias_name = "test_outlet_asset_alias_asset_not_exists_asset_alias" diff --git a/task-sdk/src/airflow/sdk/definitions/asset/metadata.py b/task-sdk/src/airflow/sdk/definitions/asset/metadata.py index ee8c42b3ad4ae..583ebedfdd52d 100644 --- a/task-sdk/src/airflow/sdk/definitions/asset/metadata.py +++ b/task-sdk/src/airflow/sdk/definitions/asset/metadata.py @@ -31,8 +31,15 @@ @attrs.define(init=True) class Metadata: - """Metadata to attach to an AssetEvent.""" + """ + Metadata to attach to an AssetEvent. + + ``extra`` is a JSON-serializable mapping merged onto the outlet event. + ``partition_key`` is an optional identity for partitioned consumers; it is + not inferred from ``extra``. + """ asset: Asset extra: dict[str, JsonValue] = attrs.field(factory=dict) alias: AssetAlias | None = None + partition_key: str | None = None diff --git a/task-sdk/src/airflow/sdk/execution_time/callback_runner.py b/task-sdk/src/airflow/sdk/execution_time/callback_runner.py index 83a80bda0ece1..68df8c9685dfe 100644 --- a/task-sdk/src/airflow/sdk/execution_time/callback_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/callback_runner.py @@ -116,6 +116,8 @@ def _run(): for metadata in _run(): if isinstance(metadata, Metadata): outlet_events[metadata.asset].extra.update(metadata.extra) + if metadata.partition_key is not None: + outlet_events[metadata.asset].add_partitions(metadata.partition_key) if metadata.alias: outlet_events[metadata.alias].add(metadata.asset, extra=metadata.extra) else: @@ -164,6 +166,8 @@ async def run(*args: P.args, **kwargs: P.kwargs) -> R: async for result in func(*args, **kwargs): if isinstance(result, Metadata): outlet_events[result.asset].extra.update(result.extra) + if result.partition_key is not None: + outlet_events[result.asset].add_partitions(result.partition_key) if result.alias: outlet_events[result.alias].add(result.asset, extra=result.extra) diff --git a/task-sdk/tests/task_sdk/execution_time/test_callback_runner.py b/task-sdk/tests/task_sdk/execution_time/test_callback_runner.py new file mode 100644 index 0000000000000..74f3c6418f31a --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_callback_runner.py @@ -0,0 +1,171 @@ +# 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. + +from __future__ import annotations + +import logging + +import pytest + +from airflow.sdk.definitions.asset import Asset, AssetAlias +from airflow.sdk.definitions.asset.metadata import Metadata +from airflow.sdk.execution_time.callback_runner import ( + create_async_executable_runner, + create_executable_runner, +) +from airflow.sdk.execution_time.context import OutletEventAccessors +from airflow.sdk.execution_time.task_runner import _serialize_outlet_events + +ASSET = Asset("a") +LOGGER = logging.getLogger("test_callback_runner") + + +def _run(func, outlet_events: OutletEventAccessors | None = None) -> OutletEventAccessors: + if outlet_events is None: + outlet_events = OutletEventAccessors() + create_executable_runner(func, outlet_events, logger=LOGGER).run() + return outlet_events + + +class TestCreateExecutableRunnerMetadata: + def test_two_yields_merge_extra_and_collect_partition_keys(self): + def gen(): + yield Metadata(ASSET, extra={"a": 1}, partition_key="us") + yield Metadata(ASSET, extra={"b": 2}, partition_key="eu") + + outlet_events = _run(gen) + accessor = outlet_events[ASSET] + assert accessor.extra == {"a": 1, "b": 2} + assert accessor.partition_keys == {"us", "eu"} + + events = list(_serialize_outlet_events(outlet_events)) + assert sorted(events, key=lambda e: e["partition_key"]) == [ + { + "dest_asset_key": {"name": "a", "uri": "a"}, + "extra": {"a": 1, "b": 2}, + "partition_key": "eu", + }, + { + "dest_asset_key": {"name": "a", "uri": "a"}, + "extra": {"a": 1, "b": 2}, + "partition_key": "us", + }, + ] + + def test_add_partitions_list_matches_two_metadata_yields(self): + def via_metadata(): + yield Metadata(ASSET, extra={"row_count": 1}, partition_key="us") + yield Metadata(ASSET, extra={"row_count": 1}, partition_key="eu") + + def via_add_partitions(*, outlet_events): + outlet_events[ASSET].extra = {"row_count": 1} + outlet_events[ASSET].add_partitions(["us", "eu"]) + + metadata_events = OutletEventAccessors() + add_events = OutletEventAccessors() + _run(via_metadata, metadata_events) + via_add_partitions(outlet_events=add_events) + + assert metadata_events[ASSET].extra == add_events[ASSET].extra == {"row_count": 1} + assert metadata_events[ASSET].partition_keys == add_events[ASSET].partition_keys == {"us", "eu"} + assert sorted(_serialize_outlet_events(metadata_events), key=lambda e: e["partition_key"]) == sorted( + _serialize_outlet_events(add_events), key=lambda e: e["partition_key"] + ) + + def test_extra_does_not_imply_partition_key(self): + def gen(): + yield Metadata(ASSET, extra={"section": "us"}) + + outlet_events = _run(gen) + accessor = outlet_events[ASSET] + assert accessor.extra == {"section": "us"} + assert accessor.partition_keys == set() + events = list(_serialize_outlet_events(outlet_events)) + assert events == [{"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {"section": "us"}}] + assert "partition_key" not in events[0] + + @pytest.mark.parametrize( + "key", + ["", "a" * 251], + ids=["empty", "too_long"], + ) + def test_invalid_partition_key_raises(self, key): + def gen(): + yield Metadata(ASSET, extra={"a": 1}, partition_key=key) + + with pytest.raises(ValueError, match="partition_key"): + _run(gen) + + def test_alias_and_partition_key_records_key_on_asset(self): + alias = AssetAlias("outputs") + + def gen(): + yield Metadata(ASSET, extra={"k": "v"}, alias=alias, partition_key="us") + + outlet_events = _run(gen) + assert outlet_events[ASSET].extra == {"k": "v"} + assert outlet_events[ASSET].partition_keys == {"us"} + assert outlet_events[alias].partition_keys == set() + assert len(outlet_events[alias].asset_alias_events) == 1 + assert outlet_events[alias].asset_alias_events[0].extra == {"k": "v"} + + events = list(_serialize_outlet_events(outlet_events)) + asset_payloads = [e for e in events if "source_alias_name" not in e] + alias_payloads = [e for e in events if "source_alias_name" in e] + assert asset_payloads == [ + { + "dest_asset_key": {"name": "a", "uri": "a"}, + "extra": {"k": "v"}, + "partition_key": "us", + } + ] + assert len(alias_payloads) == 1 + assert "partition_key" not in alias_payloads[0] + + def test_alias_as_metadata_asset_with_partition_key_raises_type_error(self): + alias = AssetAlias("outputs") + + def gen(): + yield Metadata(alias, extra={"k": 1}, partition_key="us") + + with pytest.raises(TypeError, match="not supported on asset alias"): + _run(gen) + + +class TestCreateAsyncExecutableRunnerMetadata: + @pytest.mark.asyncio + async def test_two_yields_merge_extra_and_collect_partition_keys(self): + outlet_events = OutletEventAccessors() + + async def gen(): + yield Metadata(ASSET, extra={"a": 1}, partition_key="us") + yield Metadata(ASSET, extra={"b": 2}, partition_key="eu") + + await create_async_executable_runner(gen, outlet_events, logger=LOGGER).run() + accessor = outlet_events[ASSET] + assert accessor.extra == {"a": 1, "b": 2} + assert accessor.partition_keys == {"us", "eu"} + + @pytest.mark.asyncio + async def test_invalid_partition_key_raises(self): + outlet_events = OutletEventAccessors() + + async def gen(): + yield Metadata(ASSET, extra={"a": 1}, partition_key="") + + with pytest.raises(ValueError, match="partition_key"): + await create_async_executable_runner(gen, outlet_events, logger=LOGGER).run() diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index cc9fb77e08921..2bdb8bb94eb79 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -2181,6 +2181,29 @@ def test_dedupes_partition_keys_at_serialization(self): {"dest_asset_key": {"name": "a", "uri": "a"}, "extra": {}, "partition_key": "us"}, ] + def test_emits_shared_extra_on_each_partition_key(self): + """Extra plus two keys serializes to two payloads with the same extra and distinct keys.""" + accessors = OutletEventAccessors() + accessor = accessors[Asset(name="a")] + accessor.extra.update({"row_count": 1}) + accessor.add_partitions("us") + accessor.add_partitions("eu") + + events = list(_serialize_outlet_events(accessors)) + + assert sorted(events, key=lambda e: e["partition_key"]) == [ + { + "dest_asset_key": {"name": "a", "uri": "a"}, + "extra": {"row_count": 1}, + "partition_key": "eu", + }, + { + "dest_asset_key": {"name": "a", "uri": "a"}, + "extra": {"row_count": 1}, + "partition_key": "us", + }, + ] + class TestRuntimeTaskInstance: def test_get_context_without_ti_context_from_server(self, mocked_parse, make_ti_context):