diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index b5561392587db..2a7990d233bd1 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -247,6 +247,26 @@ Inlet asset events can be read with the ``inlet_events`` accessor in the executi Each value in the ``inlet_events`` mapping is a sequence-like object that orders past events of a given asset by ``timestamp``, earliest to latest. It supports most of Python's list interface, so you can use ``[-1]`` to access the last event, ``[-2:]`` for the last two, etc. The accessor is lazy and only hits the database when you access items inside it. +The accessor also supports chaining methods to filter events before fetching them. For example, to retrieve only events where specific ``extra`` keys match given values: + +.. code-block:: python + + @task(inlets=[regional_sales]) + def process_high_priority(*, inlet_events): + events = inlet_events[regional_sales].extra("region", "us").extra("env", "prod") + for event in events: + print(event.extra) + +The ``.extra(key, value)`` method can be chained multiple times; all conditions are combined with AND logic. Other chaining methods include ``.after(timestamp)``, ``.before(timestamp)``, ``.ascending()``, and ``.limit(n)``. + +The REST API also supports filtering asset events by extra key-value pairs using the ``extra`` query parameter (repeat for multiple conditions): + +.. code-block:: bash + + curl "http:///api/v2/assets/events?extra=region%3Dus&extra=env%3Dprod" + +Each ``extra`` value uses ``key=value`` format. Multiple entries are combined with AND logic. + Dependency between ``@asset``, ``@task``, and classic operators --------------------------------------------------------------- diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index d2cfde623523f..975d35981feb5 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``d2f4e1b3c5a7`` (head) | ``9ff64e1c35d3`` | ``3.3.0`` | Add partition_date to asset_partition_dag_run. | +| ``5a5d3253e946`` (head) | ``d2f4e1b3c5a7`` | ``3.4.0`` | Add GIN index on asset_event.extra for PostgreSQL. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``d2f4e1b3c5a7`` | ``9ff64e1c35d3`` | ``3.3.0`` | Add partition_date to asset_partition_dag_run. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``9ff64e1c35d3`` | ``dd5f3a8e2b91`` | ``3.3.0`` | Add indexes on dag_run.created_dag_version_id and | | | | | task_instance.dag_version_id. | diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 935b547c1a7cf..7b235f07de306 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -68,6 +68,7 @@ from airflow.models.variable import Variable from airflow.models.xcom import XComModel from airflow.typing_compat import Self +from airflow.utils.sqlalchemy import JsonContains from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunType @@ -553,6 +554,60 @@ def depends_prefix_search( return depends_prefix_search +class _JsonKVFilter(BaseParam[dict[str, str]]): + """ + Filter on a JSON column by multiple key-value pairs (AND logic). + + Uses dialect-aware SQL: ``@>`` (JSONB containment, GIN-indexable) on + PostgreSQL, ``JSON_CONTAINS`` on MySQL, and ``JSON_EXTRACT`` on SQLite. + """ + + def __init__( + self, + attribute: ColumnElement, + value: dict[str, str] | None = None, + skip_none: bool = True, + ) -> None: + super().__init__(skip_none=skip_none) + self.attribute: ColumnElement = attribute + self.value = value + + def to_orm(self, select: Select) -> Select: + if not self.value: + return select + return select.where(JsonContains(self.attribute, self.value)) + + @classmethod + def depends(cls, *args: Any, **kwargs: Any) -> Self: + raise NotImplementedError("Use json_kv_filter_factory instead.") + + +def json_kv_filter_factory( + attribute: ColumnElement, + param_name: str = "extra", +) -> Callable[[list[str]], _JsonKVFilter]: + DESCRIPTION = ( + "Filter by JSON key-value pairs. Repeat for multiple conditions (AND logic). " + "Format: key=value (e.g. extra=region=us&extra=env=prod)." + ) + + def depends_json_kv( + values: list[str] = Query(alias=param_name, default_factory=list, description=DESCRIPTION), + ) -> _JsonKVFilter: + kv_dict: dict[str, str] = {} + for item in values: + if "=" not in item: + raise HTTPException( + status_code=HTTP_422_UNPROCESSABLE_CONTENT, + detail=f"Invalid {param_name} parameter format: {item!r}. Expected 'key=value'.", + ) + k, v = item.split("=", 1) + kv_dict[k] = v + return _JsonKVFilter(attribute, kv_dict or None) + + return depends_json_kv + + class SortParam(BaseParam[list[str]]): """Order result by the attribute.""" @@ -1587,6 +1642,7 @@ def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | N QueryAssetDagIdPatternSearch = Annotated[ _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter.depends) ] +QueryAssetEventExtraFilter = Annotated[_JsonKVFilter, Depends(json_kv_filter_factory(AssetEvent.extra))] QueryPartitionedDagRunHasCreatedDagRunIdFilter = Annotated[ FilterParam[bool | None], Depends( diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 44f797e1d3c36..be7f2d1f936b4 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -478,6 +478,18 @@ paths: \ stays index-compatible under locale-aware collations \u2014 e.g. `test_`\ \ effectively matches items starting with `test`, and `s3://` matches items\ \ starting with `s3`." + - name: extra + in: query + required: false + schema: + type: array + items: + type: string + description: 'Filter by JSON key-value pairs. Repeat for multiple conditions + (AND logic). Format: key=value (e.g. extra=region=us&extra=env=prod).' + title: Extra + description: 'Filter by JSON key-value pairs. Repeat for multiple conditions + (AND logic). Format: key=value (e.g. extra=region=us&extra=env=prod).' - name: timestamp_gte in: query required: false diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py index 15af36445dd95..44f1e3e196741 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py @@ -37,6 +37,7 @@ QueryAssetAliasNamePatternSearch, QueryAssetAliasNamePrefixPatternSearch, QueryAssetDagIdPatternSearch, + QueryAssetEventExtraFilter, QueryAssetNamePatternSearch, QueryAssetNamePrefixPatternSearch, QueryLimit, @@ -321,6 +322,7 @@ def get_asset_events( ], name_pattern: QueryAssetNamePatternSearch, name_prefix_pattern: QueryAssetNamePrefixPatternSearch, + extra_filter: QueryAssetEventExtraFilter, timestamp_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("timestamp", AssetEvent))], session: SessionDep, ) -> AssetEventCollectionResponse: @@ -339,6 +341,7 @@ def get_asset_events( source_map_index, name_pattern, name_prefix_pattern, + extra_filter, timestamp_range, ], order_by=order_by, 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 d0ecc3d3adaf2..b2a98a9de60b7 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 @@ -30,6 +30,7 @@ AssetEventsResponse, ) from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel +from airflow.utils.sqlalchemy import JsonContains router = APIRouter( responses={ @@ -44,7 +45,7 @@ def _get_asset_events_through_sql_clauses( ) -> AssetEventsResponse: 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: + if limit is not None: asset_events_query = asset_events_query.limit(limit) asset_events = session.scalars(asset_events_query) return AssetEventsResponse.model_validate( @@ -73,6 +74,23 @@ def _get_asset_events_through_sql_clauses( ) +def _parse_extra_params(extra: list[str] | None) -> dict[str, str]: + """Parse repeated ``key=value`` query params into a dict.""" + result: dict[str, str] = {} + for item in extra or []: + if "=" not in item: + raise HTTPException( + status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, + detail={ + "reason": "Invalid parameter", + "message": f"Invalid extra parameter format: {item!r}. Expected 'key=value'.", + }, + ) + k, v = item.split("=", 1) + result[k] = v + return result + + @router.get("/by-asset") def get_asset_event_by_asset_name_uri( name: Annotated[str | None, Query(description="The name of the Asset")], @@ -82,6 +100,12 @@ def get_asset_event_by_asset_name_uri( 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, + extra: Annotated[ + list[str] | None, + Query( + description="Filter by extra JSON key-value pairs. Format: key=value. Repeat for AND logic.", + ), + ] = None, ) -> AssetEventsResponse: if name and uri: where_clause = and_(AssetModel.name == name, AssetModel.uri == uri) @@ -102,6 +126,9 @@ def get_asset_event_by_asset_name_uri( where_clause = and_(where_clause, AssetEvent.timestamp >= after) if before: where_clause = and_(where_clause, AssetEvent.timestamp <= before) + extra_dict = _parse_extra_params(extra) + if extra_dict: + where_clause = and_(where_clause, JsonContains(AssetEvent.extra, extra_dict)) return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.asset, @@ -120,12 +147,21 @@ def get_asset_event_by_asset_alias( 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, + extra: Annotated[ + list[str] | None, + Query( + description="Filter by extra JSON key-value pairs. Format: key=value. Repeat for AND logic.", + ), + ] = 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) + extra_dict = _parse_extra_params(extra) + if extra_dict: + where_clause = and_(where_clause, JsonContains(AssetEvent.extra, extra_dict)) return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.source_aliases, diff --git a/airflow-core/src/airflow/migrations/env.py b/airflow-core/src/airflow/migrations/env.py index dfb46ab6c37e0..f32a003a4151d 100644 --- a/airflow-core/src/airflow/migrations/env.py +++ b/airflow-core/src/airflow/migrations/env.py @@ -27,6 +27,12 @@ from airflow import models, settings from airflow.utils.db import compare_server_default, compare_type +_POSTGRES_ONLY_INDEXES = frozenset( + { + "idx_asset_event_extra_gin", + } +) + def include_object(_, name, type_, *args): """Filter objects for autogenerating revisions.""" @@ -36,6 +42,10 @@ def include_object(_, name, type_, *args): # Only create migrations for objects that are in the target metadata if type_ == "table" and name not in target_metadata.tables: return False + # Indexes created by raw SQL in migrations (e.g. Postgres GIN) are not + # represented in the SQLAlchemy model; hide them from autogenerate. + if type_ == "index" and name in _POSTGRES_ONLY_INDEXES: + return False return True diff --git a/airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py b/airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py new file mode 100644 index 0000000000000..bffa7eb0f0348 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0124_3_4_0_add_gin_index_on_asset_event_extra.py @@ -0,0 +1,57 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add GIN index on asset_event.extra for PostgreSQL. + +Revision ID: 5a5d3253e946 +Revises: d2f4e1b3c5a7 +Create Date: 2026-04-01 23:00:00.000000 + +""" + +from __future__ import annotations + +from alembic import op +from sqlalchemy import text + +# revision identifiers, used by Alembic. +revision = "5a5d3253e946" +down_revision = "d2f4e1b3c5a7" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add GIN index on asset_event.extra for PostgreSQL only.""" + conn = op.get_bind() + if conn.dialect.name == "postgresql": + op.execute( + text( + "CREATE INDEX IF NOT EXISTS idx_asset_event_extra_gin " + "ON asset_event USING GIN ((extra::jsonb) jsonb_ops)" + ) + ) + + +def downgrade(): + """Remove GIN index on asset_event.extra.""" + conn = op.get_bind() + if conn.dialect.name == "postgresql": + op.execute(text("DROP INDEX IF EXISTS idx_asset_event_extra_gin")) diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index 69372a087a36b..544907fe3843b 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -36,8 +36,9 @@ export const UseAssetServiceGetAssetAliasKeyFn = ({ assetAliasId }: { export type AssetServiceGetAssetEventsDefaultResponse = Awaited>; export type AssetServiceGetAssetEventsQueryResult = UseQueryResult; export const useAssetServiceGetAssetEventsKey = "AssetServiceGetAssetEvents"; -export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; + extra?: string[]; limit?: number; namePattern?: string; namePrefixPattern?: string; @@ -51,7 +52,7 @@ export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, limit, namePattern timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}, queryKey?: Array) => [useAssetServiceGetAssetEventsKey, ...(queryKey ?? [{ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }])]; +} = {}, queryKey?: Array) => [useAssetServiceGetAssetEventsKey, ...(queryKey ?? [{ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }])]; export type AssetServiceGetAssetQueuedEventsDefaultResponse = Awaited>; export type AssetServiceGetAssetQueuedEventsQueryResult = UseQueryResult; export const useAssetServiceGetAssetQueuedEventsKey = "AssetServiceGetAssetQueuedEvents"; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index ff0174b77e946..0f37dbf009f96 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -83,6 +83,7 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * * **Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and most of the time prevents the database from using B-tree indexes, which can be very slow on large tables. Prefer the equivalent ``name_prefix_pattern`` parameter when possible. * @param data.namePrefixPattern Prefix match — returns items whose value starts with the given string (case-sensitive, index-friendly). Use the pipe `|` operator for OR logic (e.g. `dag1|dag2`). Use `~` to match all. Wildcard characters (`%`, `_`) are treated as literal characters. Trailing non-alphanumeric characters in the prefix are stripped before matching so the range scan stays index-compatible under locale-aware collations — e.g. `test_` effectively matches items starting with `test`, and `s3://` matches items starting with `s3`. +* @param data.extra Filter by JSON key-value pairs. Repeat for multiple conditions (AND logic). Format: key=value (e.g. extra=region=us&extra=env=prod). * @param data.timestampGte * @param data.timestampGt * @param data.timestampLte @@ -90,8 +91,9 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * @returns AssetEventCollectionResponse Successful Response * @throws ApiError */ -export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient, { assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; + extra?: string[]; limit?: number; namePattern?: string; namePrefixPattern?: string; @@ -105,7 +107,7 @@ export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}) => queryClient.ensureQueryData({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) }); +} = {}) => queryClient.ensureQueryData({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) }); /** * Get Asset Queued Events * Get queued asset events for an asset. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index 1f8db278d9d3a..afef265e8884e 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -83,6 +83,7 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * * **Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and most of the time prevents the database from using B-tree indexes, which can be very slow on large tables. Prefer the equivalent ``name_prefix_pattern`` parameter when possible. * @param data.namePrefixPattern Prefix match — returns items whose value starts with the given string (case-sensitive, index-friendly). Use the pipe `|` operator for OR logic (e.g. `dag1|dag2`). Use `~` to match all. Wildcard characters (`%`, `_`) are treated as literal characters. Trailing non-alphanumeric characters in the prefix are stripped before matching so the range scan stays index-compatible under locale-aware collations — e.g. `test_` effectively matches items starting with `test`, and `s3://` matches items starting with `s3`. +* @param data.extra Filter by JSON key-value pairs. Repeat for multiple conditions (AND logic). Format: key=value (e.g. extra=region=us&extra=env=prod). * @param data.timestampGte * @param data.timestampGt * @param data.timestampLte @@ -90,8 +91,9 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * @returns AssetEventCollectionResponse Successful Response * @throws ApiError */ -export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, { assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; + extra?: string[]; limit?: number; namePattern?: string; namePrefixPattern?: string; @@ -105,7 +107,7 @@ export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) }); +} = {}) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) }); /** * Get Asset Queued Events * Get queued asset events for an asset. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index c5ec74825be02..c4bf24c745376 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -83,6 +83,7 @@ export const useAssetServiceGetAssetAlias = = unknown[]>({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const useAssetServiceGetAssetEvents = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; + extra?: string[]; limit?: number; namePattern?: string; namePrefixPattern?: string; @@ -105,7 +107,7 @@ export const useAssetServiceGetAssetEvents = , "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) as TData, ...options }); +} = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) as TData, ...options }); /** * Get Asset Queued Events * Get queued asset events for an asset. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index 90f75e1cfb359..2bf895cd5c992 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -83,6 +83,7 @@ export const useAssetServiceGetAssetAliasSuspense = = unknown[]>({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const useAssetServiceGetAssetEventsSuspense = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; + extra?: string[]; limit?: number; namePattern?: string; namePrefixPattern?: string; @@ -105,7 +107,7 @@ export const useAssetServiceGetAssetEventsSuspense = , "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) as TData, ...options }); +} = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }) as TData, ...options }); /** * Get Asset Queued Events * Get queued asset events for an asset. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index ebb30ffa68cd4..9d0af4e80c9ca 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -124,6 +124,7 @@ export class AssetService { * * **Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and most of the time prevents the database from using B-tree indexes, which can be very slow on large tables. Prefer the equivalent ``name_prefix_pattern`` parameter when possible. * @param data.namePrefixPattern Prefix match — returns items whose value starts with the given string (case-sensitive, index-friendly). Use the pipe `|` operator for OR logic (e.g. `dag1|dag2`). Use `~` to match all. Wildcard characters (`%`, `_`) are treated as literal characters. Trailing non-alphanumeric characters in the prefix are stripped before matching so the range scan stays index-compatible under locale-aware collations — e.g. `test_` effectively matches items starting with `test`, and `s3://` matches items starting with `s3`. + * @param data.extra Filter by JSON key-value pairs. Repeat for multiple conditions (AND logic). Format: key=value (e.g. extra=region=us&extra=env=prod). * @param data.timestampGte * @param data.timestampGt * @param data.timestampLte @@ -146,6 +147,7 @@ export class AssetService { source_map_index: data.sourceMapIndex, name_pattern: data.namePattern, name_prefix_pattern: data.namePrefixPattern, + extra: data.extra, timestamp_gte: data.timestampGte, timestamp_gt: data.timestampGt, timestamp_lte: data.timestampLte, diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index c070c9d74d71f..bee323f680b82 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -2816,6 +2816,10 @@ export type GetAssetAliasResponse = unknown; export type GetAssetEventsData = { assetId?: number | null; + /** + * Filter by JSON key-value pairs. Repeat for multiple conditions (AND logic). Format: key=value (e.g. extra=region=us&extra=env=prod). + */ + extra?: Array<(string)>; limit?: number; /** * SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index f24f3636132ae..f02ef037174f2 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,6 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", + "3.4.0": "5a5d3253e946", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index c47d8fd4796b0..095d0de4741ee 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -20,6 +20,7 @@ import contextlib import copy import datetime +import json import logging from collections.abc import Generator from typing import TYPE_CHECKING, Any @@ -28,8 +29,9 @@ from sqlalchemy.dialects import mysql from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.ext.compiler import compiles +from sqlalchemy.sql.expression import ColumnElement from sqlalchemy.sql.functions import FunctionElement -from sqlalchemy.types import JSON, Text, TypeDecorator +from sqlalchemy.types import JSON, NullType, Text, TypeDecorator from airflow._shared.timezones.timezone import make_naive, utc from airflow.configuration import conf @@ -45,7 +47,6 @@ from sqlalchemy.exc import OperationalError from sqlalchemy.orm import Session from sqlalchemy.sql import Select - from sqlalchemy.sql.elements import ColumnElement from sqlalchemy.types import TypeEngine from airflow.typing_compat import Self @@ -133,6 +134,56 @@ def _random_db_uuid_sqlite(element, compiler, **kw): return "uuid4()" +class JsonContains(ColumnElement): + """ + Dialect-aware JSON containment check. + + Compiles to ``@>`` on PostgreSQL (GIN-indexable), ``JSON_CONTAINS`` on + MySQL, and per-key ``json_extract`` comparisons on SQLite. + + All dialects use bound parameters to avoid SQL injection. + """ + + inherit_cache = False + type = NullType() + + def __init__(self, column, kv_dict: dict[str, str]): + self.column = column + self.kv_dict = kv_dict + + +@compiles(JsonContains, "postgresql") +def _pg_json_contains(element, compiler, **kw): + from sqlalchemy import cast, literal + + col = cast(element.column, JSONB) + param = literal(json.dumps(element.kv_dict)).cast(JSONB) + expr = col.contains(param) + return compiler.process(expr, **kw) + + +@compiles(JsonContains, "mysql") +def _mysql_json_contains(element, compiler, **kw): + from sqlalchemy import bindparam, func + + param = bindparam(None, json.dumps(element.kv_dict), expanding=False) + expr = func.JSON_CONTAINS(element.column, param) + return compiler.process(expr == 1, **kw) + + +@compiles(JsonContains) +def _default_json_contains(element, compiler, **kw): + from sqlalchemy import and_, func, literal + + clauses = [] + for k, v in element.kv_dict.items(): + path = f"$.{k}" + clauses.append(func.json_extract(element.column, literal(path)) == literal(v)) + if len(clauses) == 1: + return compiler.process(clauses[0], **kw) + return compiler.process(and_(*clauses), **kw) + + class UtcDateTime(TypeDecorator): """ Similar to :class:`~sqlalchemy.types.TIMESTAMP` with ``timezone=True`` option, with some differences. diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py index 33c43785a55a8..718dbfe772651 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py @@ -1078,6 +1078,87 @@ def test_should_mask_sensitive_extra(self, test_client, session): } +class TestGetAssetEventsExtraFilter(TestAssets): + @pytest.fixture + def _setup(self, session): + self.create_assets(num=2, session=session) + events = [ + AssetEvent( + asset_id=1, + extra={"region": "us", "env": "prod"}, + source_task_id="t1", + source_dag_id="d1", + source_run_id="r1", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=1, + extra={"region": "eu", "env": "prod"}, + source_task_id="t1", + source_dag_id="d1", + source_run_id="r2", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=2, + extra={"region": "us", "env": "staging"}, + source_task_id="t2", + source_dag_id="d2", + source_run_id="r3", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=1, + extra={}, + source_task_id="t1", + source_dag_id="d1", + source_run_id="r4", + timestamp=DEFAULT_DATE, + ), + ] + session.add_all(events) + session.commit() + + @pytest.mark.usefixtures("_setup") + @pytest.mark.parametrize( + ("params", "expected_count"), + [ + ({"extra": "region=us"}, 2), + ({"extra": "region=eu"}, 1), + ({"extra": "env=prod"}, 2), + ({"extra": "env=staging"}, 1), + ({"extra": "region=ap"}, 0), + ({"extra": "nonexistent=us"}, 0), + ({}, 4), + ], + ) + def test_extra_filter(self, test_client, params, expected_count): + response = test_client.get("/assets/events", params=params) + assert response.status_code == 200 + assert response.json()["total_entries"] == expected_count + + @pytest.mark.usefixtures("_setup") + def test_extra_filter_combined_with_asset_id(self, test_client): + response = test_client.get("/assets/events", params={"extra": "region=us", "asset_id": "1"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + + @pytest.mark.usefixtures("_setup") + @pytest.mark.parametrize( + ("params", "expected_count"), + [ + ([("extra", "region=us"), ("extra", "env=prod")], 1), + ([("extra", "region=eu"), ("extra", "env=prod")], 1), + ([("extra", "region=us"), ("extra", "env=staging")], 1), + ([("extra", "region=eu"), ("extra", "env=staging")], 0), + ], + ) + def test_extra_filter_multiple_keys(self, test_client, params, expected_count): + response = test_client.get("/assets/events", params=params) + assert response.status_code == 200 + assert response.json()["total_entries"] == expected_count + + class TestGetAssetEndpoint(TestAssets): @provide_session def test_should_respond_200(self, test_client, *, session): 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 e3839f19eafe8..6300265472a0a 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 @@ -521,3 +521,210 @@ def test_get_by_asset(self, client): }, ] } + + +class TestGetAssetEventByAssetExtraFilter: + @pytest.fixture + def test_events_with_extra(self, session): + asset = AssetModel( + id=1, + name="test_asset", + uri="s3://bucket/key", + group="asset", + extra={}, + created_at=DEFAULT_DATE, + updated_at=DEFAULT_DATE, + ) + asset_active = AssetActive.for_asset(asset) + session.add_all([asset, asset_active]) + session.flush() + + events = [ + AssetEvent( + id=1, + asset_id=1, + extra={"region": "us", "env": "prod"}, + source_dag_id="d1", + source_task_id="t1", + source_run_id="r1", + source_map_index=-1, + timestamp=DEFAULT_DATE, + ), + AssetEvent( + id=2, + asset_id=1, + extra={"region": "eu", "env": "prod"}, + source_dag_id="d1", + source_task_id="t1", + source_run_id="r2", + source_map_index=-1, + timestamp=DEFAULT_DATE, + ), + AssetEvent( + id=3, + asset_id=1, + extra={"region": "us", "env": "staging"}, + source_dag_id="d1", + source_task_id="t1", + source_run_id="r3", + source_map_index=-1, + timestamp=DEFAULT_DATE, + ), + ] + session.add_all(events) + session.commit() + yield events + for e in events: + session.delete(e) + session.delete(asset_active) + session.delete(asset) + session.commit() + + @pytest.mark.usefixtures("test_events_with_extra") + def test_filter_by_extra_key_value(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params=[("name", "test_asset"), ("uri", "s3://bucket/key"), ("extra", "region=us")], + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 2 + assert all(e["extra"]["region"] == "us" for e in data["asset_events"]) + + @pytest.mark.usefixtures("test_events_with_extra") + def test_filter_by_extra_env(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params=[("name", "test_asset"), ("uri", "s3://bucket/key"), ("extra", "env=prod")], + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 2 + + @pytest.mark.usefixtures("test_events_with_extra") + def test_no_extra_filter_returns_all(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": "test_asset", "uri": "s3://bucket/key"}, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 3 + + @pytest.mark.usefixtures("test_events_with_extra") + def test_filter_nonexistent_key(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params=[("name", "test_asset"), ("uri", "s3://bucket/key"), ("extra", "missing=val")], + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 0 + + @pytest.mark.usefixtures("test_events_with_extra") + @pytest.mark.parametrize( + ("extra_params", "expected_count"), + [ + ([("extra", "region=us"), ("extra", "env=prod")], 1), + ([("extra", "region=eu"), ("extra", "env=prod")], 1), + ([("extra", "region=us"), ("extra", "env=staging")], 1), + ([("extra", "region=eu"), ("extra", "env=staging")], 0), + ], + ) + def test_filter_multiple_extra_keys(self, client, extra_params, expected_count): + response = client.get( + "/execution/asset-events/by-asset", + params=[("name", "test_asset"), ("uri", "s3://bucket/key"), *extra_params], + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == expected_count + + +class TestGetAssetEventByAssetAliasExtraFilter: + @pytest.fixture + def test_alias_events_with_extra(self, session): + asset = AssetModel( + id=1, + name="test_asset", + uri="s3://bucket/key", + group="asset", + extra={}, + created_at=DEFAULT_DATE, + updated_at=DEFAULT_DATE, + ) + asset_active = AssetActive.for_asset(asset) + session.add_all([asset, asset_active]) + session.flush() + + events = [ + AssetEvent( + id=1, + asset_id=1, + extra={"tier": "gold", "region": "us"}, + source_dag_id="d1", + source_task_id="t1", + source_run_id="r1", + source_map_index=-1, + timestamp=DEFAULT_DATE, + ), + AssetEvent( + id=2, + asset_id=1, + extra={"tier": "silver", "region": "eu"}, + source_dag_id="d1", + source_task_id="t1", + source_run_id="r2", + source_map_index=-1, + timestamp=DEFAULT_DATE, + ), + ] + session.add_all(events) + session.flush() + + alias = AssetAliasModel(id=1, name="test_alias") + alias.asset_events = events + alias.assets.append(asset) + session.add(alias) + session.commit() + yield events + session.delete(alias) + for e in events: + session.delete(e) + session.delete(asset_active) + session.delete(asset) + session.commit() + + @pytest.mark.usefixtures("test_alias_events_with_extra") + def test_filter_by_extra_key_value(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params=[("name", "test_alias"), ("extra", "tier=gold")], + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["extra"]["tier"] == "gold" + + @pytest.mark.usefixtures("test_alias_events_with_extra") + def test_no_extra_filter_returns_all(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "test_alias"}, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 2 + + @pytest.mark.usefixtures("test_alias_events_with_extra") + @pytest.mark.parametrize( + ("extra_params", "expected_count"), + [ + ([("extra", "tier=gold"), ("extra", "region=us")], 1), + ([("extra", "tier=gold"), ("extra", "region=eu")], 0), + ([("extra", "tier=silver"), ("extra", "region=eu")], 1), + ], + ) + def test_filter_multiple_extra_keys(self, client, extra_params, expected_count): + response = client.get( + "/execution/asset-events/by-asset-alias", + params=[("name", "test_alias"), *extra_params], + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == expected_count diff --git a/airflow-core/tests/unit/utils/test_db.py b/airflow-core/tests/unit/utils/test_db.py index e54cc6bf4bb15..c81859edbdcfa 100644 --- a/airflow-core/tests/unit/utils/test_db.py +++ b/airflow-core/tests/unit/utils/test_db.py @@ -180,6 +180,8 @@ def test_database_schema_and_sqlalchemy_model_are_in_sync(self, initialized_db): lambda t: t[0] == "add_index" and t[1].name == "idx_ab_user_username", lambda t: t[0] == "remove_index" and t[1].name == "idx_ab_register_user_username", lambda t: t[0] == "remove_index" and t[1].name == "idx_ab_user_username", + # Postgres-only GIN index created by raw SQL in migration, not in SQLAlchemy model + lambda t: t[0] == "remove_index" and t[1].name == "idx_asset_event_extra_gin", ] for ignore in ignores: diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 037ddd7eb5feb..fb6fda21f88fb 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -862,6 +862,7 @@ def get( before: datetime | None = None, ascending: bool = True, limit: int | None = None, + extra: dict[str, str] | None = None, ) -> AssetEventsResponse: """Get Asset event from the API server.""" common_params: dict[str, Any] = {} @@ -870,15 +871,20 @@ def get( if before: common_params["before"] = before.isoformat() common_params["ascending"] = ascending - if limit: + if limit is not None: common_params["limit"] = limit + extra_params: list[tuple[str, str]] = [] + if extra: + extra_params = [("extra", f"{k}={v}") for k, v in extra.items()] if name or uri: resp = self.client.get( - "asset-events/by-asset", params={"name": name, "uri": uri, **common_params} + "asset-events/by-asset", + params=[*{"name": name, "uri": uri, **common_params}.items(), *extra_params], ) elif alias_name: resp = self.client.get( - "asset-events/by-asset-alias", params={"name": alias_name, **common_params} + "asset-events/by-asset-alias", + params=[*{"name": alias_name, **common_params}.items(), *extra_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 c70b3e9b518d8..2141919610f20 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1102,6 +1102,7 @@ class GetAssetEventByAsset(BaseModel): before: AwareDatetime | None = None limit: int | None = None ascending: bool = True + extra: dict[str, str] | None = None type: Literal["GetAssetEventByAsset"] = "GetAssetEventByAsset" @@ -1111,6 +1112,7 @@ class GetAssetEventByAssetAlias(BaseModel): before: AwareDatetime | None = None limit: int | None = None ascending: bool = True + extra: dict[str, str] | None = None 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 d41306009eac7..42daa0cb863d4 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -1081,6 +1081,7 @@ class InletEventsAccessor(Sequence["AssetEventResult"]): _before: str | datetime | None _ascending: bool _limit: int | None + _extra: dict[str, str] _asset_name: str | None _asset_uri: str | None _alias_name: str | None @@ -1095,6 +1096,7 @@ def __init__( self._before = None self._ascending = True self._limit = None + self._extra: dict[str, str] = {} def after(self, after: str) -> Self: self._after = after @@ -1116,6 +1118,11 @@ def limit(self, limit: int) -> Self: self._reset_cache() return self + def extra(self, key: str, value: str) -> Self: + self._extra[key] = value + self._reset_cache() + return self + @functools.cached_property def _asset_events(self) -> list[AssetEventResult]: from airflow.sdk.execution_time.comms import ( @@ -1131,6 +1138,7 @@ def _asset_events(self) -> list[AssetEventResult]: "before": self._before, "ascending": self._ascending, "limit": self._limit, + "extra": self._extra or None, } msg: ToSupervisor diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index e6ce8aa3d066e..643bda0a5c6d1 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1846,6 +1846,21 @@ "title": "Ascending", "type": "boolean" }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, "type": { "const": "GetAssetEventByAsset", "default": "GetAssetEventByAsset", @@ -1909,6 +1924,21 @@ "title": "Ascending", "type": "boolean" }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, "type": { "const": "GetAssetEventByAssetAlias", "default": "GetAssetEventByAssetAlias", diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 76dfd2009ac18..ed269c3ad7adf 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1758,6 +1758,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: before=msg.before, ascending=msg.ascending, limit=msg.limit, + extra=msg.extra, ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) resp = asset_event_result @@ -1769,6 +1770,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: before=msg.before, ascending=msg.ascending, limit=msg.limit, + extra=msg.extra, ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) resp = asset_event_result diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 92b970f7ee51a..205b12bd4f2d8 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1216,6 +1216,68 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert result.asset_events[0].asset.name == "this_asset" assert result.asset_events[0].asset.uri == "s3://bucket/key" + def test_extra_dict_param_passed(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + extra_pairs = [v for k, v in params.multi_items() if k == "extra"] + assert sorted(extra_pairs) == sorted(["region=us", "env=prod"]) + return httpx.Response( + status_code=200, + json={ + "asset_events": [ + { + "id": 1, + "asset": {"name": "a", "uri": "s3://b", "group": "asset"}, + "created_dagruns": [], + "timestamp": "2023-01-01T00:00:00Z", + } + ] + }, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.asset_events.get(name="a", uri="s3://b", extra={"region": "us", "env": "prod"}) + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 1 + + def test_extra_dict_with_alias(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert request.url.path == "/asset-events/by-asset-alias" + extra_pairs = [v for k, v in params.multi_items() if k == "extra"] + assert extra_pairs == ["env=prod"] + return httpx.Response( + status_code=200, + json={ + "asset_events": [ + { + "id": 1, + "asset": {"name": "a", "uri": "s3://b", "group": "asset"}, + "created_dagruns": [], + "timestamp": "2023-01-01T00:00:00Z", + } + ] + }, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.asset_events.get(alias_name="my_alias", extra={"env": "prod"}) + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 1 + + def test_extra_none_not_sent(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert "extra" not in params + return httpx.Response( + status_code=200, + json={"asset_events": []}, + ) + + client = make_client(transport=httpx.MockTransport(handle_request)) + result = client.asset_events.get(name="a", uri="s3://b") + assert isinstance(result, AssetEventsResponse) + class TestAssetOperations: @pytest.mark.parametrize( 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 eb1ffa0013e22..fafa7dacc0674 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -904,6 +904,51 @@ def test__get_item__with_filters(self, sample_inlet_evnets_accessor, mock_superv name="test_uri", uri="test://test/", after=None, before=None, limit=10, ascending=False ) + def test__get_item__with_extra_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] * 3 + + list(sample_inlet_evnets_accessor[TEST_ASSET].extra("region", "us")) + list(sample_inlet_evnets_accessor[TEST_ASSET].extra("region", "us").extra("env", "prod")) + list(sample_inlet_evnets_accessor[TEST_ASSET].extra("region", "us").extra("env", "prod").limit(5)) + + assert mock_supervisor_comms.send.call_count == 3 + + 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, + extra={"region": "us"}, + ) + assert calls[1][0][0] == GetAssetEventByAsset( + name="test_uri", + uri="test://test/", + after=None, + before=None, + limit=None, + ascending=True, + extra={"region": "us", "env": "prod"}, + ) + assert calls[2][0][0] == GetAssetEventByAsset( + name="test_uri", + uri="test://test/", + after=None, + before=None, + limit=5, + ascending=True, + extra={"region": "us", "env": "prod"}, + ) + @pytest.mark.parametrize( ("name", "uri", "expected_key"), ( 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 8dea3d0793f49..2b834d2e6d537 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1987,6 +1987,7 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2030,6 +2031,7 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2066,6 +2068,7 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2109,6 +2112,7 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2145,6 +2149,7 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2188,6 +2193,7 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2223,6 +2229,7 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2264,6 +2271,7 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "extra": None, }, response=AssetEventsResult( asset_events=[ @@ -2278,6 +2286,86 @@ class RequestTestCase: ), test_id="get_asset_events_by_asset_alias_with_filters", ), + RequestTestCase( + message=GetAssetEventByAsset( + uri="s3://bucket/obj", + name="test", + extra={"region": "us"}, + ), + 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": None, + "before": None, + "limit": None, + "ascending": True, + "extra": {"region": "us"}, + }, + 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_with_extra_filter", + ), + RequestTestCase( + message=GetAssetEventByAssetAlias( + alias_name="test_alias", + extra={"env": "prod"}, + ), + 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": None, + "before": None, + "limit": None, + "ascending": True, + "extra": {"env": "prod"}, + }, + 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_alias_with_extra_filter", + ), RequestTestCase( message=ValidateInletsAndOutlets(ti_id=TI_ID), expected_body={