Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions task-sdk/src/airflow/sdk/execution_time/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -977,7 +977,14 @@ def add_partitions(self, keys: str | list[str]) -> None:

:raises ValueError: If any key is empty/whitespace-only or longer than
``_PARTITION_KEY_MAX_LENGTH`` characters.
:raises TypeError: If this accessor is for an asset alias, since partition
keys are only attached to concrete asset events, not alias events.
"""
if isinstance(self.key, AssetAliasUniqueKey):
raise TypeError(
"add_partitions() is not supported on asset alias outlet events; "
"partition keys can only be attached to a concrete asset."
)
if isinstance(keys, str):
keys = [keys]
for key in keys:
Expand Down
8 changes: 8 additions & 0 deletions task-sdk/tests/task_sdk/execution_time/test_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -536,6 +536,14 @@ def test_add_partitions_rejects_any_invalid_in_list(self, accessor):
accessor.add_partitions(["us", ""])
assert accessor.partition_keys == set()

def test_add_partitions_rejects_asset_alias_accessor(self):
alias_accessor = OutletEventAccessor(
key=AssetAliasUniqueKey.from_asset_alias(AssetAlias("test_alias"))
)
with pytest.raises(TypeError, match="not supported on asset alias"):
alias_accessor.add_partitions("us")
assert alias_accessor.partition_keys == set()


class TestTriggeringAssetEventsAccessor:
@pytest.fixture(autouse=True)
Expand Down
Loading