diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 2a7990d233bd1..3bd84a054d06e 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -247,7 +247,19 @@ 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: +The accessor also supports chaining methods to filter events before fetching them. For example, to retrieve only events matching a specific partition key regular expression: + +.. code-block:: python + + @task(inlets=[regional_sales]) + def process_us_sales(*, inlet_events): + us_events = inlet_events[regional_sales].partition_key_regexp_pattern(r"^us\|") + for event in us_events: + print(event.extra, event.partition_key) + +For an exact partition key match, use ``.partition_key(value)`` instead. Regexp filtering is opt-in: it is enabled only by setting :ref:`[api] regexp_query_timeout ` to a positive number of seconds, which also bounds the query runtime; see the config for the security trade-off. + +You can also filter events by their ``extra`` key-value pairs: .. code-block:: python @@ -975,6 +987,45 @@ When a runtime run emits exactly one partition key, the producing these events the same way as timetable-produced partitions, through ``PartitionedAssetTimetable``. +You can also query asset events filtered by partition key using the REST API. +Two parameters are available: + +- ``partition_key`` for **exact match** — uses the B-tree index for fast lookups: + +.. code-block:: bash + + curl -G "http:///api/v2/assets/events" \ + --data-urlencode "partition_key=us|2026-03-10" + +- ``partition_key_regexp_pattern`` for **regular-expression filtering**: + +.. code-block:: bash + + curl -G "http:///api/v2/assets/events" \ + --data-urlencode "partition_key_regexp_pattern=^us" + +Both parameters can be combined; the conditions are applied with AND logic. + +.. note:: + + ``partition_key_regexp_pattern`` is evaluated by the database's own regular-expression engine, + which is a Regular expression Denial of Service (ReDoS) surface. For that reason it is **disabled + by default**: it is enabled only by setting :ref:`[api] regexp_query_timeout ` + to a positive number of seconds (fractional values allowed), which simultaneously bounds the query + runtime (enforced as a ``statement_timeout`` on PostgreSQL and as ``max_execution_time`` on MySQL). + Prefer the exact-match ``partition_key`` (which uses the B-tree index and is always enabled) + whenever a full key is known. + +The same filters are available in the ``InletEventsAccessor``: + +.. code-block:: python + + # Exact match + events = inlet_events[Asset("my_asset")].partition_key("us|2026-03-10").limit(1) + + # Regular-expression pattern + events = inlet_events[Asset("my_asset")].partition_key_regexp_pattern(r"^us\|2026-03-").limit(10) + Fan-out mappers ~~~~~~~~~~~~~~~ diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 3dd6e96a7d00b..de64f8123f059 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 | +=========================+==================+===================+==============================================================+ -| ``c4e7a1f9b2d0`` (head) | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index on asset_event (asset_id, partition_key). | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``c4e7a1f9b2d0`` | ``436dc127462c`` | ``3.4.0`` | Add index on asset.uri. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``436dc127462c`` | ``5a5d3253e946`` | ``3.4.0`` | Drop span_status column. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 060d9478a390e..31471d9e26a96 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -17,8 +17,9 @@ from __future__ import annotations +import re from abc import ABC, abstractmethod -from collections.abc import Callable, Iterable, Sequence +from collections.abc import Callable, Generator, Iterable, Sequence from datetime import datetime from enum import Enum from typing import ( @@ -41,6 +42,7 @@ from sqlalchemy.sql.functions import FunctionElement from airflow._shared.timezones import timezone +from airflow.api_fastapi.common.db.common import SessionDep from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT from airflow.api_fastapi.core_api.base import OrmClause from airflow.api_fastapi.core_api.security import GetUserDep @@ -68,7 +70,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.sqlalchemy import JsonContains, apply_regex_query_timeout from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunType @@ -516,6 +518,75 @@ def depends_search( return depends_search +class _RegexParam(BaseParam[str]): + """ + Filter using database-level regex matching (regexp_match). + + The pattern is handed to the database's own regex engine (via SQLAlchemy's + ``regexp_match``), so this filter is gated behind the ``[api] regexp_query_timeout`` + setting to contain the ReDoS attack surface: it cannot be instantiated with a value + unless a positive timeout is configured (which both enables the feature and bounds it). + + Use :func:`regex_param_factory` to build the FastAPI dependency for this filter. That + dependency also applies :func:`airflow.utils.sqlalchemy.apply_regex_query_timeout` to the + request's session, so the query runtime is bounded automatically and callers never need to + remember to do it in the view. + """ + + def __init__(self, attribute: ColumnElement, value: str | None = None, skip_none: bool = True) -> None: + super().__init__(value=value, skip_none=skip_none) + self.attribute: ColumnElement = attribute + if value is not None and conf.getfloat("api", "regexp_query_timeout") <= 0: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Regexp query filters are disabled. " + "Set [api] regexp_query_timeout to a positive number of seconds to enable them.", + ) + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + return select.where(self.attribute.regexp_match(self.value)) + + @classmethod + def depends(cls, *args: Any, **kwargs: Any) -> Self: + raise NotImplementedError("Use regex_param_factory instead, depends is not implemented.") + + +_DEFAULT_REGEX_DESCRIPTION = "Filter results by matching this regular expression against the field value." + + +def regex_param_factory( + attribute: ColumnElement, + pattern_name: str, + skip_none: bool = True, + description: str = _DEFAULT_REGEX_DESCRIPTION, +) -> Callable[..., Generator[_RegexParam, None, None]]: + def depends_regex( + session: SessionDep, + value: str | None = Query(alias=pattern_name, default=None, description=description), + ) -> Generator[_RegexParam, None, None]: + if value is not None: + try: + re.compile(value) + except re.error as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid regular expression: {e}", + ) + # ``__init__`` rejects the request (400) when a pattern is supplied while the feature is off. + param = _RegexParam(attribute, value, skip_none) + if value is None: + yield param + return + # Bound the query runtime for the whole request so a view can never forget to do it; the + # previous timeout is restored on teardown (after the response) before the session closes. + with apply_regex_query_timeout(session): + yield param + + return depends_regex + + def prefix_search_param_factory( attribute: ColumnElement, prefix_pattern_name: str, @@ -1687,6 +1758,15 @@ def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | N QueryAssetAliasNamePrefixPatternSearch = Annotated[ _PrefixSearchParam, Depends(prefix_search_param_factory(AssetAliasModel.name, "name_prefix_pattern")) ] +QueryAssetEventPartitionKeyFilter = Annotated[ + FilterParam[str | None], + Depends(filter_param_factory(AssetEvent.partition_key, str | None, filter_name="partition_key")), +] +QueryAssetEventPartitionKeyRegex = Annotated[ + _RegexParam, + # ``function`` scope so the dependency can depend on the (function-scoped) session it bounds. + Depends(regex_param_factory(AssetEvent.partition_key, "partition_key_regexp_pattern"), scope="function"), +] QueryAssetDagIdPatternSearch = Annotated[ _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter.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 005da0c6f4b5e..073f8dcfb9f28 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 @@ -447,6 +447,26 @@ paths: - type: integer - type: 'null' title: Source Map Index + - name: partition_key + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Partition Key + - name: partition_key_regexp_pattern + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + description: Filter results by matching this regular expression against + the field value. + title: Partition Key Regexp Pattern + description: Filter results by matching this regular expression against the + field value. - name: name_pattern 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 424938af86d7f..434ab8f1fdd66 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 @@ -38,6 +38,8 @@ QueryAssetAliasNamePrefixPatternSearch, QueryAssetDagIdPatternSearch, QueryAssetEventExtraFilter, + QueryAssetEventPartitionKeyFilter, + QueryAssetEventPartitionKeyRegex, QueryAssetNamePatternSearch, QueryAssetNamePrefixPatternSearch, QueryLimit, @@ -331,6 +333,8 @@ def get_asset_events( source_map_index: Annotated[ FilterParam[int | None], Depends(filter_param_factory(AssetEvent.source_map_index, int | None)) ], + partition_key: QueryAssetEventPartitionKeyFilter, + partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex, name_pattern: QueryAssetNamePatternSearch, name_prefix_pattern: QueryAssetNamePrefixPatternSearch, extra_filter: QueryAssetEventExtraFilter, @@ -338,6 +342,8 @@ def get_asset_events( session: SessionDep, ) -> AssetEventCollectionResponse: """Get asset events.""" + # The regexp partition-key filter bounds the query runtime automatically (its dependency applies + # ``apply_regex_query_timeout`` to this request's session), so no explicit wrapping is needed here. base_statement = select(AssetEvent) if name_pattern.value or name_prefix_pattern.value: base_statement = base_statement.join(AssetModel, AssetEvent.asset_id == AssetModel.id) @@ -350,6 +356,8 @@ def get_asset_events( source_task_id, source_run_id, source_map_index, + partition_key, + partition_key_regexp_pattern, name_pattern, name_prefix_pattern, extra_filter, @@ -364,7 +372,9 @@ def get_asset_events( assets_event_select = assets_event_select.options( subqueryload(AssetEvent.created_dagruns), joinedload(AssetEvent.asset) ) - assets_events = session.scalars(assets_event_select) + # Materialize here (not lazily during response serialization) so the regexp query runs while the + # dependency-applied timeout is still active. + assets_events = session.scalars(assets_event_select).all() return AssetEventCollectionResponse( asset_events=assets_events, 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 b2a98a9de60b7..4fd73b3067208 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,11 @@ from sqlalchemy import and_, select from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.parameters import ( + BaseParam, + QueryAssetEventPartitionKeyFilter, + QueryAssetEventPartitionKeyRegex, +) 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 ( @@ -41,13 +46,23 @@ def _get_asset_events_through_sql_clauses( - *, join_clause, where_clause, session: SessionDep, ascending: bool = True, limit: int | None = None + *, + join_clause, + where_clause, + session: SessionDep, + ascending: bool = True, + limit: int | None = None, + filters: list[BaseParam] | None = None, ) -> 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 is not None: + for filter_ in filters or []: + asset_events_query = filter_.to_orm(asset_events_query) + if limit: asset_events_query = asset_events_query.limit(limit) - asset_events = session.scalars(asset_events_query) + # A regexp partition-key filter bounds the query runtime automatically (its dependency applies + # apply_regex_query_timeout to this request's session), so no explicit wrapping is needed here. + asset_events = session.scalars(asset_events_query).all() return AssetEventsResponse.model_validate( { "asset_events": [ @@ -96,6 +111,8 @@ 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, + partition_key: QueryAssetEventPartitionKeyFilter, + partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex, 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, @@ -136,6 +153,7 @@ def get_asset_event_by_asset_name_uri( session=session, ascending=ascending, limit=limit, + filters=[partition_key, partition_key_regexp_pattern], ) @@ -143,6 +161,8 @@ 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, + partition_key: QueryAssetEventPartitionKeyFilter, + partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex, 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, @@ -169,4 +189,5 @@ def get_asset_event_by_asset_alias( session=session, ascending=ascending, limit=limit, + filters=[partition_key, partition_key_regexp_pattern], ) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 545cded21d71d..b93e7a727112f 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1635,6 +1635,29 @@ api: type: boolean example: ~ default: "True" + regexp_query_timeout: + description: | + Timeout, in seconds, for regular-expression based query filters in the API. Fractional + values are allowed (e.g. ``0.5``). This single value both enables the feature and bounds its + runtime: a value of ``0`` (the default) disables regexp filtering entirely, and any positive + value enables it while capping how long such a query may run. + + When enabled, it allows additional filtering features that require database-side regexp + matching, such as the ``partition_key_regexp_pattern`` filter on ``GET /assets/events`` + and on the Execution API asset-event endpoints. It is disabled by default for security + reasons: a user-supplied regular expression is passed to the database's own regex engine, + which is a Regular expression Denial of Service (ReDoS) vector: a crafted pattern can + force the engine into pathological backtracking and consume database backend CPU. + + The timeout is the primary runtime mitigation for that risk: it is enforced as a + transaction-local ``statement_timeout`` on PostgreSQL and as the session ``max_execution_time`` + on MySQL, so a pathological pattern is aborted instead of pinning a database backend. Keep it + as low as your legitimate queries allow. Coupling the on/off switch with the timeout + intentionally makes it impossible to enable regexp filtering without a runtime bound in place. + version_added: 3.4.0 + type: float + example: ~ + default: "0" secret_key: description: | Secret key used to run your api server. It should be as random as possible. However, when running diff --git a/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py b/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py new file mode 100644 index 0000000000000..d8763debbb9c3 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py @@ -0,0 +1,51 @@ +# +# 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 index on asset_event (asset_id, partition_key). + +Revision ID: 7a98f1b7dbd3 +Revises: c4e7a1f9b2d0 +Create Date: 2026-04-01 00:00:00.000000 + +""" + +from __future__ import annotations + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "7a98f1b7dbd3" +down_revision = "c4e7a1f9b2d0" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Apply Add index on asset_event (asset_id, partition_key).""" + op.create_index( + "idx_asset_event_asset_id_partition_key", + "asset_event", + ["asset_id", "partition_key"], + ) + + +def downgrade(): + """Unapply Add index on asset_event (asset_id, partition_key).""" + op.drop_index("idx_asset_event_asset_id_partition_key", table_name="asset_event") diff --git a/airflow-core/src/airflow/models/asset.py b/airflow-core/src/airflow/models/asset.py index 4354034e649fe..9b1856d268959 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -828,6 +828,7 @@ class AssetEvent(Base): __tablename__ = "asset_event" __table_args__ = ( Index("idx_asset_id_timestamp", asset_id, timestamp), + Index("idx_asset_event_asset_id_partition_key", asset_id, partition_key), {"sqlite_autoincrement": True}, # ensures PK values not reused ) 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 a8f08d3df78c0..709cbfa80124f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -37,7 +37,7 @@ export const UseAssetServiceGetAssetAliasKeyFn = ({ assetAliasId }: { export type AssetServiceGetAssetEventsDefaultResponse = Awaited>; export type AssetServiceGetAssetEventsQueryResult = UseQueryResult; export const useAssetServiceGetAssetEventsKey = "AssetServiceGetAssetEvents"; -export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; extra?: string[]; limit?: number; @@ -45,6 +45,8 @@ export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit, name namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyRegexpPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -53,7 +55,7 @@ export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit, name timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}, queryKey?: Array) => [useAssetServiceGetAssetEventsKey, ...(queryKey ?? [{ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }])]; +} = {}, queryKey?: Array) => [useAssetServiceGetAssetEventsKey, ...(queryKey ?? [{ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, 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 34eff20314070..771c8522f3837 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -81,6 +81,8 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * @param data.sourceTaskId * @param data.sourceRunId * @param data.sourceMapIndex +* @param data.partitionKey +* @param data.partitionKeyRegexpPattern Filter results by matching this regular expression against the field value. * @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. * * **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. @@ -93,7 +95,7 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * @returns AssetEventCollectionResponse Successful Response * @throws ApiError */ -export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient, { assetId, extra, 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, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; extra?: string[]; limit?: number; @@ -101,6 +103,8 @@ export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyRegexpPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -109,7 +113,7 @@ export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}) => 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 }) }); +} = {}) => queryClient.ensureQueryData({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, 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 170a17f9e58ae..46d1ea706dbd4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -81,6 +81,8 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * @param data.sourceTaskId * @param data.sourceRunId * @param data.sourceMapIndex +* @param data.partitionKey +* @param data.partitionKeyRegexpPattern Filter results by matching this regular expression against the field value. * @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. * * **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. @@ -93,7 +95,7 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * @returns AssetEventCollectionResponse Successful Response * @throws ApiError */ -export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, { assetId, extra, 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, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; extra?: string[]; limit?: number; @@ -101,6 +103,8 @@ export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyRegexpPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -109,7 +113,7 @@ export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}) => 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 }) }); +} = {}) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, 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 7efadc2e4f8fa..55776f5fb3311 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -81,6 +81,8 @@ export const useAssetServiceGetAssetAlias = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const useAssetServiceGetAssetEvents = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; extra?: string[]; limit?: number; @@ -101,6 +103,8 @@ export const useAssetServiceGetAssetEvents = , "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 }); +} = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, 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 d4e694554d0ff..8fefbeec56a1a 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -81,6 +81,8 @@ export const useAssetServiceGetAssetAliasSuspense = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const useAssetServiceGetAssetEventsSuspense = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; extra?: string[]; limit?: number; @@ -101,6 +103,8 @@ export const useAssetServiceGetAssetEventsSuspense = , "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 }); +} = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern, 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 a53ae4d51cea8..ff7f3c0b0afdf 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 @@ -122,6 +122,8 @@ export class AssetService { * @param data.sourceTaskId * @param data.sourceRunId * @param data.sourceMapIndex + * @param data.partitionKey + * @param data.partitionKeyRegexpPattern Filter results by matching this regular expression against the field value. * @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. * * **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. @@ -147,6 +149,8 @@ export class AssetService { source_task_id: data.sourceTaskId, source_run_id: data.sourceRunId, source_map_index: data.sourceMapIndex, + partition_key: data.partitionKey, + partition_key_regexp_pattern: data.partitionKeyRegexpPattern, name_pattern: data.namePattern, name_prefix_pattern: data.namePrefixPattern, extra: data.extra, 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 24b1fbe5f2e9b..26be8198d6969 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 @@ -2878,6 +2878,11 @@ export type GetAssetEventsData = { * Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `source_task_id, source_dag_id, source_run_id, source_map_index, timestamp` */ orderBy?: Array<(string)>; + partitionKey?: string | null; + /** + * Filter results by matching this regular expression against the field value. + */ + partitionKeyRegexpPattern?: string | null; sourceDagId?: string | null; sourceMapIndex?: number | null; sourceRunId?: string | null; diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index ce4ef036d20fd..1076715ee7180 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "c4e7a1f9b2d0", + "3.4.0": "7a98f1b7dbd3", } # 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 7ca788ffe1283..0064c0ec3c874 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -62,6 +62,58 @@ def get_dialect_name(session: Session) -> str | None: return getattr(bind.dialect, "name", None) +@contextlib.contextmanager +def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]: + """ + Bound the runtime of a user-supplied regex filter, scoped to the wrapped query. + + Reads the ``[api] regexp_query_timeout`` config (in seconds, fractional values allowed) and sets + a database-side timeout for the duration of the ``with`` block, then restores the previous value + so it does not affect any other statement running later in the same transaction/session: + + * PostgreSQL: transaction-local ``statement_timeout`` (via ``set_config(..., is_local=True)``). + * MySQL: session ``max_execution_time`` (applies to read-only ``SELECT`` statements). + + The previous value is captured and restored (rather than reset to ``0``) so a server- or + role-level global timeout is preserved instead of being cleared. This is a ReDoS safeguard: a + malicious pattern is aborted instead of pinning a database backend. No-op on other backends + (e.g. SQLite) and when the configured timeout is ``0`` (which also means regexp filtering is + disabled). + """ + timeout_seconds = conf.getfloat("api", "regexp_query_timeout") + if timeout_seconds <= 0: + yield + return + # ``timeout_ms`` is derived from the (operator-controlled) config, not from user input. + timeout_ms = int(timeout_seconds * 1000) + dialect_name = get_dialect_name(session) + if dialect_name == "postgresql": + previous = session.execute(text("SELECT current_setting('statement_timeout')")).scalar() + session.execute( + text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) + ) + try: + yield + finally: + # Restore the previous value (e.g. a global statement_timeout) instead of clearing it, so + # the bound only applies to the regexp query and not later statements in the transaction. + session.execute( + text("SELECT set_config('statement_timeout', :timeout, true)").bindparams( + timeout=str(previous) if previous is not None else "0" + ) + ) + elif dialect_name == "mysql": + previous = session.execute(text("SELECT @@SESSION.max_execution_time")).scalar() + session.execute(text(f"SET SESSION max_execution_time = {timeout_ms}")) + try: + yield + finally: + # Restore the previous value so a global max_execution_time is preserved. + session.execute(text(f"SET SESSION max_execution_time = {int(previous or 0)}")) + else: + yield + + def build_upsert_stmt( dialect: str | None, model: Any, diff --git a/airflow-core/tests/unit/api_fastapi/common/test_parameters.py b/airflow-core/tests/unit/api_fastapi/common/test_parameters.py index 4d2f473f4f830..68df51a3f728f 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_parameters.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_parameters.py @@ -21,6 +21,7 @@ from datetime import datetime, timezone from types import SimpleNamespace from typing import Annotated +from unittest import mock import pytest from fastapi import Depends, FastAPI, HTTPException @@ -41,11 +42,15 @@ _TaskDisplayNamePrefixPatternParam, datetime_range_filter_factory, filter_param_factory, + regex_param_factory, ) from airflow.models import DagModel, DagRun, Log +from airflow.models.asset import AssetEvent from airflow.models.errors import ParseImportError from airflow.models.taskinstance import TaskInstance +from tests_common.test_utils.config import conf_vars + class TestFilterParam: def test_filter_param_factory_description(self): @@ -537,3 +542,51 @@ def test_no_coalesce_for_start_date(self): rf = _make_datetime_filter("start_date", upper_bound_lte=bound) sql = _compile(rf.to_orm(select(TaskInstance))) assert "coalesce" not in sql + + +class TestRegexParamFactory: + """The regexp filter dependency must apply the query timeout itself (callers can't forget).""" + + @staticmethod + def _depends(): + return regex_param_factory(AssetEvent.partition_key, "partition_key_regexp_pattern") + + def test_dependency_applies_timeout_when_pattern_provided(self): + session = mock.MagicMock() + with conf_vars({("api", "regexp_query_timeout"): "30"}): + with mock.patch("airflow.api_fastapi.common.parameters.apply_regex_query_timeout") as apply_mock: + gen = self._depends()(session=session, value="^us") + param = next(gen) + apply_mock.assert_called_once_with(session) + assert param.value == "^us" + with pytest.raises(StopIteration): + next(gen) + + def test_dependency_skips_timeout_without_pattern(self): + session = mock.MagicMock() + with conf_vars({("api", "regexp_query_timeout"): "30"}): + with mock.patch("airflow.api_fastapi.common.parameters.apply_regex_query_timeout") as apply_mock: + gen = self._depends()(session=session, value=None) + param = next(gen) + apply_mock.assert_not_called() + assert param.value is None + with pytest.raises(StopIteration): + next(gen) + + def test_dependency_rejects_pattern_when_disabled(self): + session = mock.MagicMock() + with conf_vars({("api", "regexp_query_timeout"): "0"}): + gen = self._depends()(session=session, value="^us") + with pytest.raises(HTTPException) as exc: + next(gen) + assert exc.value.status_code == 400 + assert "disabled" in exc.value.detail + + def test_dependency_rejects_invalid_regex(self): + session = mock.MagicMock() + with conf_vars({("api", "regexp_query_timeout"): "30"}): + gen = self._depends()(session=session, value="[invalid(") + with pytest.raises(HTTPException) as exc: + next(gen) + assert exc.value.status_code == 400 + assert "Invalid regular expression" in exc.value.detail 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 d5c56d133ad3c..835a8ccb6dcba 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 @@ -1087,6 +1087,162 @@ def test_should_mask_sensitive_extra(self, test_client, session): } +class TestGetAssetEventsPartitionKeyRegex(TestAssets): + """Tests for partition_key_regexp_pattern regex filter on GET /assets/events. + + Patterns are written to work consistently across PostgreSQL (~), + MySQL (REGEXP), and SQLite (re.match), including both anchored and + unanchored expressions where appropriate. + """ + + @pytest.fixture(autouse=True) + def _enable_regexp_query_filters(self): + with conf_vars({("api", "regexp_query_timeout"): "30"}): + yield + + @pytest.fixture(autouse=True) + def _create_partition_key_test_data(self, setup, session): + _create_assets(session=session) + events = [ + AssetEvent( + asset_id=1, + extra={}, + source_task_id="t", + source_dag_id="d", + source_run_id="r1", + partition_key="2024-01-01", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=2, + extra={}, + source_task_id="t", + source_dag_id="d", + source_run_id="r2", + partition_key="2024-01-02", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=1, + extra={}, + source_task_id="t", + source_dag_id="d", + source_run_id="r3", + partition_key="us|2024-01-01", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=2, + extra={}, + source_task_id="t", + source_dag_id="d", + source_run_id="r4", + partition_key="eu|2024-01-01", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=1, + extra={}, + source_task_id="t", + source_dag_id="d", + source_run_id="r5", + partition_key="apac|2024-03-20", + timestamp=DEFAULT_DATE, + ), + AssetEvent( + asset_id=1, + extra={}, + source_task_id="t", + source_dag_id="d", + source_run_id="r6", + partition_key=None, + timestamp=DEFAULT_DATE, + ), + ] + session.add_all(events) + session.commit() + + @pytest.mark.parametrize( + ("partition_key_regexp_pattern", "expected_count"), + [ + ("^2024-01-01$", 1), + ("^2024-01-", 2), + ("^us\\|", 1), + (".*\\|2024-01-01$", 2), + ("^(us|eu)\\|", 2), + ("^nonexistent", 0), + ], + ) + def test_partition_key_regexp_pattern_filtering( + self, test_client, partition_key_regexp_pattern, expected_count + ): + response = test_client.get( + "/assets/events", params={"partition_key_regexp_pattern": partition_key_regexp_pattern} + ) + assert response.status_code == 200 + assert response.json()["total_entries"] == expected_count + + @pytest.mark.parametrize( + ("params", "expected_count"), + [ + ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "1"}, 1), + ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "2"}, 0), + ({"partition_key_regexp_pattern": ".*\\|2024-01-01$", "source_dag_id": "d"}, 2), + ({"partition_key_regexp_pattern": ".*\\|2024-01-01$", "source_dag_id": "other"}, 0), + ], + ) + def test_partition_key_regexp_pattern_combined_filters(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 + + def test_partition_key_regexp_pattern_invalid_regex_returns_400(self, test_client): + response = test_client.get( + "/assets/events", params={"partition_key_regexp_pattern": "[invalid(regex"} + ) + assert response.status_code == 400 + assert "Invalid regular expression" in response.json()["detail"] + + def test_partition_key_regexp_pattern_disabled_returns_400(self, test_client): + with conf_vars({("api", "regexp_query_timeout"): "0"}): + response = test_client.get("/assets/events", params={"partition_key_regexp_pattern": "^2024-"}) + assert response.status_code == 400 + assert "disabled" in response.json()["detail"] + + def test_exact_match_works_when_regex_disabled(self, test_client): + with conf_vars({("api", "regexp_query_timeout"): "0"}): + response = test_client.get("/assets/events", params={"partition_key": "2024-01-01"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + + def test_partition_key_exact_match_via_regex(self, test_client): + response = test_client.get("/assets/events", params={"partition_key_regexp_pattern": "^2024-01-01$"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + + @pytest.mark.parametrize( + ("partition_key", "expected_count"), + [ + ("2024-01-01", 1), + ("us|2024-01-01", 1), + ("nonexistent", 0), + ], + ) + def test_partition_key_exact_match(self, test_client, partition_key, expected_count): + response = test_client.get("/assets/events", params={"partition_key": partition_key}) + assert response.status_code == 200 + assert response.json()["total_entries"] == expected_count + + def test_partition_key_and_pattern_combined(self, test_client): + # Both filters are allowed and combine with AND: a disjoint pair yields no results. + response = test_client.get( + "/assets/events", + params={"partition_key": "2024-01-01", "partition_key_regexp_pattern": "^2025-"}, + ) + assert response.status_code == 200 + assert response.json()["total_entries"] == 0 + + class TestGetAssetEventsExtraFilter(TestAssets): @pytest.fixture def _setup(self, 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 6300265472a0a..da1e3705daf7b 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 @@ -24,6 +24,8 @@ from airflow._shared.timezones import timezone from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent, AssetModel +from tests_common.test_utils.config import conf_vars + DEFAULT_DATE = timezone.parse("2021-01-01T00:00:00") pytestmark = pytest.mark.db_test @@ -523,6 +525,355 @@ def test_get_by_asset(self, client): } +class TestGetAssetEventByAssetPartitionKey: + """Tests for partition_key_regexp_pattern regex filter on execution API. + + Patterns are written to work consistently across PostgreSQL (~), + MySQL (REGEXP), and SQLite (re.match). They are not necessarily + ^-anchored; some cases use explicit prefixes such as ``.*`` to + achieve SQLite ``re.match``-compatible behavior. + """ + + @pytest.fixture(autouse=True) + def _enable_regexp_query_filters(self): + with conf_vars({("api", "regexp_query_timeout"): "30"}): + yield + + @pytest.fixture + def test_partitioned_events(self, session, test_asset): + def make_timestamp(day): + return datetime(2021, 1, day, tzinfo=timezone.utc) + + common = { + "asset_id": 1, + "extra": {"foo": "bar"}, + "source_dag_id": "foo", + "source_task_id": "bar", + "source_run_id": "custom", + "source_map_index": -1, + } + + events = [ + AssetEvent(id=10, timestamp=make_timestamp(1), partition_key="2024-01-01", **common), + AssetEvent(id=11, timestamp=make_timestamp(2), partition_key="2024-01-02", **common), + AssetEvent(id=12, timestamp=make_timestamp(3), partition_key="us|2024-01-01", **common), + AssetEvent(id=13, timestamp=make_timestamp(4), partition_key="eu|2024-01-01", **common), + AssetEvent(id=14, timestamp=make_timestamp(5), partition_key="apac|2024-03-20", **common), + AssetEvent(id=15, timestamp=make_timestamp(6), partition_key=None, **common), + ] + session.add_all(events) + session.commit() + yield events + + for event in events: + session.delete(event) + session.commit() + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_exact_partition_key(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": "^2024-01-01$", + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["partition_key"] == "2024-01-01" + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_prefix_partition_key(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": "^2024-01-", + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 2 + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_composite_partition_key_multiple_regions(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": r"^(us|eu)\|2024-01-.*", + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 2 + keys = {e["partition_key"] for e in data["asset_events"]} + assert keys == {"us|2024-01-01", "eu|2024-01-01"} + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_composite_partition_key_date_across_regions(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": r".*\|2024-01-01$", + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 2 + keys = {e["partition_key"] for e in data["asset_events"]} + assert keys == {"us|2024-01-01", "eu|2024-01-01"} + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_partition_key_no_match(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": "^2025-", + }, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 0 + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_without_partition_key_returns_all(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={"name": "test_get_asset_by_name", "uri": None}, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 6 + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_partition_key_and_time_range(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": "^2024-01-", + "after": "2021-01-01T12:00:00Z", + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["partition_key"] == "2024-01-02" + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_with_invalid_regex_returns_400(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": "[invalid(regex", + }, + ) + assert response.status_code == 400 + assert "Invalid regular expression" in response.json()["detail"] + + def test_get_by_asset_with_pattern_disabled_returns_400(self, client): + with conf_vars({("api", "regexp_query_timeout"): "0"}): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_regexp_pattern": "^2021-", + }, + ) + assert response.status_code == 400 + assert "disabled" in response.json()["detail"] + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + @pytest.mark.parametrize( + ("partition_key", "expected_keys"), + [ + ("2024-01-01", ["2024-01-01"]), + ("us|2024-01-01", ["us|2024-01-01"]), + ("nonexistent", []), + ], + ) + def test_get_by_asset_exact_partition_key(self, client, partition_key, expected_keys): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key": partition_key, + }, + ) + assert response.status_code == 200 + keys = [event["partition_key"] for event in response.json()["asset_events"]] + assert keys == expected_keys + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_partition_key_and_pattern_combined(self, client): + # Both filters are allowed and combine with AND: a disjoint pair yields no results. + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key": "2024-01-01", + "partition_key_regexp_pattern": "^2025-", + }, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 0 + + +class TestGetAssetEventByAssetAliasPartitionKey: + """Tests for partition_key_regexp_pattern regex filter on by-asset-alias endpoint. + + All patterns use ^-anchored regex so they work consistently across + PostgreSQL (~), MySQL (REGEXP), and SQLite (re.match). + """ + + @pytest.fixture(autouse=True) + def _enable_regexp_query_filters(self): + with conf_vars({("api", "regexp_query_timeout"): "30"}): + yield + + @pytest.fixture + def test_partitioned_alias_events(self, session, test_asset): + def make_timestamp(day): + return datetime(2021, 1, day, tzinfo=timezone.utc) + + common = { + "asset_id": 1, + "extra": {"foo": "bar"}, + "source_dag_id": "foo", + "source_task_id": "bar", + "source_run_id": "custom", + "source_map_index": -1, + } + + events = [ + AssetEvent(id=20, timestamp=make_timestamp(1), partition_key="us|2024-01-01", **common), + AssetEvent(id=21, timestamp=make_timestamp(2), partition_key="eu|2024-01-01", **common), + AssetEvent(id=22, timestamp=make_timestamp(3), partition_key=None, **common), + ] + session.add_all(events) + session.commit() + + alias = AssetAliasModel(id=10, name="partitioned_alias") + alias.asset_events = events + alias.assets.append(test_asset) + session.add(alias) + session.commit() + + yield events, alias + + session.delete(alias) + for event in events: + session.delete(event) + session.commit() + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_with_partition_key_region(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key_regexp_pattern": r"^us\|.*"}, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["partition_key"] == "us|2024-01-01" + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_with_partition_key_all_regions(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key_regexp_pattern": r".*\|2024-01-01$"}, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 2 + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_without_partition_key_returns_all(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias"}, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 3 + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_with_invalid_regex_returns_400(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key_regexp_pattern": "[invalid(regex"}, + ) + assert response.status_code == 400 + assert "Invalid regular expression" in response.json()["detail"] + + def test_get_by_alias_with_pattern_disabled_returns_400(self, client): + with conf_vars({("api", "regexp_query_timeout"): "0"}): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key_regexp_pattern": "^us"}, + ) + assert response.status_code == 400 + assert "disabled" in response.json()["detail"] + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_with_exact_partition_key_regex(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key_regexp_pattern": r"^us\|2024-01-01$"}, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["partition_key"] == "us|2024-01-01" + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_exact_partition_key(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key": "us|2024-01-01"}, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["partition_key"] == "us|2024-01-01" + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_exact_partition_key_no_match(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key": "nonexistent"}, + ) + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 0 + + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") + def test_get_by_alias_partition_key_and_pattern_combined(self, client): + # Both filters are allowed and combine with AND. + response = client.get( + "/execution/asset-events/by-asset-alias", + params={ + "name": "partitioned_alias", + "partition_key": "us|2024-01-01", + "partition_key_regexp_pattern": r"^us\|", + }, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["asset_events"]) == 1 + assert data["asset_events"][0]["partition_key"] == "us|2024-01-01" + + class TestGetAssetEventByAssetExtraFilter: @pytest.fixture def test_events_with_extra(self, session): diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 1ef7b42c58733..8dbcd41313055 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -35,6 +35,7 @@ from airflow.settings import Session from airflow.utils.sqlalchemy import ( ExecutorConfigType, + apply_regex_query_timeout, ensure_pod_is_valid_after_unpickling, get_dialect_name, prohibit_commit, @@ -43,6 +44,7 @@ from airflow.utils.state import State from airflow.utils.types import DagRunTriggeredByType, DagRunType +from tests_common.test_utils.config import conf_vars from tests_common.test_utils.dag import sync_dag_to_db pytestmark = pytest.mark.db_test @@ -380,3 +382,64 @@ def _refresh_api_key(config): pickle.dumps(fixed_pod) assert fixed_pod.local_vars_configuration.refresh_api_key_hook is None assert fixed_pod.spec.containers[0].local_vars_configuration.refresh_api_key_hook is None + + +class TestApplyRegexQueryTimeout: + @staticmethod + def _mock_session(dialect_name): + session = mock.MagicMock() + session.get_bind.return_value.dialect.name = dialect_name + return session + + def test_sets_and_restores_statement_timeout_on_postgresql(self): + session = self._mock_session("postgresql") + # Pre-existing (e.g. global) statement_timeout captured before we override it. + session.execute.return_value.scalar.return_value = "10s" + with conf_vars({("api", "regexp_query_timeout"): "5"}): + with apply_regex_query_timeout(session): + # Capture-then-set on enter. + assert session.execute.call_count == 2 + set_stmt = session.execute.call_args_list[1].args[0] + # 5 seconds -> 5000 ms, passed as a bound parameter (no SQL injection). + assert set_stmt.compile().params == {"timeout": "5000"} + # Restored on exit to the previous value (not reset to 0), so a global timeout is preserved. + assert session.execute.call_count == 3 + restore_stmt = session.execute.call_args_list[2].args[0] + assert restore_stmt.compile().params == {"timeout": "10s"} + + def test_sets_and_restores_max_execution_time_on_mysql(self): + session = self._mock_session("mysql") + # Pre-existing (e.g. global) max_execution_time captured before we override it. + session.execute.return_value.scalar.return_value = 1000 + with conf_vars({("api", "regexp_query_timeout"): "5"}): + with apply_regex_query_timeout(session): + # Capture-then-set on enter (5 seconds -> 5000 ms). + assert session.execute.call_count == 2 + assert "max_execution_time = 5000" in str(session.execute.call_args.args[0]) + # Restored on exit to the previous value, so a global limit is preserved. + assert session.execute.call_count == 3 + assert "max_execution_time = 1000" in str(session.execute.call_args.args[0]) + + def test_fractional_seconds_are_converted_to_milliseconds(self): + session = self._mock_session("postgresql") + session.execute.return_value.scalar.return_value = "0" + with conf_vars({("api", "regexp_query_timeout"): "0.5"}): + with apply_regex_query_timeout(session): + # Set is the second call (after capturing the previous value). + set_stmt = session.execute.call_args_list[1].args[0] + # 0.5 seconds -> 500 ms. + assert set_stmt.compile().params == {"timeout": "500"} + + def test_noop_on_sqlite(self): + session = self._mock_session("sqlite") + with conf_vars({("api", "regexp_query_timeout"): "5"}): + with apply_regex_query_timeout(session): + pass + session.execute.assert_not_called() + + def test_noop_when_timeout_disabled(self): + session = self._mock_session("postgresql") + with conf_vars({("api", "regexp_query_timeout"): "0"}): + with apply_regex_query_timeout(session): + pass + session.execute.assert_not_called() diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 74a680053711c..7e681c1114956 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -858,6 +858,8 @@ def get( before: datetime | None = None, ascending: bool = True, limit: int | None = None, + partition_key: str | None = None, + partition_key_regexp_pattern: str | None = None, extra: dict[str, str] | None = None, ) -> AssetEventsResponse: """Get Asset event from the API server.""" @@ -869,6 +871,10 @@ def get( common_params["ascending"] = ascending if limit is not None: common_params["limit"] = limit + if partition_key is not None: + common_params["partition_key"] = partition_key + if partition_key_regexp_pattern is not None: + common_params["partition_key_regexp_pattern"] = partition_key_regexp_pattern extra_params: list[tuple[str, str]] = [] if extra: extra_params = [("extra", f"{k}={v}") for k, v in extra.items()] diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 6e9e49b2676c4..5c4ae14dff1cf 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1142,6 +1142,8 @@ class GetAssetEventByAsset(BaseModel): before: AwareDatetime | None = None limit: int | None = None ascending: bool = True + partition_key: str | None = None + partition_key_regexp_pattern: str | None = None extra: dict[str, str] | None = None type: Literal["GetAssetEventByAsset"] = "GetAssetEventByAsset" @@ -1152,6 +1154,8 @@ class GetAssetEventByAssetAlias(BaseModel): before: AwareDatetime | None = None limit: int | None = None ascending: bool = True + partition_key: str | None = None + partition_key_regexp_pattern: str | None = None 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 2828525a7bb28..2c2364a49b2f5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -1114,6 +1114,8 @@ class InletEventsAccessor(Sequence["AssetEventResult"]): _before: str | datetime | None _ascending: bool _limit: int | None + _partition_key: str | None + _partition_key_regexp_pattern: str | None _extra: dict[str, str] _asset_name: str | None _asset_uri: str | None @@ -1129,6 +1131,8 @@ def __init__( self._before = None self._ascending = True self._limit = None + self._partition_key = None + self._partition_key_regexp_pattern = None self._extra: dict[str, str] = {} def after(self, after: str) -> Self: @@ -1151,6 +1155,18 @@ def limit(self, limit: int) -> Self: self._reset_cache() return self + def partition_key(self, key: str) -> Self: + """Filter by exact partition key match.""" + self._partition_key = key + self._reset_cache() + return self + + def partition_key_regexp_pattern(self, pattern: str) -> Self: + """Filter by partition key regexp pattern.""" + self._partition_key_regexp_pattern = pattern + self._reset_cache() + return self + def extra(self, key: str, value: str) -> Self: self._extra[key] = value self._reset_cache() @@ -1165,6 +1181,8 @@ def _asset_events(self) -> list[AssetEventResult]: "before": self._before, "ascending": self._ascending, "limit": self._limit, + "partition_key": self._partition_key, + "partition_key_regexp_pattern": self._partition_key_regexp_pattern, "extra": self._extra or None, } 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 01e3c8970b5cb..0ec8fe4e49a1c 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1872,6 +1872,30 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_regexp_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Regexp Pattern" + }, "extra": { "anyOf": [ { @@ -1950,6 +1974,30 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_regexp_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Regexp Pattern" + }, "extra": { "anyOf": [ { diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 447f7c7594de6..87311f02da7a1 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1774,6 +1774,8 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: before=msg.before, ascending=msg.ascending, limit=msg.limit, + partition_key=msg.partition_key, + partition_key_regexp_pattern=msg.partition_key_regexp_pattern, extra=msg.extra, ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) @@ -1786,6 +1788,8 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: before=msg.before, ascending=msg.ascending, limit=msg.limit, + partition_key=msg.partition_key, + partition_key_regexp_pattern=msg.partition_key_regexp_pattern, extra=msg.extra, ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) diff --git a/task-sdk/tests/task_sdk/api/test_client.py b/task-sdk/tests/task_sdk/api/test_client.py index 205b12bd4f2d8..0f9b8130e2654 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1216,6 +1216,95 @@ 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_partition_key_exact_match_param_passed(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert params.get("partition_key") == "2024-01-01" + assert "partition_key_regexp_pattern" not in params + return httpx.Response( + status_code=200, + json={ + "asset_events": [ + { + "id": 1, + "asset": { + "name": "this_asset", + "uri": "s3://bucket/key", + "group": "asset", + }, + "created_dagruns": [], + "timestamp": "2023-01-01T00:00:00Z", + "partition_key": "2024-01-01", + } + ] + }, + ) + + client = make_client(httpx.MockTransport(handle_request)) + result = client.asset_events.get(name="this_asset", partition_key="2024-01-01") + + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 1 + + def test_partition_key_regexp_pattern_param_passed(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert params.get("partition_key_regexp_pattern") == "^2024-01-" + assert "partition_key" not in params + return httpx.Response( + status_code=200, + json={"asset_events": []}, + ) + + client = make_client(httpx.MockTransport(handle_request)) + result = client.asset_events.get(name="this_asset", partition_key_regexp_pattern="^2024-01-") + + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 0 + + def test_partition_key_regexp_pattern_with_other_params(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert params.get("partition_key_regexp_pattern") == r"^us\|2024-.*" + assert "partition_key" not in params + assert params.get("after") == "2023-06-01T00:00:00+00:00" + assert params.get("limit") == "5" + assert params.get("ascending") == "false" + return httpx.Response( + status_code=200, + json={"asset_events": []}, + ) + + client = make_client(httpx.MockTransport(handle_request)) + result = client.asset_events.get( + name="this_asset", + partition_key_regexp_pattern=r"^us\|2024-.*", + after=datetime(2023, 6, 1, tzinfo=dt_timezone.utc), + limit=5, + ascending=False, + ) + + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 0 + + def test_partition_key_with_alias(self): + def handle_request(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/asset-events/by-asset-alias" + params = request.url.params + assert params.get("name") == "my_alias" + assert params.get("partition_key") == "us-east|2024-01-01" + assert "partition_key_regexp_pattern" not in params + return httpx.Response( + status_code=200, + json={"asset_events": []}, + ) + + client = make_client(httpx.MockTransport(handle_request)) + result = client.asset_events.get(alias_name="my_alias", partition_key="us-east|2024-01-01") + + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 0 + def test_extra_dict_param_passed(self): def handle_request(request: httpx.Request) -> httpx.Response: params = request.url.params 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 870b2e6deedca..c39ceb92b0e3b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1940,6 +1940,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -1984,6 +1986,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2021,6 +2025,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2065,6 +2071,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2102,6 +2110,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2146,6 +2156,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2182,6 +2194,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2224,6 +2238,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2239,6 +2255,49 @@ class RequestTestCase: ), test_id="get_asset_events_by_asset_alias_with_filters", ), + RequestTestCase( + message=GetAssetEventByAsset( + uri="s3://bucket/obj", + name="test", + partition_key="us|2024-01-15", + ), + 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, + "partition_key": "us|2024-01-15", + "partition_key_regexp_pattern": None, + "extra": None, + }, + 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_partition_key", + ), RequestTestCase( message=GetAssetEventByAsset( uri="s3://bucket/obj", @@ -2265,6 +2324,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": {"region": "us"}, }, response=AssetEventsResult( @@ -2280,6 +2341,47 @@ class RequestTestCase: ), test_id="get_asset_events_with_extra_filter", ), + RequestTestCase( + message=GetAssetEventByAssetAlias( + alias_name="test_alias", + partition_key="eu|2024-03", + ), + 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, + "partition_key": "eu|2024-03", + "partition_key_regexp_pattern": None, + "extra": None, + }, + 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_partition_key", + ), RequestTestCase( message=GetAssetEventByAssetAlias( alias_name="test_alias", @@ -2304,6 +2406,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_regexp_pattern": None, "extra": {"env": "prod"}, }, response=AssetEventsResult( diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 4766c0e8a7949..ab2632831ab95 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -372,6 +372,8 @@ export type After = string | null; export type Before = string | null; export type Limit = number | null; export type Ascending = boolean; +export type PartitionKey5 = string | null; +export type PartitionKeyRegexpPattern = string | null; export type Extra7 = { [k: string]: string; } | null; @@ -381,6 +383,8 @@ export type After1 = string | null; export type Before1 = string | null; export type Limit1 = number | null; export type Ascending1 = boolean; +export type PartitionKey6 = string | null; +export type PartitionKeyRegexpPattern1 = string | null; export type Extra8 = { [k: string]: string; } | null; @@ -589,7 +593,7 @@ export type Conf2 = { [k: string]: unknown; } | null; export type ResetDagRun = boolean | null; -export type PartitionKey5 = string | null; +export type PartitionKey7 = string | null; export type Note2 = string | null; export type DagId22 = string; export type DagRunId = string; @@ -1223,6 +1227,8 @@ export interface GetAssetEventByAsset { before?: Before; limit?: Limit; ascending?: Ascending; + partition_key?: PartitionKey5; + partition_key_regexp_pattern?: PartitionKeyRegexpPattern; extra?: Extra7; type?: Type29; } @@ -1236,6 +1242,8 @@ export interface GetAssetEventByAssetAlias { before?: Before1; limit?: Limit1; ascending?: Ascending1; + partition_key?: PartitionKey6; + partition_key_regexp_pattern?: PartitionKeyRegexpPattern1; extra?: Extra8; type?: Type30; } @@ -1803,7 +1811,7 @@ export interface TriggerDagRun { run_after?: RunAfter2; conf?: Conf2; reset_dag_run?: ResetDagRun; - partition_key?: PartitionKey5; + partition_key?: PartitionKey7; note?: Note2; dag_id: DagId22; run_id: DagRunId;