Allow to_downstream() to return more than one key

airflow-core/src/airflow/assets/manager.py

@@ -41,6 +41,7 @@
     DagScheduleAssetUriReference,
     PartitionedAssetKeyLog,
 )
+from airflow.utils.helpers import is_container
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks
 
@@ -416,22 +417,32 @@ def _queue_partitioned_dags(
             target_key = timetable.get_partition_mapper(
                 name=asset_model.name, uri=asset_model.uri
             ).to_downstream(partition_key)
-
-            apdr = cls._get_or_create_apdr(
-                target_key=target_key,
-                target_dag=target_dag,
-                asset_id=asset_id,
-                session=session,
-            )
-            log_record = PartitionedAssetKeyLog(
-                asset_id=asset_id,
-                asset_event_id=event.id,
-                asset_partition_dag_run_id=apdr.id,
-                source_partition_key=partition_key,
-                target_dag_id=target_dag.dag_id,
-                target_partition_key=target_key,
-            )
-            session.add(log_record)
+            if is_container(target_key):
+                # TODO (AIP-76): This never happens now. When we implement
+                # one-to-many partition key mapping, this should also add a
+                # config to cap the iterable size so the scheduler does not
+                # blow up with an incorrectly implemented PartitionMapper.
+                target_keys: Iterable[str] = target_key
+            else:
+                target_keys = [target_key]
+            del target_key
+
+            for target_key in target_keys:
+                apdr = cls._get_or_create_apdr(
+                    target_key=target_key,
+                    target_dag=target_dag,
+                    asset_id=asset_id,
+                    session=session,
+                )
+                log_record = PartitionedAssetKeyLog(
+                    asset_id=asset_id,
+                    asset_event_id=event.id,
+                    asset_partition_dag_run_id=apdr.id,
+                    source_partition_key=partition_key,
+                    target_dag_id=target_dag.dag_id,
+                    target_partition_key=target_key,
+                )
+                session.add(log_record)
 
     @classmethod
     def _get_or_create_apdr(

airflow-core/src/airflow/partition_mappers/base.py

@@ -18,7 +18,10 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
 
 
 class PartitionMapper(ABC):
@@ -29,7 +32,7 @@ class PartitionMapper(ABC):
     """
 
     @abstractmethod
-    def to_downstream(self, key: str) -> str:
+    def to_downstream(self, key: str) -> str | Iterable[str]:
         """Return the target key that the given source partition key maps to."""
 
     def serialize(self) -> dict[str, Any]:

airflow-core/src/airflow/utils/helpers.py

@@ -34,9 +34,9 @@
 
 if TYPE_CHECKING:
     from datetime import datetime
-    from typing import TypeGuard
 
     import jinja2
+    from typing_extensions import TypeIs
 
     from airflow.models.taskinstance import TaskInstance
 
@@ -95,11 +95,11 @@ def handler(signum, frame):
 
 
 @overload
-def is_container(obj: None | int | Iterable[int] | range) -> TypeGuard[Iterable[int]]: ...
+def is_container(obj: None | int | Iterable[int] | range) -> TypeIs[Iterable[int]]: ...
 
 
 @overload
-def is_container(obj: None | CT | Iterable[CT]) -> TypeGuard[Iterable[CT]]: ...
+def is_container(obj: None | CT | Iterable[CT]) -> TypeIs[Iterable[CT]]: ...
 
 
 def is_container(obj) -> bool: