ashb commented on code in PR #45960:
URL: https://github.com/apache/airflow/pull/45960#discussion_r1969542538
##########
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:
I don't quite understand the question. Can you elaborate on it a bit please?
Also we shouldn't really call this method "_from_db" -- as its not hitting
the DB, it's hitting the API server
--
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]
