diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py index 09fb0fd5e679f..205e7efad328b 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py @@ -23,6 +23,7 @@ from sqlalchemy import and_, select from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.types import UtcDateTime from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse from airflow.api_fastapi.execution_api.datamodels.asset_event import ( AssetEventResponse, @@ -40,11 +41,13 @@ def _get_asset_events_through_sql_clauses( - *, join_clause, where_clause, session: SessionDep + *, join_clause, where_clause, session: SessionDep, ascending: bool = True, limit: int | None = None ) -> AssetEventsResponse: - asset_events = session.scalars( - select(AssetEvent).join(join_clause).where(where_clause).order_by(AssetEvent.timestamp) - ) + order_by_clause = AssetEvent.timestamp.asc() if ascending else AssetEvent.timestamp.desc() + asset_events_query = select(AssetEvent).join(join_clause).where(where_clause).order_by(order_by_clause) + if limit: + asset_events_query = asset_events_query.limit(limit) + asset_events = session.scalars(asset_events_query) return AssetEventsResponse.model_validate( { "asset_events": [ @@ -75,6 +78,10 @@ def get_asset_event_by_asset_name_uri( name: Annotated[str | None, Query(description="The name of the Asset")], uri: Annotated[str | None, Query(description="The URI of the Asset")], session: SessionDep, + after: Annotated[UtcDateTime | None, Query(description="The start of the time range")] = None, + before: Annotated[UtcDateTime | None, Query(description="The end of the time range")] = None, + ascending: Annotated[bool, Query(description="Whether to sort results in ascending order")] = True, + limit: Annotated[int | None, Query(description="The maximum number of results to return")] = None, ) -> AssetEventsResponse: if name and uri: where_clause = and_(AssetModel.name == name, AssetModel.uri == uri) @@ -91,10 +98,17 @@ def get_asset_event_by_asset_name_uri( }, ) + if after: + where_clause = and_(where_clause, AssetEvent.timestamp >= after) + if before: + where_clause = and_(where_clause, AssetEvent.timestamp <= before) + return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.asset, where_clause=where_clause, session=session, + ascending=ascending, + limit=limit, ) @@ -102,9 +116,21 @@ def get_asset_event_by_asset_name_uri( def get_asset_event_by_asset_alias( name: Annotated[str, Query(description="The name of the Asset Alias")], session: SessionDep, + after: Annotated[UtcDateTime | None, Query(description="The start of the time range")] = None, + before: Annotated[UtcDateTime | None, Query(description="The end of the time range")] = None, + ascending: Annotated[bool, Query(description="Whether to sort results in ascending order")] = True, + limit: Annotated[int | None, Query(description="The maximum number of results to return")] = None, ) -> AssetEventsResponse: + where_clause = AssetAliasModel.name == name + if after: + where_clause = and_(where_clause, AssetEvent.timestamp >= after) + if before: + where_clause = and_(where_clause, AssetEvent.timestamp <= before) + return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.source_aliases, - where_clause=(AssetAliasModel.name == name), + where_clause=where_clause, session=session, + ascending=ascending, + limit=limit, ) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py index 4bc313d85d96c..f2879d8375677 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py @@ -17,6 +17,8 @@ from __future__ import annotations +from datetime import datetime + import pytest from airflow._shared.timezones import timezone @@ -29,6 +31,9 @@ @pytest.fixture def test_asset_events(session): + def make_timestamp(day): + return datetime(2021, 1, day, tzinfo=timezone.utc) + common = { "asset_id": 1, "extra": {"foo": "bar"}, @@ -38,7 +43,7 @@ def test_asset_events(session): "source_map_index": -1, } - events = [AssetEvent(id=i, timestamp=DEFAULT_DATE, **common) for i in (1, 2)] + events = [AssetEvent(id=i, timestamp=make_timestamp(i), **common) for i in (1, 2, 3)] session.add_all(events) session.commit() yield events @@ -132,11 +137,312 @@ def test_get_by_asset(self, uri, name, client): "uri": "s3://bucket/key", }, "created_dagruns": [], + "timestamp": "2021-01-02T00:00:00Z", + }, + { + "id": 3, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-03T00:00:00Z", + }, + ] + } + + @pytest.mark.parametrize( + "uri, name", + [ + (None, "test_get_asset_by_name"), + ("s3://bucket/key", None), + ("s3://bucket/key", "test_get_asset_by_name"), + ], + ) + @pytest.mark.usefixtures("test_asset", "test_asset_events") + def test_get_by_asset_with_after_filter(self, uri, name, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": name, "uri": uri, "after": "2021-01-02T00:00:00Z"}, + ) + assert response.status_code == 200 + assert response.json() == { + "asset_events": [ + { + "id": 2, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-02T00:00:00Z", + }, + { + "id": 3, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-03T00:00:00Z", + }, + ] + } + + @pytest.mark.parametrize( + "uri, name", + [ + (None, "test_get_asset_by_name"), + ("s3://bucket/key", None), + ("s3://bucket/key", "test_get_asset_by_name"), + ], + ) + @pytest.mark.usefixtures("test_asset", "test_asset_events") + def test_get_by_asset_with_before_filter(self, uri, name, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": name, "uri": uri, "before": "2021-01-02T00:00:00Z"}, + ) + assert response.status_code == 200 + assert response.json() == { + "asset_events": [ + { + "id": 1, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-01T00:00:00Z", + }, + { + "id": 2, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-02T00:00:00Z", + }, + ] + } + + @pytest.mark.parametrize( + "uri, name", + [ + (None, "test_get_asset_by_name"), + ("s3://bucket/key", None), + ("s3://bucket/key", "test_get_asset_by_name"), + ], + ) + @pytest.mark.usefixtures("test_asset", "test_asset_events") + def test_get_by_asset_with_before_and_after_filters(self, uri, name, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": name, + "uri": uri, + "before": "2021-01-02T12:00:00Z", + "after": "2021-01-01T12:00:00Z", + }, + ) + assert response.status_code == 200 + assert response.json() == { + "asset_events": [ + { + "id": 2, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-02T00:00:00Z", + }, + ] + } + + @pytest.mark.parametrize( + "uri, name", + [ + (None, "test_get_asset_by_name"), + ("s3://bucket/key", None), + ("s3://bucket/key", "test_get_asset_by_name"), + ], + ) + @pytest.mark.usefixtures("test_asset", "test_asset_events") + def test_get_by_asset_with_descending_order(self, uri, name, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": name, "uri": uri, "ascending": False}, + ) + assert response.status_code == 200 + assert response.json() == { + "asset_events": [ + { + "id": 3, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-03T00:00:00Z", + }, + { + "id": 2, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-02T00:00:00Z", + }, + { + "id": 1, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-01T00:00:00Z", + }, + ] + } + + @pytest.mark.parametrize( + "uri, name", + [ + (None, "test_get_asset_by_name"), + ("s3://bucket/key", None), + ("s3://bucket/key", "test_get_asset_by_name"), + ], + ) + @pytest.mark.usefixtures("test_asset", "test_asset_events") + def test_get_by_asset_get_first(self, uri, name, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": name, "uri": uri, "limit": 1}, + ) + assert response.status_code == 200 + assert response.json() == { + "asset_events": [ + { + "id": 1, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], "timestamp": "2021-01-01T00:00:00Z", }, ] } + @pytest.mark.parametrize( + "uri, name", + [ + (None, "test_get_asset_by_name"), + ("s3://bucket/key", None), + ("s3://bucket/key", "test_get_asset_by_name"), + ], + ) + @pytest.mark.usefixtures("test_asset", "test_asset_events") + def test_get_by_asset_get_last(self, uri, name, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": name, "uri": uri, "limit": 1, "ascending": False}, + ) + assert response.status_code == 200 + assert response.json() == { + "asset_events": [ + { + "id": 3, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-03T00:00:00Z", + }, + ] + } + class TestGetAssetEventByAssetAlias: @pytest.mark.usefixtures("test_asset_alias") @@ -178,7 +484,23 @@ def test_get_by_asset(self, client): "uri": "s3://bucket/key", }, "created_dagruns": [], - "timestamp": "2021-01-01T00:00:00Z", + "timestamp": "2021-01-02T00:00:00Z", + }, + { + "id": 3, + "extra": {"foo": "bar"}, + "source_task_id": "bar", + "source_dag_id": "foo", + "source_run_id": "custom", + "source_map_index": -1, + "asset": { + "extra": {"foo": "bar"}, + "group": "asset", + "name": "test_get_asset_by_name", + "uri": "s3://bucket/key", + }, + "created_dagruns": [], + "timestamp": "2021-01-03T00:00:00Z", }, ] } diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 04575fa770d05..3b023e4666a58 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -605,13 +605,32 @@ def __init__(self, client: Client): self.client = client def get( - self, name: str | None = None, uri: str | None = None, alias_name: str | None = None + self, + name: str | None = None, + uri: str | None = None, + alias_name: str | None = None, + after: datetime | None = None, + before: datetime | None = None, + ascending: bool = True, + limit: int | None = None, ) -> AssetEventsResponse: """Get Asset event from the API server.""" + common_params: dict[str, Any] = {} + if after: + common_params["after"] = after.isoformat() + if before: + common_params["before"] = before.isoformat() + common_params["ascending"] = ascending + if limit: + common_params["limit"] = limit if name or uri: - resp = self.client.get("asset-events/by-asset", params={"name": name, "uri": uri}) + resp = self.client.get( + "asset-events/by-asset", params={"name": name, "uri": uri, **common_params} + ) elif alias_name: - resp = self.client.get("asset-events/by-asset-alias", params={"name": alias_name}) + resp = self.client.get( + "asset-events/by-asset-alias", params={"name": alias_name, **common_params} + ) else: raise ValueError("Either `name`, `uri` or `alias_name` must be provided") diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index bfa9fb012ae17..ca7dcefa02419 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -824,11 +824,19 @@ class GetAssetByUri(BaseModel): class GetAssetEventByAsset(BaseModel): name: str | None uri: str | None + after: AwareDatetime | None = None + before: AwareDatetime | None = None + limit: int | None = None + ascending: bool = True type: Literal["GetAssetEventByAsset"] = "GetAssetEventByAsset" class GetAssetEventByAssetAlias(BaseModel): alias_name: str + after: AwareDatetime | None = None + before: AwareDatetime | None = None + limit: int | None = None + ascending: bool = True type: Literal["GetAssetEventByAssetAlias"] = "GetAssetEventByAssetAlias" diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 570cd25d9a3ef..9e406dd83f313 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -18,7 +18,9 @@ import collections import contextlib +import functools from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence +from datetime import datetime from functools import cache from typing import TYPE_CHECKING, Any, Generic, TypeVar, overload @@ -60,6 +62,7 @@ VariableResult, ) from airflow.sdk.types import OutletEventAccessorsProtocol + from airflow.typing_compat import Self DEFAULT_FORMAT_PREFIX = "airflow.ctx." @@ -503,10 +506,106 @@ def __getitem__(self, key: Asset | AssetAlias | AssetRef) -> OutletEventAccessor return self._dict[hashable_key] +@attrs.define(init=False) +class InletEventsAccessor(Sequence["AssetEventResult"]): + _after: str | datetime | None + _before: str | datetime | None + _ascending: bool + _limit: int | None + _asset_name: str | None + _asset_uri: str | None + _alias_name: str | None + + def __init__( + self, asset_name: str | None = None, asset_uri: str | None = None, alias_name: str | None = None + ): + self._asset_name = asset_name + self._asset_uri = asset_uri + self._alias_name = alias_name + self._after = None + self._before = None + self._ascending = True + self._limit = None + + def after(self, after: str) -> Self: + self._after = after + self._reset_cache() + return self + + def before(self, before: str) -> Self: + self._before = before + self._reset_cache() + return self + + def ascending(self, ascending: bool = True) -> Self: + self._ascending = ascending + self._reset_cache() + return self + + def limit(self, limit: int) -> Self: + self._limit = limit + self._reset_cache() + return self + + @functools.cached_property + def _asset_events(self) -> list[AssetEventResult]: + from airflow.sdk.execution_time.comms import ( + ErrorResponse, + GetAssetEventByAsset, + GetAssetEventByAssetAlias, + ToSupervisor, + ) + from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS + + query_dict: dict[str, Any] = { + "after": self._after, + "before": self._before, + "ascending": self._ascending, + "limit": self._limit, + } + + msg: ToSupervisor + if self._alias_name is not None: + msg = GetAssetEventByAssetAlias(alias_name=self._alias_name, **query_dict) + else: + if self._asset_name is None and self._asset_uri is None: + raise ValueError("Either asset_name or asset_uri must be provided") + msg = GetAssetEventByAsset(name=self._asset_name, uri=self._asset_uri, **query_dict) + resp = SUPERVISOR_COMMS.send(msg) + if isinstance(resp, ErrorResponse): + raise AirflowRuntimeError(resp) + + if TYPE_CHECKING: + assert isinstance(resp, AssetEventsResult) + + return list(resp.iter_asset_event_results()) + + def _reset_cache(self) -> None: + try: + del self._asset_events + except AttributeError: + pass + + def __iter__(self) -> Iterator[AssetEventResult]: + return iter(self._asset_events) + + def __len__(self) -> int: + return len(self._asset_events) + + @overload + def __getitem__(self, key: int) -> AssetEventResult: ... + + @overload + def __getitem__(self, key: slice) -> Sequence[AssetEventResult]: ... + + def __getitem__(self, key: int | slice) -> AssetEventResult | Sequence[AssetEventResult]: + return self._asset_events[key] + + @attrs.define(init=False) class InletEventsAccessors( Mapping["int | Asset | AssetAlias | AssetRef", Any], - _AssetEventAccessorsMixin[list["AssetEventResult"]], + _AssetEventAccessorsMixin[Sequence["AssetEventResult"]], ): """Lazy mapping of inlet asset event accessors.""" @@ -537,17 +636,12 @@ def __iter__(self) -> Iterator[Asset | AssetAlias]: def __len__(self) -> int: return len(self._inlets) - def __getitem__(self, key: int | Asset | AssetAlias | AssetRef) -> list[AssetEventResult]: + def __getitem__( + self, + key: int | Asset | AssetAlias | AssetRef, + ) -> InletEventsAccessor: from airflow.sdk.definitions.asset import Asset - from airflow.sdk.execution_time.comms import ( - ErrorResponse, - GetAssetEventByAsset, - GetAssetEventByAssetAlias, - ToSupervisor, - ) - from airflow.sdk.execution_time.task_runner import SUPERVISOR_COMMS - msg: ToSupervisor if isinstance(key, int): # Support index access; it's easier for trivial cases. obj = self._inlets[key] if not isinstance(obj, (Asset, AssetAlias, AssetRef)): @@ -557,33 +651,23 @@ def __getitem__(self, key: int | Asset | AssetAlias | AssetRef) -> list[AssetEve if isinstance(obj, Asset): asset = self._assets[AssetUniqueKey.from_asset(obj)] - msg = GetAssetEventByAsset(name=asset.name, uri=asset.uri) - elif isinstance(obj, AssetNameRef): + return InletEventsAccessor(asset_name=asset.name, asset_uri=asset.uri) + if isinstance(obj, AssetNameRef): try: asset = next(a for k, a in self._assets.items() if k.name == obj.name) except StopIteration: raise KeyError(obj) from None - msg = GetAssetEventByAsset(name=asset.name, uri=None) - elif isinstance(obj, AssetUriRef): + return InletEventsAccessor(asset_name=asset.name) + if isinstance(obj, AssetUriRef): try: asset = next(a for k, a in self._assets.items() if k.uri == obj.uri) except StopIteration: raise KeyError(obj) from None - msg = GetAssetEventByAsset(name=None, uri=asset.uri) - elif isinstance(obj, AssetAlias): + return InletEventsAccessor(asset_uri=asset.uri) + if isinstance(obj, AssetAlias): asset_alias = self._asset_aliases[AssetAliasUniqueKey.from_asset_alias(obj)] - msg = GetAssetEventByAssetAlias(alias_name=asset_alias.name) - else: - raise TypeError(f"`key` is of unknown type ({type(key).__name__})") - - resp = SUPERVISOR_COMMS.send(msg) - if isinstance(resp, ErrorResponse): - raise AirflowRuntimeError(resp) - - if TYPE_CHECKING: - assert isinstance(resp, AssetEventsResult) - - return list(resp.iter_asset_event_results()) + return InletEventsAccessor(alias_name=asset_alias.name) + raise TypeError(f"`key` is of unknown type ({type(key).__name__})") @attrs.define diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 440443b8c9c36..b6daa0cb9e47e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1258,12 +1258,25 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: else: resp = asset_resp elif isinstance(msg, GetAssetEventByAsset): - asset_event_resp = self.client.asset_events.get(uri=msg.uri, name=msg.name) + asset_event_resp = self.client.asset_events.get( + uri=msg.uri, + name=msg.name, + after=msg.after, + before=msg.before, + ascending=msg.ascending, + limit=msg.limit, + ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) resp = asset_event_result dump_opts = {"exclude_unset": True} elif isinstance(msg, GetAssetEventByAssetAlias): - asset_event_resp = self.client.asset_events.get(alias_name=msg.alias_name) + asset_event_resp = self.client.asset_events.get( + alias_name=msg.alias_name, + after=msg.after, + before=msg.before, + ascending=msg.ascending, + limit=msg.limit, + ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) resp = asset_event_result dump_opts = {"exclude_unset": True} diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py b/task-sdk/tests/task_sdk/execution_time/test_context.py index 54e2c66bee82f..9c1b394d6e89c 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -654,13 +654,98 @@ def test__get_item__(self, key, sample_inlet_evnets_accessor, mock_supervisor_co events_result = AssetEventsResult(asset_events=[asset_event_resp]) mock_supervisor_comms.send.side_effect = [events_result] * 4 - assert sample_inlet_evnets_accessor[key] == [asset_event_resp] + assert list(sample_inlet_evnets_accessor[key]) == [asset_event_resp] @pytest.mark.usefixtures("mock_supervisor_comms") def test__get_item__out_of_index(self, sample_inlet_evnets_accessor): with pytest.raises(IndexError): sample_inlet_evnets_accessor[5] + def test__get_item__with_filters(self, sample_inlet_evnets_accessor, mock_supervisor_comms): + asset_event_resp = AssetEventResult( + id=1, + created_dagruns=[], + timestamp=timezone.utcnow(), + asset=AssetResponse(name="test_uri", uri="test_uri", group="asset"), + ) + events_result = AssetEventsResult(asset_events=[asset_event_resp]) + mock_supervisor_comms.send.side_effect = [events_result] * 10 + + list(sample_inlet_evnets_accessor[TEST_ASSET]) + list(sample_inlet_evnets_accessor[TEST_ASSET].after("2024-01-01T00:00:00Z")) + list(sample_inlet_evnets_accessor[TEST_ASSET].before("2024-01-01T00:00:00Z")) + list(sample_inlet_evnets_accessor[TEST_ASSET].limit(10)) + list( + sample_inlet_evnets_accessor[TEST_ASSET] + .after("2024-01-01T00:00:00Z") + .before("2024-01-02T00:00:00Z") + .limit(10) + ) + list(sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)) + + assert mock_supervisor_comms.send.call_count == 6 + + # test accessing the accessor without list() or [] + sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10) + + assert mock_supervisor_comms.send.call_count == 6 + + # test accessing one of the elements + res = sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10)[0] + assert res == asset_event_resp + assert mock_supervisor_comms.send.call_count == 7 + + # test evaluating the accessor multiple times with the same filters + res = sample_inlet_evnets_accessor[TEST_ASSET].ascending(False).limit(10) + assert res[0] == asset_event_resp + assert res[0] == asset_event_resp + + assert mock_supervisor_comms.send.call_count == 8 + + # test changing one of the filters + assert res.after("2024-01-01T00:00:00Z")[0] == asset_event_resp + + assert mock_supervisor_comms.send.call_count == 9 + + # test len() + assert len(sample_inlet_evnets_accessor[TEST_ASSET].ascending(True).limit(10)) == 1 + assert mock_supervisor_comms.send.call_count == 10 + + calls = mock_supervisor_comms.send.call_args_list + assert calls[0][0][0] == GetAssetEventByAsset( + name="test_uri", uri="test://test/", after=None, before=None, limit=None, ascending=True + ) + assert calls[1][0][0] == GetAssetEventByAsset( + name="test_uri", + uri="test://test/", + after="2024-01-01T00:00:00Z", + before=None, + limit=None, + ascending=True, + ) + assert calls[2][0][0] == GetAssetEventByAsset( + name="test_uri", + uri="test://test/", + after=None, + before="2024-01-01T00:00:00Z", + limit=None, + ascending=True, + ) + assert calls[3][0][0] == GetAssetEventByAsset( + name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=True + ) + assert calls[4][0][0] == GetAssetEventByAsset( + name="test_uri", + uri="test://test/", + after="2024-01-01T00:00:00Z", + before="2024-01-02T00:00:00Z", + limit=10, + ascending=True, + ) + assert calls[5][0][0] == GetAssetEventByAsset( + name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=False + ) + @pytest.mark.parametrize( "name, uri, expected_key", ( @@ -702,8 +787,17 @@ def test_source_task_instance_xcom_pull(self, sample_inlet_evnets_accessor, mock ], ) ] - events = sample_inlet_evnets_accessor[Asset.ref(name="test_uri")] - mock_supervisor_comms.send.assert_called_once_with(GetAssetEventByAsset(name="test_uri", uri=None)) + events = list(sample_inlet_evnets_accessor[Asset.ref(name="test_uri")]) + mock_supervisor_comms.send.assert_called_once_with( + GetAssetEventByAsset( + name="test_uri", + uri=None, + after=None, + before=None, + limit=None, + ascending=True, + ) + ) assert len(events) == 2 assert events[1].source_task_instance is None diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 1c04adafcc32d..33af0105563b2 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -29,6 +29,7 @@ import time from contextlib import nullcontext from dataclasses import dataclass, field +from datetime import datetime from operator import attrgetter from random import randint from time import sleep @@ -1629,7 +1630,14 @@ class RequestTestCase: }, client_mock=ClientMock( method_path="asset_events.get", - kwargs={"uri": "s3://bucket/obj", "name": "test"}, + kwargs={ + "uri": "s3://bucket/obj", + "name": "test", + "after": None, + "before": None, + "limit": None, + "ascending": True, + }, response=AssetEventsResult( asset_events=[ AssetEventResponse( @@ -1643,6 +1651,49 @@ class RequestTestCase: ), test_id="get_asset_events_by_uri_and_name", ), + RequestTestCase( + message=GetAssetEventByAsset( + uri="s3://bucket/obj", + name="test", + after=datetime(2024, 10, 1, 12, 0, 0, tzinfo=timezone.utc), + before=datetime(2024, 10, 15, 12, 0, 0, tzinfo=timezone.utc), + limit=5, + ascending=False, + ), + expected_body={ + "asset_events": [ + { + "id": 1, + "timestamp": timezone.parse("2024-10-31T12:00:00Z"), + "asset": {"name": "asset", "uri": "s3://bucket/obj", "group": "asset"}, + "created_dagruns": [], + } + ], + "type": "AssetEventsResult", + }, + client_mock=ClientMock( + method_path="asset_events.get", + kwargs={ + "uri": "s3://bucket/obj", + "name": "test", + "after": timezone.parse("2024-10-01T12:00:00Z"), + "before": timezone.parse("2024-10-15T12:00:00Z"), + "limit": 5, + "ascending": False, + }, + response=AssetEventsResult( + asset_events=[ + AssetEventResponse( + id=1, + asset=AssetResponse(name="asset", uri="s3://bucket/obj", group="asset"), + created_dagruns=[], + timestamp=timezone.parse("2024-10-31T12:00:00Z"), + ), + ], + ), + ), + test_id="get_asset_events_by_uri_and_name_with_filters", + ), RequestTestCase( message=GetAssetEventByAsset(uri="s3://bucket/obj", name=None), expected_body={ @@ -1658,7 +1709,14 @@ class RequestTestCase: }, client_mock=ClientMock( method_path="asset_events.get", - kwargs={"uri": "s3://bucket/obj", "name": None}, + kwargs={ + "uri": "s3://bucket/obj", + "name": None, + "after": None, + "before": None, + "limit": None, + "ascending": True, + }, response=AssetEventsResult( asset_events=[ AssetEventResponse( @@ -1672,6 +1730,49 @@ class RequestTestCase: ), test_id="get_asset_events_by_uri", ), + RequestTestCase( + message=GetAssetEventByAsset( + uri="s3://bucket/obj", + name=None, + after=datetime(2024, 10, 1, 12, 0, 0, tzinfo=timezone.utc), + before=datetime(2024, 10, 15, 12, 0, 0, tzinfo=timezone.utc), + limit=5, + ascending=False, + ), + expected_body={ + "asset_events": [ + { + "id": 1, + "timestamp": timezone.parse("2024-10-31T12:00:00Z"), + "asset": {"name": "asset", "uri": "s3://bucket/obj", "group": "asset"}, + "created_dagruns": [], + } + ], + "type": "AssetEventsResult", + }, + client_mock=ClientMock( + method_path="asset_events.get", + kwargs={ + "uri": "s3://bucket/obj", + "name": None, + "after": timezone.parse("2024-10-01T12:00:00Z"), + "before": timezone.parse("2024-10-15T12:00:00Z"), + "limit": 5, + "ascending": False, + }, + response=AssetEventsResult( + asset_events=[ + AssetEventResponse( + id=1, + asset=AssetResponse(name="asset", uri="s3://bucket/obj", group="asset"), + created_dagruns=[], + timestamp=timezone.parse("2024-10-31T12:00:00Z"), + ) + ], + ), + ), + test_id="get_asset_events_by_uri_with_filters", + ), RequestTestCase( message=GetAssetEventByAsset(uri=None, name="test"), expected_body={ @@ -1687,7 +1788,14 @@ class RequestTestCase: }, client_mock=ClientMock( method_path="asset_events.get", - kwargs={"uri": None, "name": "test"}, + kwargs={ + "uri": None, + "name": "test", + "after": None, + "before": None, + "limit": None, + "ascending": True, + }, response=AssetEventsResult( asset_events=[ AssetEventResponse( @@ -1701,6 +1809,49 @@ class RequestTestCase: ), test_id="get_asset_events_by_name", ), + RequestTestCase( + message=GetAssetEventByAsset( + uri=None, + name="test", + after=datetime(2024, 10, 1, 12, 0, 0, tzinfo=timezone.utc), + before=datetime(2024, 10, 15, 12, 0, 0, tzinfo=timezone.utc), + limit=5, + ascending=False, + ), + expected_body={ + "asset_events": [ + { + "id": 1, + "timestamp": timezone.parse("2024-10-31T12:00:00Z"), + "asset": {"name": "asset", "uri": "s3://bucket/obj", "group": "asset"}, + "created_dagruns": [], + } + ], + "type": "AssetEventsResult", + }, + client_mock=ClientMock( + method_path="asset_events.get", + kwargs={ + "uri": None, + "name": "test", + "after": timezone.parse("2024-10-01T12:00:00Z"), + "before": timezone.parse("2024-10-15T12:00:00Z"), + "limit": 5, + "ascending": False, + }, + response=AssetEventsResult( + asset_events=[ + AssetEventResponse( + id=1, + asset=AssetResponse(name="asset", uri="s3://bucket/obj", group="asset"), + created_dagruns=[], + timestamp=timezone.parse("2024-10-31T12:00:00Z"), + ) + ] + ), + ), + test_id="get_asset_events_by_name_with_filters", + ), RequestTestCase( message=GetAssetEventByAssetAlias(alias_name="test_alias"), expected_body={ @@ -1716,7 +1867,13 @@ class RequestTestCase: }, client_mock=ClientMock( method_path="asset_events.get", - kwargs={"alias_name": "test_alias"}, + kwargs={ + "alias_name": "test_alias", + "after": None, + "before": None, + "limit": None, + "ascending": True, + }, response=AssetEventsResult( asset_events=[ AssetEventResponse( @@ -1730,6 +1887,47 @@ class RequestTestCase: ), test_id="get_asset_events_by_asset_alias", ), + RequestTestCase( + message=GetAssetEventByAssetAlias( + alias_name="test_alias", + after=datetime(2024, 10, 1, 12, 0, 0, tzinfo=timezone.utc), + before=datetime(2024, 10, 15, 12, 0, 0, tzinfo=timezone.utc), + limit=5, + ascending=False, + ), + expected_body={ + "asset_events": [ + { + "id": 1, + "timestamp": timezone.parse("2024-10-31T12:00:00Z"), + "asset": {"name": "asset", "uri": "s3://bucket/obj", "group": "asset"}, + "created_dagruns": [], + } + ], + "type": "AssetEventsResult", + }, + client_mock=ClientMock( + method_path="asset_events.get", + kwargs={ + "alias_name": "test_alias", + "after": timezone.parse("2024-10-01T12:00:00Z"), + "before": timezone.parse("2024-10-15T12:00:00Z"), + "limit": 5, + "ascending": False, + }, + response=AssetEventsResult( + asset_events=[ + AssetEventResponse( + id=1, + asset=AssetResponse(name="asset", uri="s3://bucket/obj", group="asset"), + created_dagruns=[], + timestamp=timezone.parse("2024-10-31T12:00:00Z"), + ) + ] + ), + ), + test_id="get_asset_events_by_asset_alias_with_filters", + ), RequestTestCase( message=ValidateInletsAndOutlets(ti_id=TI_ID), expected_body={ 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 841137c07b3c8..4606771520416 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 @@ -1137,14 +1137,14 @@ def test_run_with_asset_inlets(create_runtime_ti, mock_supervisor_comms): inlet_events = ti.get_template_context()["inlet_events"] # access the asset events of Asset(name="test", uri="test://uri") - assert inlet_events[0] == [asset_event_resp] - assert inlet_events[-2] == [asset_event_resp] - assert inlet_events[Asset(name="test", uri="test://uri")] == [asset_event_resp] + assert list(inlet_events[0]) == [asset_event_resp] + assert list(inlet_events[-2]) == [asset_event_resp] + assert list(inlet_events[Asset(name="test", uri="test://uri")]) == [asset_event_resp] # access the asset events of AssetAlias(name="alias-name") - assert inlet_events[1] == [asset_event_resp] - assert inlet_events[-1] == [asset_event_resp] - assert inlet_events[AssetAlias(name="alias-name")] == [asset_event_resp] + assert list(inlet_events[1]) == [asset_event_resp] + assert list(inlet_events[-1]) == [asset_event_resp] + assert list(inlet_events[AssetAlias(name="alias-name")]) == [asset_event_resp] # access with invalid index with pytest.raises(IndexError):