Lee-W commented on code in PR #45960:
URL: https://github.com/apache/airflow/pull/45960#discussion_r1969550324
##########
task_sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -281,6 +283,84 @@ def _get_asset_from_db(name: str | None = None, uri: str |
None = None) -> Asset
return Asset(**msg.model_dump(exclude={"type"}))
[email protected](init=False)
+class InletEventsAccessors(Mapping[Union[int, Asset, AssetAlias, AssetRef],
Any]):
+ _inlets: list[Any]
+ _assets: dict[AssetUniqueKey, Asset]
+ _asset_aliases: dict[AssetAliasUniqueKey, AssetAlias]
+
+ def __init__(self, inlets: list) -> None:
+ self._inlets = inlets
+ self._assets = {}
+ self._asset_aliases = {}
+
+ for inlet in inlets:
+ if isinstance(inlet, Asset):
+ self._assets[AssetUniqueKey.from_asset(inlet)] = inlet
+ elif isinstance(inlet, AssetAlias):
+
self._asset_aliases[AssetAliasUniqueKey.from_asset_alias(inlet)] = inlet
+ elif isinstance(inlet, AssetNameRef):
+ asset =
OutletEventAccessors._get_asset_from_db(name=inlet.name)
+ self._assets[AssetUniqueKey.from_asset(asset)] = asset
+ elif isinstance(inlet, AssetUriRef):
+ asset = OutletEventAccessors._get_asset_from_db(uri=inlet.uri)
+ self._assets[AssetUniqueKey.from_asset(asset)] = asset
+
+ def __iter__(self) -> Iterator[Asset | AssetAlias]:
+ return iter(self._inlets)
+
+ def __len__(self) -> int:
+ return len(self._inlets)
+
+ def __getitem__(self, key: int | Asset | AssetAlias | AssetRef):
+ from airflow.sdk.definitions.asset import Asset
+
+ if isinstance(key, int): # Support index access; it's easier for
trivial cases.
+ obj = self._inlets[key]
+ if not isinstance(obj, (Asset, AssetAlias, AssetRef)):
+ raise IndexError(key)
+ else:
+ obj = key
+
+ return self._get_asset_events_from_db(obj)
+
+ # TODO: This is temporary to avoid code duplication between here &
airflow/models/taskinstance.py
Review Comment:
This also fails on the outlet end as well
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
