From 02b4b83fc8e824db6d370d453ee81b3639e49152 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Mon, 27 Apr 2026 14:24:14 +0200 Subject: [PATCH 01/13] Add partition_key and partition_key_pattern filters for asset events Add two new query parameters to the asset events API endpoints: - partition_key: exact-match filter leveraging the new composite B-tree index on (asset_id, partition_key) - partition_key_pattern: regex-based filter using database-native regexp_match for flexible pattern matching Includes Core API, Execution API, Task SDK (InletEventsAccessor), client, documentation, migration for the composite index, and tests. Made-with: Cursor --- .../docs/authoring-and-scheduling/assets.rst | 44 +++ airflow-core/docs/migrations-ref.rst | 4 +- .../airflow/api_fastapi/common/parameters.py | 68 ++++ .../openapi/v2-rest-api-generated.yaml | 22 ++ .../core_api/routes/public/assets.py | 12 + .../execution_api/routes/asset_events.py | 64 +++- ...3_3_0_add_idx_asset_event_partition_key.py | 51 +++ airflow-core/src/airflow/models/asset.py | 1 + .../airflow/ui/openapi-gen/queries/common.ts | 6 +- .../ui/openapi-gen/queries/ensureQueryData.ts | 8 +- .../ui/openapi-gen/queries/prefetch.ts | 8 +- .../airflow/ui/openapi-gen/queries/queries.ts | 8 +- .../ui/openapi-gen/queries/suspense.ts | 8 +- .../ui/openapi-gen/requests/services.gen.ts | 4 + .../ui/openapi-gen/requests/types.gen.ts | 5 + airflow-core/src/airflow/utils/db.py | 2 +- .../core_api/routes/public/test_assets.py | 149 ++++++++ .../versions/head/test_asset_events.py | 334 ++++++++++++++++++ task-sdk/src/airflow/sdk/api/client.py | 10 +- .../src/airflow/sdk/execution_time/comms.py | 4 + .../src/airflow/sdk/execution_time/context.py | 20 ++ .../airflow/sdk/execution_time/supervisor.py | 4 + task-sdk/tests/task_sdk/api/test_client.py | 89 +++++ .../execution_time/test_supervisor.py | 98 +++++ 24 files changed, 1009 insertions(+), 14 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0123_3_3_0_add_idx_asset_event_partition_key.py diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 7bf609e9d9834..8547dd0de2b12 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -247,6 +247,20 @@ 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 matching a specific partition key pattern (using database-native regex): + +.. code-block:: python + + @task(inlets=[regional_sales]) + def process_us_sales(*, inlet_events): + us_events = inlet_events[regional_sales].partition_key_pattern(r"^us\|") + for event in us_events: + print(event.extra, event.partition_key) + +For an exact partition key match (no regex), use ``.partition_key(value)`` instead. The two are mutually exclusive; setting one clears the other. + +Other chaining methods include ``.after(timestamp)``, ``.before(timestamp)``, ``.ascending()``, and ``.limit(n)``. + Dependency between ``@asset``, ``@task``, and classic operators --------------------------------------------------------------- @@ -775,5 +789,35 @@ 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_pattern`` for **regex filtering** — uses database-native regex + (PostgreSQL ``~`` operator, MySQL ``REGEXP``, SQLite ``re.match``): + +.. code-block:: bash + + curl -G "http:///api/v2/assets/events" \ + --data-urlencode "partition_key_pattern=^us" + +These parameters are mutually exclusive; providing both returns a 400 error. + +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) + + # Regex pattern + events = inlet_events[Asset("my_asset")].partition_key_pattern(r"^us\|2026-03-").limit(10) + For complete runnable examples, see ``airflow-core/src/airflow/example_dags/example_asset_partition.py``. diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 55c37ecb9a088..a00af025bfc36 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 | +=========================+==================+===================+==============================================================+ -| ``9ff64e1c35d3`` (head) | ``dd5f3a8e2b91`` | ``3.3.0`` | Add indexes on dag_run.created_dag_version_id and | +| ``7a98f1b7dbd3`` (head) | ``9ff64e1c35d3`` | ``3.3.0`` | Add index on asset_event (asset_id, partition_key). | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``9ff64e1c35d3`` | ``dd5f3a8e2b91`` | ``3.3.0`` | Add indexes on dag_run.created_dag_version_id and | | | | | task_instance.dag_version_id. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``dd5f3a8e2b91`` | ``c20871fbf23a`` | ``3.3.0`` | Add rollup_fingerprint to AssetPartitionDagRun and index | diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 685736ed2f43d..f585ad84182e8 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -17,6 +17,7 @@ from __future__ import annotations +import re from abc import ABC, abstractmethod from collections.abc import Callable, Iterable, Sequence from datetime import datetime @@ -505,6 +506,66 @@ def depends_search( return depends_search +class _RegexParam(BaseParam[str]): + """ + Filter using database-level regex matching (regexp_match). + + SQLAlchemy's ``regexp_match`` is supported on PostgreSQL (``~`` operator), + MySQL/MariaDB (``REGEXP``), and SQLite (via Python ``re.match``). + Note: SQLite uses ``re.match`` semantics (anchored at the start of the + string), while PostgreSQL uses ``re.search`` (matches anywhere). + """ + + def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: + super().__init__(skip_none=skip_none) + self.attribute: ColumnElement = attribute + + 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.") + + +def _validate_regex_pattern(value: str | None) -> str | None: + """Validate that the regex pattern is syntactically correct.""" + if value is None: + return value + try: + re.compile(value) + except re.error as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Invalid regular expression: {e}", + ) + return value + + +_DEFAULT_REGEX_DESCRIPTION = ( + "Regex filter. Uses database-native regex " + "(PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). " + "Note: on SQLite, matching is anchored at the start of the string." +) + + +def regex_param_factory( + attribute: ColumnElement, + pattern_name: str, + skip_none: bool = True, + description: str = _DEFAULT_REGEX_DESCRIPTION, +) -> Callable[[str | None], _RegexParam]: + def depends_regex( + value: str | None = Query(alias=pattern_name, default=None, description=description), + ) -> _RegexParam: + _validate_regex_pattern(value) + return _RegexParam(attribute, skip_none).set_value(value) + + return depends_regex + + def prefix_search_param_factory( attribute: ColumnElement, prefix_pattern_name: str, @@ -1560,6 +1621,13 @@ 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, Depends(regex_param_factory(AssetEvent.partition_key, "partition_key_pattern")) +] 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 87fc2a858b9bc..2725f42c7ad6b 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 @@ -433,6 +433,28 @@ 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_pattern + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + description: 'Regex filter. Uses database-native regex (PostgreSQL ~ operator, + MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored + at the start of the string.' + title: Partition Key Pattern + description: 'Regex filter. Uses database-native regex (PostgreSQL ~ operator, + MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at + the start of the string.' - 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 15af36445dd95..1dd0c7eb16c96 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,8 @@ QueryAssetAliasNamePatternSearch, QueryAssetAliasNamePrefixPatternSearch, QueryAssetDagIdPatternSearch, + QueryAssetEventPartitionKeyFilter, + QueryAssetEventPartitionKeyRegex, QueryAssetNamePatternSearch, QueryAssetNamePrefixPatternSearch, QueryLimit, @@ -319,12 +321,20 @@ 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_pattern: QueryAssetEventPartitionKeyRegex, name_pattern: QueryAssetNamePatternSearch, name_prefix_pattern: QueryAssetNamePrefixPatternSearch, timestamp_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("timestamp", AssetEvent))], session: SessionDep, ) -> AssetEventCollectionResponse: """Get asset events.""" + if partition_key.value is not None and partition_key_pattern.value is not None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="partition_key and partition_key_pattern are mutually exclusive.", + ) + base_statement = select(AssetEvent) if name_pattern.value or name_prefix_pattern.value: base_statement = base_statement.join(AssetModel, AssetEvent.asset_id == AssetModel.id) @@ -337,6 +347,8 @@ def get_asset_events( source_task_id, source_run_id, source_map_index, + partition_key, + partition_key_pattern, name_pattern, name_prefix_pattern, timestamp_range, 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..06bb5669ece3e 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 @@ -17,6 +17,7 @@ from __future__ import annotations +import re from typing import Annotated from fastapi import APIRouter, HTTPException, Query, status @@ -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,31 @@ def _get_asset_events_through_sql_clauses( ) +def _validate_partition_key_params(partition_key: str | None, partition_key_pattern: str | None) -> None: + """Validate partition_key and partition_key_pattern are not both set, and that the pattern is valid regex.""" + if partition_key is not None and partition_key_pattern is not None: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "reason": "Mutually exclusive parameters", + "message": "partition_key and partition_key_pattern cannot both be provided. " + "Use partition_key for exact match or partition_key_pattern for regex filtering.", + }, + ) + if partition_key_pattern is None: + return + try: + re.compile(partition_key_pattern) + except re.error as e: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "reason": "Invalid regex", + "message": f"The partition_key_pattern is not a valid regular expression: {e}", + }, + ) + + @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 +108,18 @@ 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, + partition_key: Annotated[ + str | None, + Query(description="Exact partition key filter."), + ] = None, + partition_key_pattern: Annotated[ + str | None, + Query( + description="Partition key regex filter. Uses database-native regex " + "(PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). " + "Note: on SQLite, matching is anchored at the start of the string." + ), + ] = None, ) -> AssetEventsResponse: if name and uri: where_clause = and_(AssetModel.name == name, AssetModel.uri == uri) @@ -103,6 +141,12 @@ def get_asset_event_by_asset_name_uri( if before: where_clause = and_(where_clause, AssetEvent.timestamp <= before) + _validate_partition_key_params(partition_key, partition_key_pattern) + if partition_key is not None: + where_clause = and_(where_clause, AssetEvent.partition_key == partition_key) + elif partition_key_pattern is not None: + where_clause = and_(where_clause, AssetEvent.partition_key.regexp_match(partition_key_pattern)) + return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.asset, where_clause=where_clause, @@ -120,6 +164,18 @@ 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, + partition_key: Annotated[ + str | None, + Query(description="Exact partition key filter."), + ] = None, + partition_key_pattern: Annotated[ + str | None, + Query( + description="Partition key regex filter. Uses database-native regex " + "(PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). " + "Note: on SQLite, matching is anchored at the start of the string." + ), + ] = None, ) -> AssetEventsResponse: where_clause = AssetAliasModel.name == name if after: @@ -127,6 +183,12 @@ def get_asset_event_by_asset_alias( if before: where_clause = and_(where_clause, AssetEvent.timestamp <= before) + _validate_partition_key_params(partition_key, partition_key_pattern) + if partition_key is not None: + where_clause = and_(where_clause, AssetEvent.partition_key == partition_key) + elif partition_key_pattern is not None: + where_clause = and_(where_clause, AssetEvent.partition_key.regexp_match(partition_key_pattern)) + return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.source_aliases, where_clause=where_clause, diff --git a/airflow-core/src/airflow/migrations/versions/0123_3_3_0_add_idx_asset_event_partition_key.py b/airflow-core/src/airflow/migrations/versions/0123_3_3_0_add_idx_asset_event_partition_key.py new file mode 100644 index 0000000000000..2392f5fc8c546 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0123_3_3_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: 9ff64e1c35d3 +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 = "9ff64e1c35d3" +branch_labels = None +depends_on = None +airflow_version = "3.3.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 34e4ffaec3e38..b9eca2b15a0c2 100644 --- a/airflow-core/src/airflow/models/asset.py +++ b/airflow-core/src/airflow/models/asset.py @@ -827,6 +827,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 0933a3fea4d30..0fd3edb0a3af4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -36,13 +36,15 @@ 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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; limit?: number; namePattern?: string; namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -51,7 +53,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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 30096ac546b62..0fe0c0abc913f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -79,6 +79,8 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * @param data.sourceTaskId * @param data.sourceRunId * @param data.sourceMapIndex +* @param data.partitionKey +* @param data.partitionKeyPattern Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. * @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). or 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. @@ -90,13 +92,15 @@ 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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; limit?: number; namePattern?: string; namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -105,7 +109,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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 8f78029a31cf0..e265575b2f2ae 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -79,6 +79,8 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * @param data.sourceTaskId * @param data.sourceRunId * @param data.sourceMapIndex +* @param data.partitionKey +* @param data.partitionKeyPattern Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. * @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). or 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. @@ -90,13 +92,15 @@ 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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; limit?: number; namePattern?: string; namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -105,7 +109,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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 8f5c660029ce2..2151631e57e3f 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -79,6 +79,8 @@ export const useAssetServiceGetAssetAlias = = unknown[]>({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const useAssetServiceGetAssetEvents = = unknown[]>({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; limit?: number; namePattern?: string; namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -105,7 +109,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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 3df940743cc12..6af008e395d3c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -79,6 +79,8 @@ export const useAssetServiceGetAssetAliasSuspense = = unknown[]>({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { +export const useAssetServiceGetAssetEventsSuspense = = unknown[]>({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: { assetId?: number; limit?: number; namePattern?: string; namePrefixPattern?: string; offset?: number; orderBy?: string[]; + partitionKey?: string; + partitionKeyPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; @@ -105,7 +109,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, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 d69a06aa60cde..4b6fdf288af93 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 @@ -120,6 +120,8 @@ export class AssetService { * @param data.sourceTaskId * @param data.sourceRunId * @param data.sourceMapIndex + * @param data.partitionKey + * @param data.partitionKeyPattern Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. * @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). or 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. @@ -144,6 +146,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_pattern: data.partitionKeyPattern, name_pattern: data.namePattern, name_prefix_pattern: data.namePrefixPattern, timestamp_gte: data.timestampGte, 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 62260fa61135f..a701fdbde937e 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 @@ -2724,6 +2724,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; + /** + * Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. + */ + partitionKeyPattern?: 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 2155ca50d33b3..2f8c863ea96e7 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -116,7 +116,7 @@ class MappedClassProtocol(Protocol): "3.1.0": "cc92b33c6709", "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", - "3.3.0": "9ff64e1c35d3", + "3.3.0": "7a98f1b7dbd3", } # Prefix used to identify tables holding data moved during migration. 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 89957e179c2e2..543b83df1f699 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,155 @@ def test_should_mask_sensitive_extra(self, test_client, session): } +class TestGetAssetEventsPartitionKeyRegex(TestAssets): + """Tests for partition_key_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. + """ + + @provide_session + def _create_partition_key_test_data(self, session=None): + _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_pattern", "expected_count"), + [ + ("^2024-01-01$", 1), + ("^2024-01-", 2), + ("^us\\|", 1), + (".*\\|2024-01-01$", 2), + ("^(us|eu)\\|", 2), + ("^nonexistent", 0), + ], + ) + @provide_session + def test_partition_key_pattern_filtering( + self, test_client, partition_key_pattern, expected_count, session + ): + self._create_partition_key_test_data() + response = test_client.get("/assets/events", params={"partition_key_pattern": partition_key_pattern}) + assert response.status_code == 200 + assert response.json()["total_entries"] == expected_count + + @pytest.mark.parametrize( + ("params", "expected_count"), + [ + ({"partition_key_pattern": "^us\\|", "asset_id": "1"}, 1), + ({"partition_key_pattern": "^us\\|", "asset_id": "2"}, 0), + ({"partition_key_pattern": ".*\\|2024-01-01$", "source_dag_id": "d"}, 2), + ({"partition_key_pattern": ".*\\|2024-01-01$", "source_dag_id": "other"}, 0), + ], + ) + @provide_session + def test_partition_key_pattern_combined_filters(self, test_client, params, expected_count, session): + self._create_partition_key_test_data() + response = test_client.get("/assets/events", params=params) + assert response.status_code == 200 + assert response.json()["total_entries"] == expected_count + + @provide_session + def test_partition_key_pattern_invalid_regex_returns_400(self, test_client, session): + response = test_client.get("/assets/events", params={"partition_key_pattern": "[invalid(regex"}) + assert response.status_code == 400 + assert "Invalid regular expression" in response.json()["detail"] + + @provide_session + def test_partition_key_exact_match_via_regex(self, test_client, session): + self._create_partition_key_test_data() + response = test_client.get("/assets/events", params={"partition_key_pattern": "^2024-01-01$"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + + @provide_session + def test_partition_key_exact_match(self, test_client, session): + self._create_partition_key_test_data() + response = test_client.get("/assets/events", params={"partition_key": "2024-01-01"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + + @provide_session + def test_partition_key_exact_match_composite(self, test_client, session): + self._create_partition_key_test_data() + response = test_client.get("/assets/events", params={"partition_key": "us|2024-01-01"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + + @provide_session + def test_partition_key_exact_match_no_match(self, test_client, session): + self._create_partition_key_test_data() + response = test_client.get("/assets/events", params={"partition_key": "nonexistent"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 0 + + @provide_session + def test_partition_key_and_pattern_mutual_exclusion(self, test_client, session): + response = test_client.get( + "/assets/events", + params={"partition_key": "2024-01-01", "partition_key_pattern": "^2024-"}, + ) + assert response.status_code == 400 + + 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..901c0bce15041 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,337 @@ def test_get_by_asset(self, client): }, ] } + + +class TestGetAssetEventByAssetPartitionKey: + """Tests for partition_key_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 + 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_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_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_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_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_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_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_pattern": "[invalid(regex", + }, + ) + assert response.status_code == 400 + assert "Invalid regex" in response.json()["detail"]["reason"] + + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") + def test_get_by_asset_exact_partition_key(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key": "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_exact_partition_key_composite(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "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_events") + def test_get_by_asset_exact_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": "nonexistent", + }, + ) + 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_mutual_exclusion_partition_key_and_pattern(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key": "2024-01-01", + "partition_key_pattern": "^2024-", + }, + ) + assert response.status_code == 400 + assert "Mutually exclusive" in response.json()["detail"]["reason"] + + +class TestGetAssetEventByAssetAliasPartitionKey: + """Tests for partition_key_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 + 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_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_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_pattern": "[invalid(regex"}, + ) + assert response.status_code == 400 + assert "Invalid regex" in response.json()["detail"]["reason"] + + @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_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_mutual_exclusion_partition_key_and_pattern(self, client): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={ + "name": "partitioned_alias", + "partition_key": "us|2024-01-01", + "partition_key_pattern": r"^us\|", + }, + ) + assert response.status_code == 400 + assert "Mutually exclusive" in response.json()["detail"]["reason"] diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 2a4bb63fd5da5..4685885757a4a 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -860,16 +860,24 @@ def get( before: datetime | None = None, ascending: bool = True, limit: int | None = None, + partition_key: str | None = None, + partition_key_pattern: str | None = None, ) -> AssetEventsResponse: """Get Asset event from the API server.""" + if partition_key is not None and partition_key_pattern is not None: + raise ValueError("partition_key and partition_key_pattern are mutually exclusive") common_params: dict[str, Any] = {} if after: common_params["after"] = after.isoformat() if before: common_params["before"] = before.isoformat() common_params["ascending"] = ascending - if limit: + if limit is not None: common_params["limit"] = limit + if partition_key is not None: + common_params["partition_key"] = partition_key + if partition_key_pattern is not None: + common_params["partition_key_pattern"] = partition_key_pattern if name or uri: resp = self.client.get( "asset-events/by-asset", params={"name": name, "uri": uri, **common_params} diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py b/task-sdk/src/airflow/sdk/execution_time/comms.py index 330fe934c7c45..0a2ea89ce580b 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1106,6 +1106,8 @@ class GetAssetEventByAsset(BaseModel): before: AwareDatetime | None = None limit: int | None = None ascending: bool = True + partition_key: str | None = None + partition_key_pattern: str | None = None type: Literal["GetAssetEventByAsset"] = "GetAssetEventByAsset" @@ -1115,6 +1117,8 @@ class GetAssetEventByAssetAlias(BaseModel): before: AwareDatetime | None = None limit: int | None = None ascending: bool = True + partition_key: str | None = None + partition_key_pattern: 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 fa83ce3a5f5f5..2151ab9d3ac5f 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -1030,6 +1030,8 @@ class InletEventsAccessor(Sequence["AssetEventResult"]): _before: str | datetime | None _ascending: bool _limit: int | None + _partition_key: str | None + _partition_key_pattern: str | None _asset_name: str | None _asset_uri: str | None _alias_name: str | None @@ -1044,6 +1046,8 @@ def __init__( self._before = None self._ascending = True self._limit = None + self._partition_key = None + self._partition_key_pattern = None def after(self, after: str) -> Self: self._after = after @@ -1065,6 +1069,20 @@ 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._partition_key_pattern = None + self._reset_cache() + return self + + def partition_key_pattern(self, pattern: str) -> Self: + """Filter by partition key regex pattern.""" + self._partition_key_pattern = pattern + self._partition_key = None + self._reset_cache() + return self + @functools.cached_property def _asset_events(self) -> list[AssetEventResult]: from airflow.sdk.execution_time.comms import ( @@ -1080,6 +1098,8 @@ def _asset_events(self) -> list[AssetEventResult]: "before": self._before, "ascending": self._ascending, "limit": self._limit, + "partition_key": self._partition_key, + "partition_key_pattern": self._partition_key_pattern, } msg: ToSupervisor diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index 6527e651041b8..727d56d5fccc0 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1757,6 +1757,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_pattern=msg.partition_key_pattern, ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) resp = asset_event_result @@ -1768,6 +1770,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_pattern=msg.partition_key_pattern, ) 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 ce4afe51af5fe..46d1aabeed44e 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_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_pattern_param_passed(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert params.get("partition_key_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_pattern="^2024-01-") + + assert isinstance(result, AssetEventsResponse) + assert len(result.asset_events) == 0 + + def test_partition_key_pattern_with_other_params(self): + def handle_request(request: httpx.Request) -> httpx.Response: + params = request.url.params + assert params.get("partition_key_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_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_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 + class TestAssetOperations: @pytest.mark.parametrize( 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 62a9d00fdf286..6ebd5ba995ec3 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -2018,6 +2018,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2061,6 +2063,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2097,6 +2101,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2140,6 +2146,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2176,6 +2184,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2219,6 +2229,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2254,6 +2266,8 @@ class RequestTestCase: "before": None, "limit": None, "ascending": True, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2295,6 +2309,8 @@ class RequestTestCase: "before": timezone.parse("2024-10-15T12:00:00Z"), "limit": 5, "ascending": False, + "partition_key": None, + "partition_key_pattern": None, }, response=AssetEventsResult( asset_events=[ @@ -2309,6 +2325,88 @@ 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_pattern": 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=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_pattern": 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=ValidateInletsAndOutlets(ti_id=TI_ID), expected_body={ From 36772095c795a84561ccf6ea319855fc7f890cc4 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Mon, 6 Jul 2026 21:06:27 +0200 Subject: [PATCH 02/13] Regenerate supervisor schema snapshot for partition_key filters --- .../sdk/execution_time/schema/schema.json | 48 +++++++++++++++++++ 1 file changed, 48 insertions(+) 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..aba757c437372 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,30 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Pattern" + }, "type": { "const": "GetAssetEventByAsset", "default": "GetAssetEventByAsset", @@ -1909,6 +1933,30 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Pattern" + }, "type": { "const": "GetAssetEventByAssetAlias", "default": "GetAssetEventByAssetAlias", From 54a42fa093831a59f7b2d223340194c349e3e6e9 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Mon, 6 Jul 2026 21:22:02 +0200 Subject: [PATCH 03/13] Add ReDoS safeguards for partition_key_pattern regex filter The pattern is evaluated by the database's regex engine, so a pathological pattern could consume DB CPU (ReDoS). Mitigations: - Cap the pattern length (MAX_REGEX_PATTERN_LENGTH) in both the Core and Execution API validators, returning 400 for over-long patterns. - Bound the regex query with a transaction-local statement_timeout on PostgreSQL (apply_regex_query_timeout); no-op elsewhere. MySQL bounds regex evaluation via its built-in regexp_time_limit; SQLite has no server-side regex and MariaDB is unsupported. - Document the safeguards in assets.rst. --- .../docs/authoring-and-scheduling/assets.rst | 8 +++++++ .../airflow/api_fastapi/common/parameters.py | 14 ++++++++++- .../core_api/routes/public/assets.py | 5 ++++ .../execution_api/routes/asset_events.py | 13 +++++++++++ airflow-core/src/airflow/utils/sqlalchemy.py | 23 +++++++++++++++++++ .../core_api/routes/public/test_assets.py | 5 ++++ .../versions/head/test_asset_events.py | 12 ++++++++++ 7 files changed, 79 insertions(+), 1 deletion(-) diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index ce1a2f4b9e4a5..08a05442501b3 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -989,6 +989,14 @@ Two parameters are available: These parameters are mutually exclusive; providing both returns a 400 error. +.. note:: + + ``partition_key_pattern`` is evaluated by the database's own regular-expression engine. + To guard against denial-of-service from pathological patterns (ReDoS), the pattern length + is capped and, on PostgreSQL, the query runs under a bounded ``statement_timeout``. MySQL + bounds regex evaluation with its built-in ``regexp_time_limit``. Prefer the exact-match + ``partition_key`` (which uses the B-tree index) whenever a full key is known. + The same filters are available in the ``InletEventsAccessor``: .. code-block:: python diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 7a5e5e0bb0426..8d6fee6e92637 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -540,10 +540,22 @@ def depends(cls, *args: Any, **kwargs: Any) -> Self: raise NotImplementedError("Use regex_param_factory instead, depends is not implemented.") +# Maximum length allowed for a user-supplied regex filter. The pattern is evaluated by the +# database's own regex engine (see ``_RegexParam``), so an overly long/complex pattern could be +# used to consume excessive DB CPU (ReDoS). Partition-key patterns are short in practice, so we +# cap the length as a cheap first line of defense against pathological input. +MAX_REGEX_PATTERN_LENGTH = 200 + + def _validate_regex_pattern(value: str | None) -> str | None: - """Validate that the regex pattern is syntactically correct.""" + """Validate that the regex pattern is syntactically correct and not excessively long.""" if value is None: return value + if len(value) > MAX_REGEX_PATTERN_LENGTH: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Regular expression is too long (max {MAX_REGEX_PATTERN_LENGTH} characters).", + ) try: re.compile(value) except re.error as e: 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 55dfb320527ad..66f14371a1901 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 @@ -85,6 +85,7 @@ TaskOutletAssetReference, ) from airflow.typing_compat import Unpack +from airflow.utils.sqlalchemy import apply_regex_query_timeout from airflow.utils.state import DagRunState from airflow.utils.types import DagRunTriggeredByType, DagRunType @@ -335,6 +336,10 @@ def get_asset_events( detail="partition_key and partition_key_pattern are mutually exclusive.", ) + if partition_key_pattern.value is not None: + # Bound the runtime of the user-supplied regex to guard against ReDoS (PostgreSQL only). + apply_regex_query_timeout(session) + base_statement = select(AssetEvent) if name_pattern.value or name_prefix_pattern.value: base_statement = base_statement.join(AssetModel, AssetEvent.asset_id == AssetModel.id) 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 06bb5669ece3e..bc83c9cd8bbb7 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 @@ -24,6 +24,7 @@ from sqlalchemy import and_, select from airflow.api_fastapi.common.db.common import SessionDep +from airflow.api_fastapi.common.parameters import MAX_REGEX_PATTERN_LENGTH 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 ( @@ -31,6 +32,7 @@ AssetEventsResponse, ) from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel +from airflow.utils.sqlalchemy import apply_regex_query_timeout router = APIRouter( responses={ @@ -87,6 +89,15 @@ def _validate_partition_key_params(partition_key: str | None, partition_key_patt ) if partition_key_pattern is None: return + if len(partition_key_pattern) > MAX_REGEX_PATTERN_LENGTH: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "reason": "Invalid regex", + "message": f"The partition_key_pattern is too long " + f"(max {MAX_REGEX_PATTERN_LENGTH} characters).", + }, + ) try: re.compile(partition_key_pattern) except re.error as e: @@ -145,6 +156,7 @@ def get_asset_event_by_asset_name_uri( if partition_key is not None: where_clause = and_(where_clause, AssetEvent.partition_key == partition_key) elif partition_key_pattern is not None: + apply_regex_query_timeout(session) where_clause = and_(where_clause, AssetEvent.partition_key.regexp_match(partition_key_pattern)) return _get_asset_events_through_sql_clauses( @@ -187,6 +199,7 @@ def get_asset_event_by_asset_alias( if partition_key is not None: where_clause = and_(where_clause, AssetEvent.partition_key == partition_key) elif partition_key_pattern is not None: + apply_regex_query_timeout(session) where_clause = and_(where_clause, AssetEvent.partition_key.regexp_match(partition_key_pattern)) return _get_asset_events_through_sql_clauses( diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index c47d8fd4796b0..e494da38f6809 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -61,6 +61,29 @@ def get_dialect_name(session: Session) -> str | None: return getattr(bind.dialect, "name", None) +# Upper bound (milliseconds) for a query that evaluates a user-supplied regular expression. +# The regex is matched by the database engine, so a pathological pattern could otherwise consume +# DB CPU indefinitely (ReDoS). PostgreSQL has no built-in bound, so we cap it with a +# transaction-local ``statement_timeout``. MySQL bounds regex evaluation with its own engine +# limits (``regexp_time_limit``), and SQLite has no server-side regex in Airflow. +REGEX_QUERY_STATEMENT_TIMEOUT_MS = 30_000 + + +def apply_regex_query_timeout(session: Session, timeout_ms: int = REGEX_QUERY_STATEMENT_TIMEOUT_MS) -> None: + """ + Bound the runtime of a user-supplied regex filter on PostgreSQL. + + Sets a transaction-local ``statement_timeout`` (via ``set_config(..., is_local=True)``) so it + applies only to the current transaction and resets automatically at commit/rollback. This is a + ReDoS safeguard: a malicious pattern is aborted instead of pinning a database backend. No-op on + non-PostgreSQL backends. + """ + if get_dialect_name(session) == "postgresql": + session.execute( + text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) + ) + + def build_upsert_stmt( dialect: str | None, model: Any, 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 b0214e81a0b3d..e80f76da51899 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 @@ -1190,6 +1190,11 @@ def test_partition_key_pattern_invalid_regex_returns_400(self, test_client, sess assert response.status_code == 400 assert "Invalid regular expression" in response.json()["detail"] + def test_partition_key_pattern_too_long_returns_400(self, test_client, session): + response = test_client.get("/assets/events", params={"partition_key_pattern": "a" * 201}) + assert response.status_code == 400 + assert "too long" in response.json()["detail"] + @provide_session def test_partition_key_exact_match_via_regex(self, test_client, session): self._create_partition_key_test_data() 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 901c0bce15041..5bf26f9447bdd 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 @@ -674,6 +674,18 @@ def test_get_by_asset_with_invalid_regex_returns_400(self, client): assert response.status_code == 400 assert "Invalid regex" in response.json()["detail"]["reason"] + def test_get_by_asset_with_too_long_pattern_returns_400(self, client): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_pattern": "a" * 201, + }, + ) + assert response.status_code == 400 + assert "too long" in response.json()["detail"]["message"] + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") def test_get_by_asset_exact_partition_key(self, client): response = client.get( From 0ee34eb48ae7830f2eb3e96feb2490620697ebeb Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Mon, 6 Jul 2026 23:30:37 +0200 Subject: [PATCH 04/13] Gate regex query filters behind a config flag with configurable timeout Add two [api] configs (Airflow 3.4.0): - enable_regexp_query_filters (bool, default False): the partition_key_pattern filter passes a user-supplied regex to the database engine, which is a ReDoS vector, so it is disabled by default and must be explicitly enabled. Exact-match partition_key filtering is unaffected. - regexp_query_timeout (int seconds, default 30): the primary runtime mitigation, enforced as a PostgreSQL transaction-local statement_timeout; 0 disables it. When the flag is off, partition_key_pattern requests return HTTP 400 on both the Core and Execution APIs. Config descriptions document the security rationale. Adds tests for the disabled path and the timeout helper. --- .../docs/authoring-and-scheduling/assets.rst | 12 ++++--- .../airflow/api_fastapi/common/parameters.py | 7 +++- .../execution_api/routes/asset_events.py | 10 ++++++ .../src/airflow/config_templates/config.yml | 33 +++++++++++++++++++ airflow-core/src/airflow/utils/sqlalchemy.py | 23 ++++++------- .../core_api/routes/public/test_assets.py | 19 +++++++++++ .../versions/head/test_asset_events.py | 20 +++++++++++ .../tests/unit/utils/test_sqlalchemy.py | 31 +++++++++++++++++ 8 files changed, 136 insertions(+), 19 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 08a05442501b3..2f4a3c5c3e9fe 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -991,11 +991,13 @@ These parameters are mutually exclusive; providing both returns a 400 error. .. note:: - ``partition_key_pattern`` is evaluated by the database's own regular-expression engine. - To guard against denial-of-service from pathological patterns (ReDoS), the pattern length - is capped and, on PostgreSQL, the query runs under a bounded ``statement_timeout``. MySQL - bounds regex evaluation with its built-in ``regexp_time_limit``. Prefer the exact-match - ``partition_key`` (which uses the B-tree index) whenever a full key is known. + ``partition_key_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** and must be enabled with ``[api] enable_regexp_query_filters = True``. When enabled, + the pattern length is capped and, on PostgreSQL, the query runs under a bounded + ``statement_timeout`` controlled by ``[api] regexp_query_timeout`` (seconds); MySQL bounds regex + evaluation with its built-in ``regexp_time_limit``. 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``: diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 8d6fee6e92637..2cca657daa6e9 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -548,9 +548,14 @@ def depends(cls, *args: Any, **kwargs: Any) -> Self: def _validate_regex_pattern(value: str | None) -> str | None: - """Validate that the regex pattern is syntactically correct and not excessively long.""" + """Validate that the regex pattern is enabled, syntactically correct and not excessively long.""" if value is None: return value + if not conf.getboolean("api", "enable_regexp_query_filters"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Regex query filters are disabled. Set [api] enable_regexp_query_filters = True to use them.", + ) if len(value) > MAX_REGEX_PATTERN_LENGTH: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, 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 bc83c9cd8bbb7..53a9a071bd41c 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 @@ -31,6 +31,7 @@ AssetEventResponse, AssetEventsResponse, ) +from airflow.configuration import conf from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel from airflow.utils.sqlalchemy import apply_regex_query_timeout @@ -89,6 +90,15 @@ def _validate_partition_key_params(partition_key: str | None, partition_key_patt ) if partition_key_pattern is None: return + if not conf.getboolean("api", "enable_regexp_query_filters"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ + "reason": "Regex query filters disabled", + "message": "Regex query filters are disabled. " + "Set [api] enable_regexp_query_filters = True to use partition_key_pattern.", + }, + ) if len(partition_key_pattern) > MAX_REGEX_PATTERN_LENGTH: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 545cded21d71d..9ba30377f425a 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1635,6 +1635,39 @@ api: type: boolean example: ~ default: "True" + enable_regexp_query_filters: + description: | + Whether to enable regular-expression based query filters in the API. Currently this + controls the ``partition_key_pattern`` filter on ``GET /assets/events`` and on the + Execution API asset-event endpoints. + + Disabled by default for security reasons. When enabled, 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. Only enable this if you + accept that trade-off, and keep ``regexp_query_timeout`` set to a sensible value. + + Exact-match ``partition_key`` filtering is always available and is not affected by this + setting. + version_added: 3.4.0 + type: boolean + example: ~ + default: "False" + regexp_query_timeout: + description: | + Maximum time, in seconds, that a query using a regular-expression filter (see + ``enable_regexp_query_filters``) is allowed to run. + + This is the primary runtime mitigation for the ReDoS risk described above: on PostgreSQL + it is enforced as a transaction-local ``statement_timeout`` so that a pathological + pattern is aborted instead of pinning a database backend (PostgreSQL has no built-in + regex step limit, unlike MySQL's ``regexp_time_limit``). Keep it as low as your + legitimate queries allow. A value of ``0`` disables the timeout, which is not + recommended while the feature is enabled. + version_added: 3.4.0 + type: integer + example: ~ + default: "30" 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/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index e494da38f6809..ff0c7a9f8f189 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -61,24 +61,21 @@ def get_dialect_name(session: Session) -> str | None: return getattr(bind.dialect, "name", None) -# Upper bound (milliseconds) for a query that evaluates a user-supplied regular expression. -# The regex is matched by the database engine, so a pathological pattern could otherwise consume -# DB CPU indefinitely (ReDoS). PostgreSQL has no built-in bound, so we cap it with a -# transaction-local ``statement_timeout``. MySQL bounds regex evaluation with its own engine -# limits (``regexp_time_limit``), and SQLite has no server-side regex in Airflow. -REGEX_QUERY_STATEMENT_TIMEOUT_MS = 30_000 - - -def apply_regex_query_timeout(session: Session, timeout_ms: int = REGEX_QUERY_STATEMENT_TIMEOUT_MS) -> None: +def apply_regex_query_timeout(session: Session) -> None: """ Bound the runtime of a user-supplied regex filter on PostgreSQL. - Sets a transaction-local ``statement_timeout`` (via ``set_config(..., is_local=True)``) so it - applies only to the current transaction and resets automatically at commit/rollback. This is a - ReDoS safeguard: a malicious pattern is aborted instead of pinning a database backend. No-op on - non-PostgreSQL backends. + Reads the ``[api] regexp_query_timeout`` config (in seconds) and sets a transaction-local + ``statement_timeout`` (via ``set_config(..., is_local=True)``) so it applies only to the current + transaction and resets automatically at commit/rollback. This is a ReDoS safeguard: a malicious + pattern is aborted instead of pinning a database backend. No-op on non-PostgreSQL backends and + when the configured timeout is ``0`` (disabled). """ + timeout_seconds = conf.getint("api", "regexp_query_timeout") + if timeout_seconds <= 0: + return if get_dialect_name(session) == "postgresql": + timeout_ms = timeout_seconds * 1000 session.execute( text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) ) 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 e80f76da51899..5b50da181a4f4 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 @@ -1086,6 +1086,11 @@ class TestGetAssetEventsPartitionKeyRegex(TestAssets): unanchored expressions where appropriate. """ + @pytest.fixture(autouse=True) + def _enable_regexp_query_filters(self): + with conf_vars({("api", "enable_regexp_query_filters"): "True"}): + yield + @provide_session def _create_partition_key_test_data(self, session=None): _create_assets(session=session) @@ -1195,6 +1200,20 @@ def test_partition_key_pattern_too_long_returns_400(self, test_client, session): assert response.status_code == 400 assert "too long" in response.json()["detail"] + def test_partition_key_pattern_disabled_returns_400(self, test_client, session): + with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + response = test_client.get("/assets/events", params={"partition_key_pattern": "^2024-"}) + assert response.status_code == 400 + assert "disabled" in response.json()["detail"] + + @provide_session + def test_exact_match_works_when_regex_disabled(self, test_client, session): + self._create_partition_key_test_data() + with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + response = test_client.get("/assets/events", params={"partition_key": "2024-01-01"}) + assert response.status_code == 200 + assert response.json()["total_entries"] == 1 + @provide_session def test_partition_key_exact_match_via_regex(self, test_client, session): self._create_partition_key_test_data() 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 5bf26f9447bdd..8a60e757046df 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 @@ -532,6 +534,11 @@ class TestGetAssetEventByAssetPartitionKey: achieve SQLite ``re.match``-compatible behavior. """ + @pytest.fixture(autouse=True) + def _enable_regexp_query_filters(self): + with conf_vars({("api", "enable_regexp_query_filters"): "True"}): + yield + @pytest.fixture def test_partitioned_events(self, session, test_asset): def make_timestamp(day): @@ -686,6 +693,19 @@ def test_get_by_asset_with_too_long_pattern_returns_400(self, client): assert response.status_code == 400 assert "too long" in response.json()["detail"]["message"] + def test_get_by_asset_with_pattern_disabled_returns_400(self, client): + with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + response = client.get( + "/execution/asset-events/by-asset", + params={ + "name": "test_get_asset_by_name", + "uri": None, + "partition_key_pattern": "^2021-", + }, + ) + assert response.status_code == 400 + assert "disabled" in response.json()["detail"]["message"] + @pytest.mark.usefixtures("test_asset", "test_partitioned_events") def test_get_by_asset_exact_partition_key(self, client): response = client.get( diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 1ef7b42c58733..fd7f74ff13b7e 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,32 @@ 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_statement_timeout_on_postgresql(self): + session = self._mock_session("postgresql") + with conf_vars({("api", "regexp_query_timeout"): "5"}): + apply_regex_query_timeout(session) + session.execute.assert_called_once() + stmt = session.execute.call_args.args[0] + # 5 seconds -> 5000 ms, passed as a bound parameter (no SQL injection). + assert stmt.compile().params == {"timeout": "5000"} + + def test_noop_on_non_postgresql(self): + session = self._mock_session("mysql") + with conf_vars({("api", "regexp_query_timeout"): "5"}): + apply_regex_query_timeout(session) + 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"}): + apply_regex_query_timeout(session) + session.execute.assert_not_called() From 8ff1f1e02064a65b98d65e1de57d0b6a6d2feec2 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Mon, 6 Jul 2026 23:33:39 +0200 Subject: [PATCH 05/13] Revert pattern length cap for regex filters Drop MAX_REGEX_PATTERN_LENGTH; the opt-in flag and the PostgreSQL statement_timeout are the meaningful ReDoS safeguards, and a length cap does not stop short catastrophic patterns while risking false rejections of legitimate ones. --- .../docs/authoring-and-scheduling/assets.rst | 8 ++++---- .../src/airflow/api_fastapi/common/parameters.py | 14 +------------- .../execution_api/routes/asset_events.py | 10 ---------- .../core_api/routes/public/test_assets.py | 5 ----- .../versions/head/test_asset_events.py | 12 ------------ 5 files changed, 5 insertions(+), 44 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 2f4a3c5c3e9fe..d98dfbd69abce 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -994,10 +994,10 @@ These parameters are mutually exclusive; providing both returns a 400 error. ``partition_key_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** and must be enabled with ``[api] enable_regexp_query_filters = True``. When enabled, - the pattern length is capped and, on PostgreSQL, the query runs under a bounded - ``statement_timeout`` controlled by ``[api] regexp_query_timeout`` (seconds); MySQL bounds regex - evaluation with its built-in ``regexp_time_limit``. Prefer the exact-match ``partition_key`` - (which uses the B-tree index and is always enabled) whenever a full key is known. + the query runs on PostgreSQL under a bounded ``statement_timeout`` controlled by + ``[api] regexp_query_timeout`` (seconds); MySQL bounds regex evaluation with its built-in + ``regexp_time_limit``. 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``: diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 2cca657daa6e9..00513074f6821 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -540,15 +540,8 @@ def depends(cls, *args: Any, **kwargs: Any) -> Self: raise NotImplementedError("Use regex_param_factory instead, depends is not implemented.") -# Maximum length allowed for a user-supplied regex filter. The pattern is evaluated by the -# database's own regex engine (see ``_RegexParam``), so an overly long/complex pattern could be -# used to consume excessive DB CPU (ReDoS). Partition-key patterns are short in practice, so we -# cap the length as a cheap first line of defense against pathological input. -MAX_REGEX_PATTERN_LENGTH = 200 - - def _validate_regex_pattern(value: str | None) -> str | None: - """Validate that the regex pattern is enabled, syntactically correct and not excessively long.""" + """Validate that the regex pattern is enabled and syntactically correct.""" if value is None: return value if not conf.getboolean("api", "enable_regexp_query_filters"): @@ -556,11 +549,6 @@ def _validate_regex_pattern(value: str | None) -> str | None: status_code=status.HTTP_400_BAD_REQUEST, detail="Regex query filters are disabled. Set [api] enable_regexp_query_filters = True to use them.", ) - if len(value) > MAX_REGEX_PATTERN_LENGTH: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Regular expression is too long (max {MAX_REGEX_PATTERN_LENGTH} characters).", - ) try: re.compile(value) except re.error as e: 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 53a9a071bd41c..07922b64f312a 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 @@ -24,7 +24,6 @@ from sqlalchemy import and_, select from airflow.api_fastapi.common.db.common import SessionDep -from airflow.api_fastapi.common.parameters import MAX_REGEX_PATTERN_LENGTH 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 ( @@ -99,15 +98,6 @@ def _validate_partition_key_params(partition_key: str | None, partition_key_patt "Set [api] enable_regexp_query_filters = True to use partition_key_pattern.", }, ) - if len(partition_key_pattern) > MAX_REGEX_PATTERN_LENGTH: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "reason": "Invalid regex", - "message": f"The partition_key_pattern is too long " - f"(max {MAX_REGEX_PATTERN_LENGTH} characters).", - }, - ) try: re.compile(partition_key_pattern) except re.error as e: 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 5b50da181a4f4..da3f33bd4239d 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 @@ -1195,11 +1195,6 @@ def test_partition_key_pattern_invalid_regex_returns_400(self, test_client, sess assert response.status_code == 400 assert "Invalid regular expression" in response.json()["detail"] - def test_partition_key_pattern_too_long_returns_400(self, test_client, session): - response = test_client.get("/assets/events", params={"partition_key_pattern": "a" * 201}) - assert response.status_code == 400 - assert "too long" in response.json()["detail"] - def test_partition_key_pattern_disabled_returns_400(self, test_client, session): with conf_vars({("api", "enable_regexp_query_filters"): "False"}): response = test_client.get("/assets/events", params={"partition_key_pattern": "^2024-"}) 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 8a60e757046df..abb85b105266c 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 @@ -681,18 +681,6 @@ def test_get_by_asset_with_invalid_regex_returns_400(self, client): assert response.status_code == 400 assert "Invalid regex" in response.json()["detail"]["reason"] - def test_get_by_asset_with_too_long_pattern_returns_400(self, client): - response = client.get( - "/execution/asset-events/by-asset", - params={ - "name": "test_get_asset_by_name", - "uri": None, - "partition_key_pattern": "a" * 201, - }, - ) - assert response.status_code == 400 - assert "too long" in response.json()["detail"]["message"] - def test_get_by_asset_with_pattern_disabled_returns_400(self, client): with conf_vars({("api", "enable_regexp_query_filters"): "False"}): response = client.get( From 81464ecd86b15a9af7705d716012b1a35479eb3e Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Tue, 7 Jul 2026 10:16:45 +0200 Subject: [PATCH 06/13] Fix CI: enable regex flag for by-alias tests, regenerate TS SDK schema - Add the enable_regexp_query_filters conf_vars fixture to TestGetAssetEventByAssetAliasPartitionKey (was missed), fixing the 4 by-alias regex tests that now hit the default-off flag; add a disabled->400 test for the alias endpoint too. - Regenerate ts-sdk/src/generated/supervisor.ts for the partition_key / partition_key_pattern comms fields (check-ts-sdk-supervisor-schema). --- .../versions/head/test_asset_events.py | 14 ++++++++++++++ ts-sdk/src/generated/supervisor.ts | 12 ++++++++++-- 2 files changed, 24 insertions(+), 2 deletions(-) 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 abb85b105266c..9fab45d492943 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 @@ -759,6 +759,11 @@ class TestGetAssetEventByAssetAliasPartitionKey: PostgreSQL (~), MySQL (REGEXP), and SQLite (re.match). """ + @pytest.fixture(autouse=True) + def _enable_regexp_query_filters(self): + with conf_vars({("api", "enable_regexp_query_filters"): "True"}): + yield + @pytest.fixture def test_partitioned_alias_events(self, session, test_asset): def make_timestamp(day): @@ -832,6 +837,15 @@ def test_get_by_alias_with_invalid_regex_returns_400(self, client): assert response.status_code == 400 assert "Invalid regex" in response.json()["detail"]["reason"] + def test_get_by_alias_with_pattern_disabled_returns_400(self, client): + with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + response = client.get( + "/execution/asset-events/by-asset-alias", + params={"name": "partitioned_alias", "partition_key_pattern": "^us"}, + ) + assert response.status_code == 400 + assert "disabled" in response.json()["detail"]["message"] + @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") def test_get_by_alias_with_exact_partition_key_regex(self, client): response = client.get( diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index fe61155f39e4c..b3110006a9c9e 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -363,12 +363,16 @@ 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 PartitionKeyPattern = string | null; export type Type29 = "GetAssetEventByAsset"; export type AliasName = string; 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 PartitionKeyPattern1 = string | null; export type Type30 = "GetAssetEventByAssetAlias"; export type Name11 = string; export type Key6 = string; @@ -574,7 +578,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; @@ -1205,6 +1209,8 @@ export interface GetAssetEventByAsset { before?: Before; limit?: Limit; ascending?: Ascending; + partition_key?: PartitionKey5; + partition_key_pattern?: PartitionKeyPattern; type?: Type29; } /** @@ -1217,6 +1223,8 @@ export interface GetAssetEventByAssetAlias { before?: Before1; limit?: Limit1; ascending?: Ascending1; + partition_key?: PartitionKey6; + partition_key_pattern?: PartitionKeyPattern1; type?: Type30; } /** @@ -1783,7 +1791,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; From 89e26adde82a5b5ee8ce402fbc4649f1df09e108 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Tue, 7 Jul 2026 23:54:20 +0200 Subject: [PATCH 07/13] Address review feedback on partition key regexp filter - Rename the public/SDK parameter partition_key_pattern -> partition_key_regexp_pattern to distinguish regexp filtering from the substring "*_pattern" filters on other APIs. - Move the enable_regexp_query_filters gate into _RegexParam construction and the regex compile check into depends_regex, so a regexp filter can never be instantiated without the setting being enabled. - Drop the mutually-exclusive 400 between partition_key and partition_key_regexp_pattern (both now combine with AND) on the Core and Execution APIs and in the Task SDK client/accessor. - Reuse the shared QueryAssetEventPartitionKeyFilter/Regex params in the Execution API asset-event routes instead of duplicating the query params and validation. - Keep backend implementation details out of the public parameter description; simplify the config docs; revert an unrelated limit check. - Regenerate OpenAPI specs, UI TS clients, and the supervisor schema snapshots. --- .../docs/authoring-and-scheduling/assets.rst | 6 +- .../airflow/api_fastapi/common/parameters.py | 56 +++++------ .../openapi/v2-rest-api-generated.yaml | 14 ++- .../core_api/routes/public/assets.py | 14 +-- .../execution_api/routes/asset_events.py | 97 +++++-------------- .../src/airflow/config_templates/config.yml | 18 ++-- .../airflow/ui/openapi-gen/queries/common.ts | 4 +- .../ui/openapi-gen/queries/ensureQueryData.ts | 4 +- .../ui/openapi-gen/queries/prefetch.ts | 4 +- .../airflow/ui/openapi-gen/queries/queries.ts | 4 +- .../ui/openapi-gen/queries/suspense.ts | 4 +- .../ui/openapi-gen/requests/services.gen.ts | 4 +- .../ui/openapi-gen/requests/types.gen.ts | 4 +- .../core_api/routes/public/test_assets.py | 45 +++++---- .../versions/head/test_asset_events.py | 58 +++++------ task-sdk/src/airflow/sdk/api/client.py | 8 +- .../src/airflow/sdk/execution_time/comms.py | 4 +- .../src/airflow/sdk/execution_time/context.py | 14 ++- .../sdk/execution_time/schema/schema.json | 8 +- .../airflow/sdk/execution_time/supervisor.py | 4 +- task-sdk/tests/task_sdk/api/test_client.py | 16 +-- .../execution_time/test_supervisor.py | 24 ++--- ts-sdk/src/generated/supervisor.ts | 16 ++- 23 files changed, 190 insertions(+), 240 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 99f57e3d1bd05..05b4acfca8bb9 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -247,17 +247,17 @@ 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 matching a specific partition key pattern (using database-native regex): +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_pattern(r"^us\|") + 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 (no regex), use ``.partition_key(value)`` instead. The two are mutually exclusive; setting one clears the other. +For an exact partition key match, use ``.partition_key(value)`` instead. Regexp filtering is opt-in and must be enabled with the ``[api] enable_regexp_query_filters`` setting; see the config for the security trade-off. You can also filter events by their ``extra`` key-value pairs: diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 39536056f9b9e..41b937c285d76 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -521,15 +521,21 @@ class _RegexParam(BaseParam[str]): """ Filter using database-level regex matching (regexp_match). - SQLAlchemy's ``regexp_match`` is supported on PostgreSQL (``~`` operator), - MySQL/MariaDB (``REGEXP``), and SQLite (via Python ``re.match``). - Note: SQLite uses ``re.match`` semantics (anchored at the start of the - string), while PostgreSQL uses ``re.search`` (matches anywhere). + The pattern is handed to the database's own regex engine (via SQLAlchemy's + ``regexp_match``), so this filter is gated behind the + ``[api] enable_regexp_query_filters`` setting to contain the ReDoS attack + surface: it cannot be instantiated with a value unless the setting is enabled. """ - def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: - super().__init__(skip_none=skip_none) + 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 not conf.getboolean("api", "enable_regexp_query_filters"): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Regexp query filters are disabled. " + "Set [api] enable_regexp_query_filters = True to use them.", + ) def to_orm(self, select: Select) -> Select: if self.value is None and self.skip_none: @@ -541,30 +547,7 @@ def depends(cls, *args: Any, **kwargs: Any) -> Self: raise NotImplementedError("Use regex_param_factory instead, depends is not implemented.") -def _validate_regex_pattern(value: str | None) -> str | None: - """Validate that the regex pattern is enabled and syntactically correct.""" - if value is None: - return value - if not conf.getboolean("api", "enable_regexp_query_filters"): - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="Regex query filters are disabled. Set [api] enable_regexp_query_filters = True to use them.", - ) - try: - re.compile(value) - except re.error as e: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=f"Invalid regular expression: {e}", - ) - return value - - -_DEFAULT_REGEX_DESCRIPTION = ( - "Regex filter. Uses database-native regex " - "(PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). " - "Note: on SQLite, matching is anchored at the start of the string." -) +_DEFAULT_REGEX_DESCRIPTION = "Filter results by matching this regular expression against the field value." def regex_param_factory( @@ -576,8 +559,15 @@ def regex_param_factory( def depends_regex( value: str | None = Query(alias=pattern_name, default=None, description=description), ) -> _RegexParam: - _validate_regex_pattern(value) - return _RegexParam(attribute, skip_none).set_value(value) + 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}", + ) + return _RegexParam(attribute, value, skip_none) return depends_regex @@ -1710,7 +1700,7 @@ def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | N Depends(filter_param_factory(AssetEvent.partition_key, str | None, filter_name="partition_key")), ] QueryAssetEventPartitionKeyRegex = Annotated[ - _RegexParam, Depends(regex_param_factory(AssetEvent.partition_key, "partition_key_pattern")) + _RegexParam, Depends(regex_param_factory(AssetEvent.partition_key, "partition_key_regexp_pattern")) ] 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 fdc6597433a62..6343821b66fe3 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 @@ -441,20 +441,18 @@ paths: - type: string - type: 'null' title: Partition Key - - name: partition_key_pattern + - name: partition_key_regexp_pattern in: query required: false schema: anyOf: - type: string - type: 'null' - description: 'Regex filter. Uses database-native regex (PostgreSQL ~ operator, - MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored - at the start of the string.' - title: Partition Key Pattern - description: 'Regex filter. Uses database-native regex (PostgreSQL ~ operator, - MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at - the start of the string.' + 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 ad752bc7ecbd7..5ef9e31f6c018 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 @@ -324,7 +324,7 @@ def get_asset_events( FilterParam[int | None], Depends(filter_param_factory(AssetEvent.source_map_index, int | None)) ], partition_key: QueryAssetEventPartitionKeyFilter, - partition_key_pattern: QueryAssetEventPartitionKeyRegex, + partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex, name_pattern: QueryAssetNamePatternSearch, name_prefix_pattern: QueryAssetNamePrefixPatternSearch, extra_filter: QueryAssetEventExtraFilter, @@ -332,14 +332,8 @@ def get_asset_events( session: SessionDep, ) -> AssetEventCollectionResponse: """Get asset events.""" - if partition_key.value is not None and partition_key_pattern.value is not None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail="partition_key and partition_key_pattern are mutually exclusive.", - ) - - if partition_key_pattern.value is not None: - # Bound the runtime of the user-supplied regex to guard against ReDoS (PostgreSQL only). + if partition_key_regexp_pattern.value is not None: + # Bound the runtime of the user-supplied regex to guard against ReDoS. apply_regex_query_timeout(session) base_statement = select(AssetEvent) @@ -355,7 +349,7 @@ def get_asset_events( source_run_id, source_map_index, partition_key, - partition_key_pattern, + partition_key_regexp_pattern, name_pattern, name_prefix_pattern, extra_filter, 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 010f778338079..3bfe7299ed8cd 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 @@ -17,20 +17,23 @@ from __future__ import annotations -import re from typing import Annotated from fastapi import APIRouter, HTTPException, Query, status 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 ( AssetEventResponse, AssetEventsResponse, ) -from airflow.configuration import conf from airflow.models.asset import AssetAliasModel, AssetEvent, AssetModel from airflow.utils.sqlalchemy import JsonContains, apply_regex_query_timeout @@ -43,11 +46,19 @@ 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) return AssetEventsResponse.model_validate( @@ -76,40 +87,6 @@ def _get_asset_events_through_sql_clauses( ) -def _validate_partition_key_params(partition_key: str | None, partition_key_pattern: str | None) -> None: - """Validate partition_key and partition_key_pattern are not both set, and that the pattern is valid regex.""" - if partition_key is not None and partition_key_pattern is not None: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "reason": "Mutually exclusive parameters", - "message": "partition_key and partition_key_pattern cannot both be provided. " - "Use partition_key for exact match or partition_key_pattern for regex filtering.", - }, - ) - if partition_key_pattern is None: - return - if not conf.getboolean("api", "enable_regexp_query_filters"): - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "reason": "Regex query filters disabled", - "message": "Regex query filters are disabled. " - "Set [api] enable_regexp_query_filters = True to use partition_key_pattern.", - }, - ) - try: - re.compile(partition_key_pattern) - except re.error as e: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail={ - "reason": "Invalid regex", - "message": f"The partition_key_pattern is not a valid regular expression: {e}", - }, - ) - - def _parse_extra_params(extra: list[str] | None) -> dict[str, str]: """Parse repeated ``key=value`` query params into a dict.""" result: dict[str, str] = {} @@ -132,22 +109,12 @@ 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, limit: Annotated[int | None, Query(description="The maximum number of results to return")] = None, - partition_key: Annotated[ - str | None, - Query(description="Exact partition key filter."), - ] = None, - partition_key_pattern: Annotated[ - str | None, - Query( - description="Partition key regex filter. Uses database-native regex " - "(PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). " - "Note: on SQLite, matching is anchored at the start of the string." - ), - ] = None, extra: Annotated[ list[str] | None, Query( @@ -178,12 +145,9 @@ def get_asset_event_by_asset_name_uri( if extra_dict: where_clause = and_(where_clause, JsonContains(AssetEvent.extra, extra_dict)) - _validate_partition_key_params(partition_key, partition_key_pattern) - if partition_key is not None: - where_clause = and_(where_clause, AssetEvent.partition_key == partition_key) - elif partition_key_pattern is not None: + if partition_key_regexp_pattern.value is not None: + # Bound the runtime of the user-supplied regex to guard against ReDoS. apply_regex_query_timeout(session) - where_clause = and_(where_clause, AssetEvent.partition_key.regexp_match(partition_key_pattern)) return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.asset, @@ -191,6 +155,7 @@ def get_asset_event_by_asset_name_uri( session=session, ascending=ascending, limit=limit, + filters=[partition_key, partition_key_regexp_pattern], ) @@ -198,22 +163,12 @@ 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, limit: Annotated[int | None, Query(description="The maximum number of results to return")] = None, - partition_key: Annotated[ - str | None, - Query(description="Exact partition key filter."), - ] = None, - partition_key_pattern: Annotated[ - str | None, - Query( - description="Partition key regex filter. Uses database-native regex " - "(PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). " - "Note: on SQLite, matching is anchored at the start of the string." - ), - ] = None, extra: Annotated[ list[str] | None, Query( @@ -230,12 +185,9 @@ def get_asset_event_by_asset_alias( if extra_dict: where_clause = and_(where_clause, JsonContains(AssetEvent.extra, extra_dict)) - _validate_partition_key_params(partition_key, partition_key_pattern) - if partition_key is not None: - where_clause = and_(where_clause, AssetEvent.partition_key == partition_key) - elif partition_key_pattern is not None: + if partition_key_regexp_pattern.value is not None: + # Bound the runtime of the user-supplied regex to guard against ReDoS. apply_regex_query_timeout(session) - where_clause = and_(where_clause, AssetEvent.partition_key.regexp_match(partition_key_pattern)) return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.source_aliases, @@ -243,4 +195,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 9ba30377f425a..dc0646220be16 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1637,18 +1637,16 @@ api: default: "True" enable_regexp_query_filters: description: | - Whether to enable regular-expression based query filters in the API. Currently this - controls the ``partition_key_pattern`` filter on ``GET /assets/events`` and on the + Whether to enable regular-expression based query filters in the API. 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. - Disabled by default for security reasons. When enabled, 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. Only enable this if you - accept that trade-off, and keep ``regexp_query_timeout`` set to a sensible value. - - Exact-match ``partition_key`` filtering is always available and is not affected by this - setting. + 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. Only enable this if you accept that trade-off, and keep + ``regexp_query_timeout`` set to a sensible value. version_added: 3.4.0 type: boolean example: ~ 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 4c094649e6359..18ab43ce73c1a 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -36,7 +36,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, partitionKey, partitionKeyPattern, 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; @@ -54,7 +54,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, partitionKey, partitionKeyPattern, 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 fc632360f490b..d7d4f9fc2ae65 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -93,7 +93,7 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * @returns AssetEventCollectionResponse Successful Response * @throws ApiError */ -export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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; @@ -111,7 +111,7 @@ export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}) => queryClient.ensureQueryData({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 b2b707130c934..cfd39b8c9fa49 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -93,7 +93,7 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * @returns AssetEventCollectionResponse Successful Response * @throws ApiError */ -export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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; @@ -111,7 +111,7 @@ export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, timestampGte?: string; timestampLt?: string; timestampLte?: string; -} = {}) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 6744807fcaae0..0756b5b16c099 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -93,7 +93,7 @@ export const useAssetServiceGetAssetAlias = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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; @@ -111,7 +111,7 @@ export const useAssetServiceGetAssetEvents = , "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 40dbd68706f19..60da340f8bbef 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -93,7 +93,7 @@ export const useAssetServiceGetAssetAliasSuspense = = unknown[]>({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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; @@ -111,7 +111,7 @@ export const useAssetServiceGetAssetEventsSuspense = , "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }, queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit, namePattern, namePrefixPattern, offset, orderBy, partitionKey, partitionKeyPattern, 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 89f66dc61394e..85e2d45c795fb 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 @@ -121,7 +121,7 @@ export class AssetService { * @param data.sourceRunId * @param data.sourceMapIndex * @param data.partitionKey - * @param data.partitionKeyPattern Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. + * @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. @@ -148,7 +148,7 @@ export class AssetService { source_run_id: data.sourceRunId, source_map_index: data.sourceMapIndex, partition_key: data.partitionKey, - partition_key_pattern: data.partitionKeyPattern, + 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 d9e81a68f71c5..8b936ccfc7557 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 @@ -2873,9 +2873,9 @@ export type GetAssetEventsData = { orderBy?: Array<(string)>; partitionKey?: string | null; /** - * Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. + * Filter results by matching this regular expression against the field value. */ - partitionKeyPattern?: string | null; + partitionKeyRegexpPattern?: string | null; sourceDagId?: string | null; sourceMapIndex?: number | null; sourceRunId?: string | null; 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 495b7c6508437..ef1f8b6ae1678 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 @@ -1079,7 +1079,7 @@ def test_should_mask_sensitive_extra(self, test_client, session): class TestGetAssetEventsPartitionKeyRegex(TestAssets): - """Tests for partition_key_pattern regex filter on GET /assets/events. + """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 @@ -1154,7 +1154,7 @@ def _create_partition_key_test_data(self, session=None): session.commit() @pytest.mark.parametrize( - ("partition_key_pattern", "expected_count"), + ("partition_key_regexp_pattern", "expected_count"), [ ("^2024-01-01$", 1), ("^2024-01-", 2), @@ -1165,39 +1165,45 @@ def _create_partition_key_test_data(self, session=None): ], ) @provide_session - def test_partition_key_pattern_filtering( - self, test_client, partition_key_pattern, expected_count, session + def test_partition_key_regexp_pattern_filtering( + self, test_client, partition_key_regexp_pattern, expected_count, session ): self._create_partition_key_test_data() - response = test_client.get("/assets/events", params={"partition_key_pattern": partition_key_pattern}) + 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_pattern": "^us\\|", "asset_id": "1"}, 1), - ({"partition_key_pattern": "^us\\|", "asset_id": "2"}, 0), - ({"partition_key_pattern": ".*\\|2024-01-01$", "source_dag_id": "d"}, 2), - ({"partition_key_pattern": ".*\\|2024-01-01$", "source_dag_id": "other"}, 0), + ({"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), ], ) @provide_session - def test_partition_key_pattern_combined_filters(self, test_client, params, expected_count, session): + def test_partition_key_regexp_pattern_combined_filters( + self, test_client, params, expected_count, session + ): self._create_partition_key_test_data() response = test_client.get("/assets/events", params=params) assert response.status_code == 200 assert response.json()["total_entries"] == expected_count @provide_session - def test_partition_key_pattern_invalid_regex_returns_400(self, test_client, session): - response = test_client.get("/assets/events", params={"partition_key_pattern": "[invalid(regex"}) + def test_partition_key_regexp_pattern_invalid_regex_returns_400(self, test_client, session): + 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_pattern_disabled_returns_400(self, test_client, session): + def test_partition_key_regexp_pattern_disabled_returns_400(self, test_client, session): with conf_vars({("api", "enable_regexp_query_filters"): "False"}): - response = test_client.get("/assets/events", params={"partition_key_pattern": "^2024-"}) + response = test_client.get("/assets/events", params={"partition_key_regexp_pattern": "^2024-"}) assert response.status_code == 400 assert "disabled" in response.json()["detail"] @@ -1212,7 +1218,7 @@ def test_exact_match_works_when_regex_disabled(self, test_client, session): @provide_session def test_partition_key_exact_match_via_regex(self, test_client, session): self._create_partition_key_test_data() - response = test_client.get("/assets/events", params={"partition_key_pattern": "^2024-01-01$"}) + 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 @@ -1238,12 +1244,15 @@ def test_partition_key_exact_match_no_match(self, test_client, session): assert response.json()["total_entries"] == 0 @provide_session - def test_partition_key_and_pattern_mutual_exclusion(self, test_client, session): + def test_partition_key_and_pattern_combined(self, test_client, session): + # Both filters are allowed and combine with AND: a disjoint pair yields no results. + self._create_partition_key_test_data() response = test_client.get( "/assets/events", - params={"partition_key": "2024-01-01", "partition_key_pattern": "^2024-"}, + params={"partition_key": "2024-01-01", "partition_key_regexp_pattern": "^2025-"}, ) - assert response.status_code == 400 + assert response.status_code == 200 + assert response.json()["total_entries"] == 0 class TestGetAssetEventsExtraFilter(TestAssets): 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 6d5956f1ecedb..6ec7402bc673c 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 @@ -526,7 +526,7 @@ def test_get_by_asset(self, client): class TestGetAssetEventByAssetPartitionKey: - """Tests for partition_key_pattern regex filter on execution API. + """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 @@ -576,7 +576,7 @@ def test_get_by_asset_with_exact_partition_key(self, client): params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": "^2024-01-01$", + "partition_key_regexp_pattern": "^2024-01-01$", }, ) assert response.status_code == 200 @@ -591,7 +591,7 @@ def test_get_by_asset_with_prefix_partition_key(self, client): params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": "^2024-01-", + "partition_key_regexp_pattern": "^2024-01-", }, ) assert response.status_code == 200 @@ -605,7 +605,7 @@ def test_get_by_asset_with_composite_partition_key_multiple_regions(self, client params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": r"^(us|eu)\|2024-01-.*", + "partition_key_regexp_pattern": r"^(us|eu)\|2024-01-.*", }, ) assert response.status_code == 200 @@ -621,7 +621,7 @@ def test_get_by_asset_with_composite_partition_key_date_across_regions(self, cli params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": r".*\|2024-01-01$", + "partition_key_regexp_pattern": r".*\|2024-01-01$", }, ) assert response.status_code == 200 @@ -637,7 +637,7 @@ def test_get_by_asset_with_partition_key_no_match(self, client): params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": "^2025-", + "partition_key_regexp_pattern": "^2025-", }, ) assert response.status_code == 200 @@ -659,7 +659,7 @@ def test_get_by_asset_with_partition_key_and_time_range(self, client): params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": "^2024-01-", + "partition_key_regexp_pattern": "^2024-01-", "after": "2021-01-01T12:00:00Z", }, ) @@ -675,11 +675,11 @@ def test_get_by_asset_with_invalid_regex_returns_400(self, client): params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": "[invalid(regex", + "partition_key_regexp_pattern": "[invalid(regex", }, ) assert response.status_code == 400 - assert "Invalid regex" in response.json()["detail"]["reason"] + assert "Invalid regular expression" in response.json()["detail"] def test_get_by_asset_with_pattern_disabled_returns_400(self, client): with conf_vars({("api", "enable_regexp_query_filters"): "False"}): @@ -688,11 +688,11 @@ def test_get_by_asset_with_pattern_disabled_returns_400(self, client): params={ "name": "test_get_asset_by_name", "uri": None, - "partition_key_pattern": "^2021-", + "partition_key_regexp_pattern": "^2021-", }, ) assert response.status_code == 400 - assert "disabled" in response.json()["detail"]["message"] + assert "disabled" in response.json()["detail"] @pytest.mark.usefixtures("test_asset", "test_partitioned_events") def test_get_by_asset_exact_partition_key(self, client): @@ -738,22 +738,23 @@ def test_get_by_asset_exact_partition_key_no_match(self, client): assert len(response.json()["asset_events"]) == 0 @pytest.mark.usefixtures("test_asset", "test_partitioned_events") - def test_get_by_asset_mutual_exclusion_partition_key_and_pattern(self, client): + 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_pattern": "^2024-", + "partition_key_regexp_pattern": "^2025-", }, ) - assert response.status_code == 400 - assert "Mutually exclusive" in response.json()["detail"]["reason"] + assert response.status_code == 200 + assert len(response.json()["asset_events"]) == 0 class TestGetAssetEventByAssetAliasPartitionKey: - """Tests for partition_key_pattern regex filter on by-asset-alias endpoint. + """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). @@ -803,7 +804,7 @@ def make_timestamp(day): 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_pattern": r"^us\|.*"}, + params={"name": "partitioned_alias", "partition_key_regexp_pattern": r"^us\|.*"}, ) assert response.status_code == 200 data = response.json() @@ -814,7 +815,7 @@ def test_get_by_alias_with_partition_key_region(self, client): 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_pattern": r".*\|2024-01-01$"}, + params={"name": "partitioned_alias", "partition_key_regexp_pattern": r".*\|2024-01-01$"}, ) assert response.status_code == 200 assert len(response.json()["asset_events"]) == 2 @@ -832,25 +833,25 @@ def test_get_by_alias_without_partition_key_returns_all(self, client): 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_pattern": "[invalid(regex"}, + params={"name": "partitioned_alias", "partition_key_regexp_pattern": "[invalid(regex"}, ) assert response.status_code == 400 - assert "Invalid regex" in response.json()["detail"]["reason"] + assert "Invalid regular expression" in response.json()["detail"] def test_get_by_alias_with_pattern_disabled_returns_400(self, client): with conf_vars({("api", "enable_regexp_query_filters"): "False"}): response = client.get( "/execution/asset-events/by-asset-alias", - params={"name": "partitioned_alias", "partition_key_pattern": "^us"}, + params={"name": "partitioned_alias", "partition_key_regexp_pattern": "^us"}, ) assert response.status_code == 400 - assert "disabled" in response.json()["detail"]["message"] + 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_pattern": r"^us\|2024-01-01$"}, + params={"name": "partitioned_alias", "partition_key_regexp_pattern": r"^us\|2024-01-01$"}, ) assert response.status_code == 200 data = response.json() @@ -878,17 +879,20 @@ def test_get_by_alias_exact_partition_key_no_match(self, client): assert len(response.json()["asset_events"]) == 0 @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events") - def test_get_by_alias_mutual_exclusion_partition_key_and_pattern(self, client): + 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_pattern": r"^us\|", + "partition_key_regexp_pattern": r"^us\|", }, ) - assert response.status_code == 400 - assert "Mutually exclusive" in response.json()["detail"]["reason"] + 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: diff --git a/task-sdk/src/airflow/sdk/api/client.py b/task-sdk/src/airflow/sdk/api/client.py index 9212f849ac880..7e681c1114956 100644 --- a/task-sdk/src/airflow/sdk/api/client.py +++ b/task-sdk/src/airflow/sdk/api/client.py @@ -859,12 +859,10 @@ def get( ascending: bool = True, limit: int | None = None, partition_key: str | None = None, - partition_key_pattern: str | None = None, + partition_key_regexp_pattern: str | None = None, extra: dict[str, str] | None = None, ) -> AssetEventsResponse: """Get Asset event from the API server.""" - if partition_key is not None and partition_key_pattern is not None: - raise ValueError("partition_key and partition_key_pattern are mutually exclusive") common_params: dict[str, Any] = {} if after: common_params["after"] = after.isoformat() @@ -875,8 +873,8 @@ def get( common_params["limit"] = limit if partition_key is not None: common_params["partition_key"] = partition_key - if partition_key_pattern is not None: - common_params["partition_key_pattern"] = partition_key_pattern + 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 f8ea5218ac968..5c4ae14dff1cf 100644 --- a/task-sdk/src/airflow/sdk/execution_time/comms.py +++ b/task-sdk/src/airflow/sdk/execution_time/comms.py @@ -1143,7 +1143,7 @@ class GetAssetEventByAsset(BaseModel): limit: int | None = None ascending: bool = True partition_key: str | None = None - partition_key_pattern: str | None = None + partition_key_regexp_pattern: str | None = None extra: dict[str, str] | None = None type: Literal["GetAssetEventByAsset"] = "GetAssetEventByAsset" @@ -1155,7 +1155,7 @@ class GetAssetEventByAssetAlias(BaseModel): limit: int | None = None ascending: bool = True partition_key: str | None = None - partition_key_pattern: 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 ab316efa783f2..4c40e1b59bb33 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -1089,7 +1089,7 @@ class InletEventsAccessor(Sequence["AssetEventResult"]): _ascending: bool _limit: int | None _partition_key: str | None - _partition_key_pattern: str | None + _partition_key_regexp_pattern: str | None _extra: dict[str, str] _asset_name: str | None _asset_uri: str | None @@ -1106,7 +1106,7 @@ def __init__( self._ascending = True self._limit = None self._partition_key = None - self._partition_key_pattern = None + self._partition_key_regexp_pattern = None self._extra: dict[str, str] = {} def after(self, after: str) -> Self: @@ -1132,14 +1132,12 @@ def limit(self, limit: int) -> Self: def partition_key(self, key: str) -> Self: """Filter by exact partition key match.""" self._partition_key = key - self._partition_key_pattern = None self._reset_cache() return self - def partition_key_pattern(self, pattern: str) -> Self: - """Filter by partition key regex pattern.""" - self._partition_key_pattern = pattern - self._partition_key = None + 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 @@ -1164,7 +1162,7 @@ def _asset_events(self) -> list[AssetEventResult]: "ascending": self._ascending, "limit": self._limit, "partition_key": self._partition_key, - "partition_key_pattern": self._partition_key_pattern, + "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 d4fe172163708..084259d95e945 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1858,7 +1858,7 @@ "default": null, "title": "Partition Key" }, - "partition_key_pattern": { + "partition_key_regexp_pattern": { "anyOf": [ { "type": "string" @@ -1868,7 +1868,7 @@ } ], "default": null, - "title": "Partition Key Pattern" + "title": "Partition Key Regexp Pattern" }, "extra": { "anyOf": [ @@ -1960,7 +1960,7 @@ "default": null, "title": "Partition Key" }, - "partition_key_pattern": { + "partition_key_regexp_pattern": { "anyOf": [ { "type": "string" @@ -1970,7 +1970,7 @@ } ], "default": null, - "title": "Partition Key Pattern" + "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 30f064b35bb4f..87311f02da7a1 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -1775,7 +1775,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: ascending=msg.ascending, limit=msg.limit, partition_key=msg.partition_key, - partition_key_pattern=msg.partition_key_pattern, + partition_key_regexp_pattern=msg.partition_key_regexp_pattern, extra=msg.extra, ) asset_event_result = AssetEventsResult.from_asset_events_response(asset_event_resp) @@ -1789,7 +1789,7 @@ def _handle_request(self, msg: ToSupervisor, log: FilteringBoundLogger, req_id: ascending=msg.ascending, limit=msg.limit, partition_key=msg.partition_key, - partition_key_pattern=msg.partition_key_pattern, + 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 7fe95dc16d766..0f9b8130e2654 100644 --- a/task-sdk/tests/task_sdk/api/test_client.py +++ b/task-sdk/tests/task_sdk/api/test_client.py @@ -1220,7 +1220,7 @@ 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_pattern" not in params + assert "partition_key_regexp_pattern" not in params return httpx.Response( status_code=200, json={ @@ -1246,10 +1246,10 @@ def handle_request(request: httpx.Request) -> httpx.Response: assert isinstance(result, AssetEventsResponse) assert len(result.asset_events) == 1 - def test_partition_key_pattern_param_passed(self): + 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_pattern") == "^2024-01-" + assert params.get("partition_key_regexp_pattern") == "^2024-01-" assert "partition_key" not in params return httpx.Response( status_code=200, @@ -1257,15 +1257,15 @@ def handle_request(request: httpx.Request) -> httpx.Response: ) client = make_client(httpx.MockTransport(handle_request)) - result = client.asset_events.get(name="this_asset", partition_key_pattern="^2024-01-") + 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_pattern_with_other_params(self): + 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_pattern") == r"^us\|2024-.*" + 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" @@ -1278,7 +1278,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: client = make_client(httpx.MockTransport(handle_request)) result = client.asset_events.get( name="this_asset", - partition_key_pattern=r"^us\|2024-.*", + partition_key_regexp_pattern=r"^us\|2024-.*", after=datetime(2023, 6, 1, tzinfo=dt_timezone.utc), limit=5, ascending=False, @@ -1293,7 +1293,7 @@ def handle_request(request: httpx.Request) -> httpx.Response: params = request.url.params assert params.get("name") == "my_alias" assert params.get("partition_key") == "us-east|2024-01-01" - assert "partition_key_pattern" not in params + assert "partition_key_regexp_pattern" not in params return httpx.Response( status_code=200, json={"asset_events": []}, 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 b7f969511a9b9..c39ceb92b0e3b 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -1941,7 +1941,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -1987,7 +1987,7 @@ class RequestTestCase: "limit": 5, "ascending": False, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2026,7 +2026,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2072,7 +2072,7 @@ class RequestTestCase: "limit": 5, "ascending": False, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2111,7 +2111,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2157,7 +2157,7 @@ class RequestTestCase: "limit": 5, "ascending": False, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2195,7 +2195,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2239,7 +2239,7 @@ class RequestTestCase: "limit": 5, "ascending": False, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2282,7 +2282,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": "us|2024-01-15", - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2325,7 +2325,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": None, - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": {"region": "us"}, }, response=AssetEventsResult( @@ -2366,7 +2366,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": "eu|2024-03", - "partition_key_pattern": None, + "partition_key_regexp_pattern": None, "extra": None, }, response=AssetEventsResult( @@ -2407,7 +2407,7 @@ class RequestTestCase: "limit": None, "ascending": True, "partition_key": None, - "partition_key_pattern": 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 b3110006a9c9e..ad6a6bff1b53c 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -364,7 +364,10 @@ export type Before = string | null; export type Limit = number | null; export type Ascending = boolean; export type PartitionKey5 = string | null; -export type PartitionKeyPattern = string | null; +export type PartitionKeyRegexpPattern = string | null; +export type Extra7 = { + [k: string]: string; +} | null; export type Type29 = "GetAssetEventByAsset"; export type AliasName = string; export type After1 = string | null; @@ -372,7 +375,10 @@ export type Before1 = string | null; export type Limit1 = number | null; export type Ascending1 = boolean; export type PartitionKey6 = string | null; -export type PartitionKeyPattern1 = string | null; +export type PartitionKeyRegexpPattern1 = string | null; +export type Extra8 = { + [k: string]: string; +} | null; export type Type30 = "GetAssetEventByAssetAlias"; export type Name11 = string; export type Key6 = string; @@ -1210,7 +1216,8 @@ export interface GetAssetEventByAsset { limit?: Limit; ascending?: Ascending; partition_key?: PartitionKey5; - partition_key_pattern?: PartitionKeyPattern; + partition_key_regexp_pattern?: PartitionKeyRegexpPattern; + extra?: Extra7; type?: Type29; } /** @@ -1224,7 +1231,8 @@ export interface GetAssetEventByAssetAlias { limit?: Limit1; ascending?: Ascending1; partition_key?: PartitionKey6; - partition_key_pattern?: PartitionKeyPattern1; + partition_key_regexp_pattern?: PartitionKeyRegexpPattern1; + extra?: Extra8; type?: Type30; } /** From e837b15ff805fed946ba622bfbdfa2d6719f6dcf Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Wed, 8 Jul 2026 00:26:25 +0200 Subject: [PATCH 08/13] Collapse regexp config into one setting and scope the query timeout Address further review feedback: - Replace the enable_regexp_query_filters + regexp_query_timeout pair with a single [api] regexp_query_timeout: 0 (default) disables regexp filtering entirely, any positive value enables it and bounds the query runtime. This removes the "enabled but unbounded" combination. - Make apply_regex_query_timeout a context manager that sets a transaction-local statement_timeout on enter and resets it on exit, so the bound is scoped to the regexp query and does not leak to other statements in the request's transaction. Both asset-event routes apply it automatically around query execution instead of relying on a manual call in the view. --- .../docs/authoring-and-scheduling/assets.rst | 27 +++++---- .../airflow/api_fastapi/common/parameters.py | 10 ++-- .../core_api/routes/public/assets.py | 55 +++++++++---------- .../execution_api/routes/asset_events.py | 12 +--- .../src/airflow/config_templates/config.yml | 45 +++++++-------- airflow-core/src/airflow/utils/sqlalchemy.py | 36 +++++++----- .../core_api/routes/public/test_assets.py | 6 +- .../versions/head/test_asset_events.py | 8 +-- .../tests/unit/utils/test_sqlalchemy.py | 21 ++++--- 9 files changed, 109 insertions(+), 111 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 05b4acfca8bb9..8a23d76655dc0 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -257,7 +257,7 @@ The accessor also supports chaining methods to filter events before fetching the 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 and must be enabled with the ``[api] enable_regexp_query_filters`` setting; see the config for the security trade-off. +For an exact partition key match, use ``.partition_key(value)`` instead. Regexp filtering is opt-in: it is enabled only by setting ``[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: @@ -997,25 +997,24 @@ Two parameters are available: curl -G "http:///api/v2/assets/events" \ --data-urlencode "partition_key=us|2026-03-10" -- ``partition_key_pattern`` for **regex filtering** — uses database-native regex - (PostgreSQL ``~`` operator, MySQL ``REGEXP``, SQLite ``re.match``): +- ``partition_key_regexp_pattern`` for **regular-expression filtering**: .. code-block:: bash curl -G "http:///api/v2/assets/events" \ - --data-urlencode "partition_key_pattern=^us" + --data-urlencode "partition_key_regexp_pattern=^us" -These parameters are mutually exclusive; providing both returns a 400 error. +Both parameters can be combined; the conditions are applied with AND logic. .. note:: - ``partition_key_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** and must be enabled with ``[api] enable_regexp_query_filters = True``. When enabled, - the query runs on PostgreSQL under a bounded ``statement_timeout`` controlled by - ``[api] regexp_query_timeout`` (seconds); MySQL bounds regex evaluation with its built-in - ``regexp_time_limit``. Prefer the exact-match ``partition_key`` (which uses the B-tree index and - is always enabled) whenever a full key is known. + ``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 ``[api] regexp_query_timeout`` to a positive number + of seconds, which simultaneously bounds the query runtime (enforced as a ``statement_timeout`` on + PostgreSQL; MySQL bounds regex evaluation with its built-in ``regexp_time_limit``). 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``: @@ -1024,8 +1023,8 @@ The same filters are available in the ``InletEventsAccessor``: # Exact match events = inlet_events[Asset("my_asset")].partition_key("us|2026-03-10").limit(1) - # Regex pattern - events = inlet_events[Asset("my_asset")].partition_key_pattern(r"^us\|2026-03-").limit(10) + # 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/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 41b937c285d76..7c5fc52eadb1a 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -522,19 +522,19 @@ 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] enable_regexp_query_filters`` setting to contain the ReDoS attack - surface: it cannot be instantiated with a value unless the setting is enabled. + ``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). """ 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 not conf.getboolean("api", "enable_regexp_query_filters"): + if value is not None and conf.getint("api", "regexp_query_timeout") <= 0: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Regexp query filters are disabled. " - "Set [api] enable_regexp_query_filters = True to use them.", + "Set [api] regexp_query_timeout to a positive number of seconds to enable them.", ) def to_orm(self, select: Select) -> Select: 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 5ef9e31f6c018..4e728ac334207 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 @@ -332,39 +332,38 @@ def get_asset_events( session: SessionDep, ) -> AssetEventCollectionResponse: """Get asset events.""" - if partition_key_regexp_pattern.value is not None: - # Bound the runtime of the user-supplied regex to guard against ReDoS. - apply_regex_query_timeout(session) - base_statement = select(AssetEvent) if name_pattern.value or name_prefix_pattern.value: base_statement = base_statement.join(AssetModel, AssetEvent.asset_id == AssetModel.id) - assets_event_select, total_entries = paginated_select( - statement=base_statement, - filters=[ - asset_id, - source_dag_id, - source_task_id, - source_run_id, - source_map_index, - partition_key, - partition_key_regexp_pattern, - name_pattern, - name_prefix_pattern, - extra_filter, - timestamp_range, - ], - order_by=order_by, - offset=offset, - limit=limit, - session=session, - ) + # Bound the runtime of the user-supplied regex to guard against ReDoS, scoped to the queries + # executed within this block (the count query in paginated_select and the fetch below). + with apply_regex_query_timeout(session): + assets_event_select, total_entries = paginated_select( + statement=base_statement, + filters=[ + asset_id, + source_dag_id, + source_task_id, + source_run_id, + source_map_index, + partition_key, + partition_key_regexp_pattern, + name_pattern, + name_prefix_pattern, + extra_filter, + timestamp_range, + ], + order_by=order_by, + offset=offset, + limit=limit, + session=session, + ) - assets_event_select = assets_event_select.options( - subqueryload(AssetEvent.created_dagruns), joinedload(AssetEvent.asset) - ) - assets_events = session.scalars(assets_event_select) + assets_event_select = assets_event_select.options( + subqueryload(AssetEvent.created_dagruns), joinedload(AssetEvent.asset) + ) + 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 3bfe7299ed8cd..31ccf59fb2997 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 @@ -60,7 +60,9 @@ def _get_asset_events_through_sql_clauses( 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) + # Bound the runtime of any user-supplied regex filter to guard against ReDoS, scoped to this query. + with apply_regex_query_timeout(session): + asset_events = session.scalars(asset_events_query).all() return AssetEventsResponse.model_validate( { "asset_events": [ @@ -145,10 +147,6 @@ def get_asset_event_by_asset_name_uri( if extra_dict: where_clause = and_(where_clause, JsonContains(AssetEvent.extra, extra_dict)) - if partition_key_regexp_pattern.value is not None: - # Bound the runtime of the user-supplied regex to guard against ReDoS. - apply_regex_query_timeout(session) - return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.asset, where_clause=where_clause, @@ -185,10 +183,6 @@ def get_asset_event_by_asset_alias( if extra_dict: where_clause = and_(where_clause, JsonContains(AssetEvent.extra, extra_dict)) - if partition_key_regexp_pattern.value is not None: - # Bound the runtime of the user-supplied regex to guard against ReDoS. - apply_regex_query_timeout(session) - return _get_asset_events_through_sql_clauses( join_clause=AssetEvent.source_aliases, where_clause=where_clause, diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index dc0646220be16..53b901588ad16 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1635,37 +1635,30 @@ api: type: boolean example: ~ default: "True" - enable_regexp_query_filters: - description: | - Whether to enable regular-expression based query filters in the API. 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. - - 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. Only enable this if you accept that trade-off, and keep - ``regexp_query_timeout`` set to a sensible value. - version_added: 3.4.0 - type: boolean - example: ~ - default: "False" regexp_query_timeout: description: | - Maximum time, in seconds, that a query using a regular-expression filter (see - ``enable_regexp_query_filters``) is allowed to run. - - This is the primary runtime mitigation for the ReDoS risk described above: on PostgreSQL - it is enforced as a transaction-local ``statement_timeout`` so that a pathological - pattern is aborted instead of pinning a database backend (PostgreSQL has no built-in - regex step limit, unlike MySQL's ``regexp_time_limit``). Keep it as low as your - legitimate queries allow. A value of ``0`` disables the timeout, which is not - recommended while the feature is enabled. + Timeout, in seconds, for regular-expression based query filters in the API. 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: on PostgreSQL it is enforced + as a transaction-local ``statement_timeout`` so a pathological pattern is aborted instead + of pinning a database backend (PostgreSQL has no built-in regex step limit, unlike MySQL's + ``regexp_time_limit``). 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: integer example: ~ - default: "30" + 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/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index 1f15557ce8edb..0eaaa04fbd2a4 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -62,24 +62,32 @@ def get_dialect_name(session: Session) -> str | None: return getattr(bind.dialect, "name", None) -def apply_regex_query_timeout(session: Session) -> None: +@contextlib.contextmanager +def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]: """ - Bound the runtime of a user-supplied regex filter on PostgreSQL. - - Reads the ``[api] regexp_query_timeout`` config (in seconds) and sets a transaction-local - ``statement_timeout`` (via ``set_config(..., is_local=True)``) so it applies only to the current - transaction and resets automatically at commit/rollback. This is a ReDoS safeguard: a malicious - pattern is aborted instead of pinning a database backend. No-op on non-PostgreSQL backends and - when the configured timeout is ``0`` (disabled). + Bound the runtime of a user-supplied regex filter on PostgreSQL, scoped to the wrapped query. + + Reads the ``[api] regexp_query_timeout`` config (in seconds) and, on PostgreSQL, sets a + transaction-local ``statement_timeout`` (via ``set_config(..., is_local=True)``) for the + duration of the ``with`` block, then resets it to the default so it does not affect any other + statement running later in the same transaction. This is a ReDoS safeguard: a malicious pattern + is aborted instead of pinning a database backend. No-op on non-PostgreSQL backends and when the + configured timeout is ``0`` (which also means regexp filtering is disabled). """ timeout_seconds = conf.getint("api", "regexp_query_timeout") - if timeout_seconds <= 0: + if timeout_seconds <= 0 or get_dialect_name(session) != "postgresql": + yield return - if get_dialect_name(session) == "postgresql": - timeout_ms = timeout_seconds * 1000 - session.execute( - text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) - ) + timeout_ms = timeout_seconds * 1000 + session.execute( + text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) + ) + try: + yield + finally: + # Reset to the default for the remainder of the transaction so the timeout only bounds the + # regexp query and not unrelated statements in the same request. + session.execute(text("SELECT set_config('statement_timeout', '0', true)")) def build_upsert_stmt( 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 ef1f8b6ae1678..fdfb5ca2bc61a 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 @@ -1088,7 +1088,7 @@ class TestGetAssetEventsPartitionKeyRegex(TestAssets): @pytest.fixture(autouse=True) def _enable_regexp_query_filters(self): - with conf_vars({("api", "enable_regexp_query_filters"): "True"}): + with conf_vars({("api", "regexp_query_timeout"): "30"}): yield @provide_session @@ -1202,7 +1202,7 @@ def test_partition_key_regexp_pattern_invalid_regex_returns_400(self, test_clien assert "Invalid regular expression" in response.json()["detail"] def test_partition_key_regexp_pattern_disabled_returns_400(self, test_client, session): - with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + 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"] @@ -1210,7 +1210,7 @@ def test_partition_key_regexp_pattern_disabled_returns_400(self, test_client, se @provide_session def test_exact_match_works_when_regex_disabled(self, test_client, session): self._create_partition_key_test_data() - with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + 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 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 6ec7402bc673c..ce835a81ad0c6 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 @@ -536,7 +536,7 @@ class TestGetAssetEventByAssetPartitionKey: @pytest.fixture(autouse=True) def _enable_regexp_query_filters(self): - with conf_vars({("api", "enable_regexp_query_filters"): "True"}): + with conf_vars({("api", "regexp_query_timeout"): "30"}): yield @pytest.fixture @@ -682,7 +682,7 @@ def test_get_by_asset_with_invalid_regex_returns_400(self, client): assert "Invalid regular expression" in response.json()["detail"] def test_get_by_asset_with_pattern_disabled_returns_400(self, client): - with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + with conf_vars({("api", "regexp_query_timeout"): "0"}): response = client.get( "/execution/asset-events/by-asset", params={ @@ -762,7 +762,7 @@ class TestGetAssetEventByAssetAliasPartitionKey: @pytest.fixture(autouse=True) def _enable_regexp_query_filters(self): - with conf_vars({("api", "enable_regexp_query_filters"): "True"}): + with conf_vars({("api", "regexp_query_timeout"): "30"}): yield @pytest.fixture @@ -839,7 +839,7 @@ def test_get_by_alias_with_invalid_regex_returns_400(self, client): assert "Invalid regular expression" in response.json()["detail"] def test_get_by_alias_with_pattern_disabled_returns_400(self, client): - with conf_vars({("api", "enable_regexp_query_filters"): "False"}): + 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"}, diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index fd7f74ff13b7e..f135ce75db131 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -391,23 +391,28 @@ def _mock_session(dialect_name): session.get_bind.return_value.dialect.name = dialect_name return session - def test_sets_statement_timeout_on_postgresql(self): + def test_sets_and_resets_statement_timeout_on_postgresql(self): session = self._mock_session("postgresql") with conf_vars({("api", "regexp_query_timeout"): "5"}): - apply_regex_query_timeout(session) - session.execute.assert_called_once() - stmt = session.execute.call_args.args[0] - # 5 seconds -> 5000 ms, passed as a bound parameter (no SQL injection). - assert stmt.compile().params == {"timeout": "5000"} + with apply_regex_query_timeout(session): + # Timeout is set on enter. + assert session.execute.call_count == 1 + set_stmt = session.execute.call_args.args[0] + # 5 seconds -> 5000 ms, passed as a bound parameter (no SQL injection). + assert set_stmt.compile().params == {"timeout": "5000"} + # Reset on exit so the timeout does not leak to other statements in the transaction. + assert session.execute.call_count == 2 def test_noop_on_non_postgresql(self): session = self._mock_session("mysql") with conf_vars({("api", "regexp_query_timeout"): "5"}): - apply_regex_query_timeout(session) + 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"}): - apply_regex_query_timeout(session) + with apply_regex_query_timeout(session): + pass session.execute.assert_not_called() From 586c210c178e3f202f517f12acf92146097a71c2 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Wed, 8 Jul 2026 11:08:59 +0200 Subject: [PATCH 09/13] Regenerate UI query client for partition_key_regexp_pattern rename The hand-applied rename missed the type-block declarations and JSDoc in the generated query hooks (they were not on lines matching the AssetEvents hook name). Regenerate with `pnpm codegen` so the UI client matches the OpenAPI spec, fixing the "Compile / format / lint UI" static check. --- airflow-core/src/airflow/ui/openapi-gen/queries/common.ts | 2 +- .../src/airflow/ui/openapi-gen/queries/ensureQueryData.ts | 4 ++-- airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts | 4 ++-- airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts | 4 ++-- airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts | 4 ++-- 5 files changed, 9 insertions(+), 9 deletions(-) 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 18ab43ce73c1a..9a378ad19355c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -45,7 +45,7 @@ export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit, name offset?: number; orderBy?: string[]; partitionKey?: string; - partitionKeyPattern?: string; + partitionKeyRegexpPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; 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 d7d4f9fc2ae65..6cdac537dc095 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -80,7 +80,7 @@ export const ensureUseAssetServiceGetAssetAliasData = (queryClient: QueryClient, * @param data.sourceRunId * @param data.sourceMapIndex * @param data.partitionKey -* @param data.partitionKeyPattern Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. +* @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. @@ -102,7 +102,7 @@ export const ensureUseAssetServiceGetAssetEventsData = (queryClient: QueryClient offset?: number; orderBy?: string[]; partitionKey?: string; - partitionKeyPattern?: string; + partitionKeyRegexpPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; 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 cfd39b8c9fa49..d08b21910bca3 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -80,7 +80,7 @@ export const prefetchUseAssetServiceGetAssetAlias = (queryClient: QueryClient, { * @param data.sourceRunId * @param data.sourceMapIndex * @param data.partitionKey -* @param data.partitionKeyPattern Regex filter. Uses database-native regex (PostgreSQL ~ operator, MySQL REGEXP, SQLite re.match). Note: on SQLite, matching is anchored at the start of the string. +* @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. @@ -102,7 +102,7 @@ export const prefetchUseAssetServiceGetAssetEvents = (queryClient: QueryClient, offset?: number; orderBy?: string[]; partitionKey?: string; - partitionKeyPattern?: string; + partitionKeyRegexpPattern?: string; sourceDagId?: string; sourceMapIndex?: number; sourceRunId?: string; 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 0756b5b16c099..f70c1738041dd 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -80,7 +80,7 @@ export const useAssetServiceGetAssetAlias = Date: Wed, 8 Jul 2026 12:30:45 +0200 Subject: [PATCH 10/13] Fix provide_session positional static check in partition-key tests The check-no-new-provide-session-positional hook flagged the partition-key regexp tests. The nine test methods never used the injected session, so drop their @provide_session decorator and unused session parameter; make session keyword-only in the _create_partition_key_test_data helper that does use it. --- .../core_api/routes/public/test_assets.py | 33 +++++++------------ 1 file changed, 11 insertions(+), 22 deletions(-) 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 575ca7524a886..e46bce67056e0 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 @@ -1093,7 +1093,7 @@ def _enable_regexp_query_filters(self): yield @provide_session - def _create_partition_key_test_data(self, session=None): + def _create_partition_key_test_data(self, *, session=None): _create_assets(session=session) events = [ AssetEvent( @@ -1165,9 +1165,8 @@ def _create_partition_key_test_data(self, session=None): ("^nonexistent", 0), ], ) - @provide_session def test_partition_key_regexp_pattern_filtering( - self, test_client, partition_key_regexp_pattern, expected_count, session + self, test_client, partition_key_regexp_pattern, expected_count ): self._create_partition_key_test_data() response = test_client.get( @@ -1185,67 +1184,57 @@ def test_partition_key_regexp_pattern_filtering( ({"partition_key_regexp_pattern": ".*\\|2024-01-01$", "source_dag_id": "other"}, 0), ], ) - @provide_session - def test_partition_key_regexp_pattern_combined_filters( - self, test_client, params, expected_count, session - ): + def test_partition_key_regexp_pattern_combined_filters(self, test_client, params, expected_count): self._create_partition_key_test_data() response = test_client.get("/assets/events", params=params) assert response.status_code == 200 assert response.json()["total_entries"] == expected_count - @provide_session - def test_partition_key_regexp_pattern_invalid_regex_returns_400(self, test_client, session): + 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, session): + 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"] - @provide_session - def test_exact_match_works_when_regex_disabled(self, test_client, session): + def test_exact_match_works_when_regex_disabled(self, test_client): self._create_partition_key_test_data() 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 - @provide_session - def test_partition_key_exact_match_via_regex(self, test_client, session): + def test_partition_key_exact_match_via_regex(self, test_client): self._create_partition_key_test_data() 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 - @provide_session - def test_partition_key_exact_match(self, test_client, session): + def test_partition_key_exact_match(self, test_client): self._create_partition_key_test_data() response = test_client.get("/assets/events", params={"partition_key": "2024-01-01"}) assert response.status_code == 200 assert response.json()["total_entries"] == 1 - @provide_session - def test_partition_key_exact_match_composite(self, test_client, session): + def test_partition_key_exact_match_composite(self, test_client): self._create_partition_key_test_data() response = test_client.get("/assets/events", params={"partition_key": "us|2024-01-01"}) assert response.status_code == 200 assert response.json()["total_entries"] == 1 - @provide_session - def test_partition_key_exact_match_no_match(self, test_client, session): + def test_partition_key_exact_match_no_match(self, test_client): self._create_partition_key_test_data() response = test_client.get("/assets/events", params={"partition_key": "nonexistent"}) assert response.status_code == 200 assert response.json()["total_entries"] == 0 - @provide_session - def test_partition_key_and_pattern_combined(self, test_client, session): + def test_partition_key_and_pattern_combined(self, test_client): # Both filters are allowed and combine with AND: a disjoint pair yields no results. self._create_partition_key_test_data() response = test_client.get( From 21bdb4485d054d23535d3d4ba6635744ad745dcf Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sat, 11 Jul 2026 20:19:13 +0200 Subject: [PATCH 11/13] Address review: MySQL query timeout, float config, doc link, test tidy-up - Enforce regexp_query_timeout on MySQL too (session max_execution_time), not just PostgreSQL statement_timeout; apply_regex_query_timeout now branches by dialect and is a no-op on SQLite. - Make [api] regexp_query_timeout a float so fractional seconds (e.g. 0.5) are allowed; read it via conf.getfloat. - Link the config option in assets.rst via :ref: and note the MySQL bound. - Tests: create partition-key rows via an autouse fixture instead of a per-test helper call, fold the exact-match cases into parametrized tests (core + execution API), and cover the MySQL/float timeout paths. --- .../docs/authoring-and-scheduling/assets.rst | 12 ++--- .../airflow/api_fastapi/common/parameters.py | 2 +- .../src/airflow/config_templates/config.yml | 21 ++++---- airflow-core/src/airflow/utils/sqlalchemy.py | 54 ++++++++++++------- .../core_api/routes/public/test_assets.py | 36 +++++-------- .../versions/head/test_asset_events.py | 45 +++++----------- .../tests/unit/utils/test_sqlalchemy.py | 21 +++++++- 7 files changed, 97 insertions(+), 94 deletions(-) diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst b/airflow-core/docs/authoring-and-scheduling/assets.rst index 8a23d76655dc0..3bd84a054d06e 100644 --- a/airflow-core/docs/authoring-and-scheduling/assets.rst +++ b/airflow-core/docs/authoring-and-scheduling/assets.rst @@ -257,7 +257,7 @@ The accessor also supports chaining methods to filter events before fetching the 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 ``[api] regexp_query_timeout`` to a positive number of seconds, which also bounds the query runtime; see the config for the security trade-off. +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: @@ -1010,11 +1010,11 @@ Both parameters can be combined; the conditions are applied with AND logic. ``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 ``[api] regexp_query_timeout`` to a positive number - of seconds, which simultaneously bounds the query runtime (enforced as a ``statement_timeout`` on - PostgreSQL; MySQL bounds regex evaluation with its built-in ``regexp_time_limit``). Prefer the - exact-match ``partition_key`` (which uses the B-tree index and is always enabled) whenever a full - key is known. + 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``: diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index 7c5fc52eadb1a..9a9a71c259e9c 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -530,7 +530,7 @@ class _RegexParam(BaseParam[str]): 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.getint("api", "regexp_query_timeout") <= 0: + 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. " diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 53b901588ad16..b93e7a727112f 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -1637,10 +1637,10 @@ api: default: "True" regexp_query_timeout: description: | - Timeout, in seconds, for regular-expression based query filters in the API. 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. + 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`` @@ -1649,14 +1649,13 @@ api: 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: on PostgreSQL it is enforced - as a transaction-local ``statement_timeout`` so a pathological pattern is aborted instead - of pinning a database backend (PostgreSQL has no built-in regex step limit, unlike MySQL's - ``regexp_time_limit``). 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. + 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: integer + type: float example: ~ default: "0" secret_key: diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index 0eaaa04fbd2a4..1333e77c9f97f 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -65,29 +65,45 @@ def get_dialect_name(session: Session) -> str | None: @contextlib.contextmanager def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]: """ - Bound the runtime of a user-supplied regex filter on PostgreSQL, scoped to the wrapped query. - - Reads the ``[api] regexp_query_timeout`` config (in seconds) and, on PostgreSQL, sets a - transaction-local ``statement_timeout`` (via ``set_config(..., is_local=True)``) for the - duration of the ``with`` block, then resets it to the default so it does not affect any other - statement running later in the same transaction. This is a ReDoS safeguard: a malicious pattern - is aborted instead of pinning a database backend. No-op on non-PostgreSQL backends and when the - configured timeout is ``0`` (which also means regexp filtering is disabled). + 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 clears it 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). + + 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.getint("api", "regexp_query_timeout") - if timeout_seconds <= 0 or get_dialect_name(session) != "postgresql": + timeout_seconds = conf.getfloat("api", "regexp_query_timeout") + if timeout_seconds <= 0: yield return - timeout_ms = timeout_seconds * 1000 - session.execute( - text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) - ) - try: + # ``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": + session.execute( + text("SELECT set_config('statement_timeout', :timeout, true)").bindparams(timeout=str(timeout_ms)) + ) + try: + yield + finally: + # Reset to the default for the remainder of the transaction so the timeout only bounds + # the regexp query and not unrelated statements in the same request. + session.execute(text("SELECT set_config('statement_timeout', '0', true)")) + elif dialect_name == "mysql": + session.execute(text(f"SET SESSION max_execution_time = {timeout_ms}")) + try: + yield + finally: + # Clear the limit so it only bounds the regexp query, not later statements in the session. + session.execute(text("SET SESSION max_execution_time = 0")) + else: yield - finally: - # Reset to the default for the remainder of the transaction so the timeout only bounds the - # regexp query and not unrelated statements in the same request. - session.execute(text("SELECT set_config('statement_timeout', '0', true)")) def build_upsert_stmt( 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 e46bce67056e0..cbda690d9293e 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 @@ -1092,8 +1092,8 @@ def _enable_regexp_query_filters(self): with conf_vars({("api", "regexp_query_timeout"): "30"}): yield - @provide_session - def _create_partition_key_test_data(self, *, session=None): + @pytest.fixture(autouse=True) + def _create_partition_key_test_data(self, setup, session): _create_assets(session=session) events = [ AssetEvent( @@ -1168,7 +1168,6 @@ def _create_partition_key_test_data(self, *, session=None): def test_partition_key_regexp_pattern_filtering( self, test_client, partition_key_regexp_pattern, expected_count ): - self._create_partition_key_test_data() response = test_client.get( "/assets/events", params={"partition_key_regexp_pattern": partition_key_regexp_pattern} ) @@ -1185,7 +1184,6 @@ def test_partition_key_regexp_pattern_filtering( ], ) def test_partition_key_regexp_pattern_combined_filters(self, test_client, params, expected_count): - self._create_partition_key_test_data() response = test_client.get("/assets/events", params=params) assert response.status_code == 200 assert response.json()["total_entries"] == expected_count @@ -1204,39 +1202,31 @@ def test_partition_key_regexp_pattern_disabled_returns_400(self, test_client): assert "disabled" in response.json()["detail"] def test_exact_match_works_when_regex_disabled(self, test_client): - self._create_partition_key_test_data() 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): - self._create_partition_key_test_data() 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 - def test_partition_key_exact_match(self, test_client): - self._create_partition_key_test_data() - 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_composite(self, test_client): - self._create_partition_key_test_data() - response = test_client.get("/assets/events", params={"partition_key": "us|2024-01-01"}) - assert response.status_code == 200 - assert response.json()["total_entries"] == 1 - - def test_partition_key_exact_match_no_match(self, test_client): - self._create_partition_key_test_data() - response = test_client.get("/assets/events", params={"partition_key": "nonexistent"}) + @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"] == 0 + 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. - self._create_partition_key_test_data() response = test_client.get( "/assets/events", params={"partition_key": "2024-01-01", "partition_key_regexp_pattern": "^2025-"}, 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 ce835a81ad0c6..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 @@ -695,47 +695,26 @@ def test_get_by_asset_with_pattern_disabled_returns_400(self, client): assert "disabled" in response.json()["detail"] @pytest.mark.usefixtures("test_asset", "test_partitioned_events") - def test_get_by_asset_exact_partition_key(self, client): - response = client.get( - "/execution/asset-events/by-asset", - params={ - "name": "test_get_asset_by_name", - "uri": None, - "partition_key": "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_exact_partition_key_composite(self, client): - response = client.get( - "/execution/asset-events/by-asset", - params={ - "name": "test_get_asset_by_name", - "uri": None, - "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_events") - def test_get_by_asset_exact_partition_key_no_match(self, client): + @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": "nonexistent", + "partition_key": partition_key, }, ) assert response.status_code == 200 - assert len(response.json()["asset_events"]) == 0 + 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): diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index f135ce75db131..1fa51e4220cb7 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -403,8 +403,27 @@ def test_sets_and_resets_statement_timeout_on_postgresql(self): # Reset on exit so the timeout does not leak to other statements in the transaction. assert session.execute.call_count == 2 - def test_noop_on_non_postgresql(self): + def test_sets_and_resets_max_execution_time_on_mysql(self): session = self._mock_session("mysql") + with conf_vars({("api", "regexp_query_timeout"): "5"}): + with apply_regex_query_timeout(session): + # max_execution_time is set (in ms) on enter. + assert session.execute.call_count == 1 + assert "max_execution_time = 5000" in str(session.execute.call_args.args[0]) + # Cleared on exit so the limit does not leak to other statements in the session. + assert session.execute.call_count == 2 + assert "max_execution_time = 0" in str(session.execute.call_args.args[0]) + + def test_fractional_seconds_are_converted_to_milliseconds(self): + session = self._mock_session("postgresql") + with conf_vars({("api", "regexp_query_timeout"): "0.5"}): + with apply_regex_query_timeout(session): + set_stmt = session.execute.call_args.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 From cffe65cf585a1b8f3ba688acc913a4db118443e2 Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sat, 11 Jul 2026 20:43:59 +0200 Subject: [PATCH 12/13] Restore previous db timeout instead of clearing it Address review nit: apply_regex_query_timeout now captures the current statement_timeout (PostgreSQL) / max_execution_time (MySQL) before overriding it and restores that value on exit, instead of resetting to 0. This preserves a server- or role-level global timeout rather than clobbering it. --- airflow-core/src/airflow/utils/sqlalchemy.py | 28 +++++++++------ .../tests/unit/utils/test_sqlalchemy.py | 34 ++++++++++++------- 2 files changed, 39 insertions(+), 23 deletions(-) diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py b/airflow-core/src/airflow/utils/sqlalchemy.py index 4a4fad617efc5..0064c0ec3c874 100644 --- a/airflow-core/src/airflow/utils/sqlalchemy.py +++ b/airflow-core/src/airflow/utils/sqlalchemy.py @@ -68,15 +68,17 @@ 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 clears it so it does not - affect any other statement running later in the same transaction/session: + 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). - 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). + 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: @@ -86,22 +88,28 @@ def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]: 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: - # Reset to the default for the remainder of the transaction so the timeout only bounds - # the regexp query and not unrelated statements in the same request. - session.execute(text("SELECT set_config('statement_timeout', '0', true)")) + # 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: - # Clear the limit so it only bounds the regexp query, not later statements in the session. - session.execute(text("SET SESSION max_execution_time = 0")) + # 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 diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py b/airflow-core/tests/unit/utils/test_sqlalchemy.py index 1fa51e4220cb7..8dbcd41313055 100644 --- a/airflow-core/tests/unit/utils/test_sqlalchemy.py +++ b/airflow-core/tests/unit/utils/test_sqlalchemy.py @@ -391,34 +391,42 @@ def _mock_session(dialect_name): session.get_bind.return_value.dialect.name = dialect_name return session - def test_sets_and_resets_statement_timeout_on_postgresql(self): + 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): - # Timeout is set on enter. - assert session.execute.call_count == 1 - set_stmt = session.execute.call_args.args[0] + # 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"} - # Reset on exit so the timeout does not leak to other statements in the transaction. - assert session.execute.call_count == 2 + # 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_resets_max_execution_time_on_mysql(self): + 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): - # max_execution_time is set (in ms) on enter. - assert session.execute.call_count == 1 + # 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]) - # Cleared on exit so the limit does not leak to other statements in the session. - assert session.execute.call_count == 2 - assert "max_execution_time = 0" 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_stmt = session.execute.call_args.args[0] + # 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"} From 87548c335523e8d63e9fd9ced08805962402b1bc Mon Sep 17 00:00:00 2001 From: Hussein Awala Date: Sat, 11 Jul 2026 22:08:19 +0200 Subject: [PATCH 13/13] Guard against unbounded regexp filters in paginated_select Address review feedback (defensive net for future misuse): paginated_select now refuses to run (HTTP 500) when a _RegexParam filter has a value but [api] regexp_query_timeout is not set, so a regexp filter wired through the shared helper can never run unbounded even if a view forgets to apply the timeout. Also document on _RegexParam that callers must time-bound the query via apply_regex_query_timeout. Add tests for the guard. --- .../airflow/api_fastapi/common/db/common.py | 28 +++++++- .../airflow/api_fastapi/common/parameters.py | 5 ++ .../unit/api_fastapi/common/db/test_common.py | 66 +++++++++++++++++++ 3 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 airflow-core/tests/unit/api_fastapi/common/db/test_common.py diff --git a/airflow-core/src/airflow/api_fastapi/common/db/common.py b/airflow-core/src/airflow/api_fastapi/common/db/common.py index ff1e3f7043e6a..5ef3f54718098 100644 --- a/airflow-core/src/airflow/api_fastapi/common/db/common.py +++ b/airflow-core/src/airflow/api_fastapi/common/db/common.py @@ -25,10 +25,11 @@ from collections.abc import AsyncGenerator, Generator, Sequence from typing import TYPE_CHECKING, Annotated, Literal, overload -from fastapi import Depends +from fastapi import Depends, HTTPException, status from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session +from airflow.configuration import conf from airflow.utils.db import get_query_count, get_query_count_async from airflow.utils.session import NEW_SESSION, create_session, create_session_async, provide_session @@ -46,6 +47,29 @@ def _get_session() -> Generator[Session, None, None]: SessionDep = Annotated[Session, Depends(_get_session, scope="function")] +def _reject_unbounded_regex_filters(filters: Sequence[OrmClause | None] | None) -> None: + """ + Refuse to run a database-side regexp filter without a configured runtime bound. + + A regexp filter hands a user-supplied pattern to the database's own regex engine, so the query + must be time-bounded (see :func:`airflow.utils.sqlalchemy.apply_regex_query_timeout`). This is a + defensive net for filters wired through :func:`paginated_select`: if a caller adds a regexp + filter but ``[api] regexp_query_timeout`` is not set, refuse to run the query rather than risk an + unbounded ReDoS. It does not cover regexp filters used without ``paginated_select``. + """ + if not filters: + return + from airflow.api_fastapi.common.parameters import _RegexParam + + if any(isinstance(f, _RegexParam) and f.value is not None for f in filters): + if conf.getfloat("api", "regexp_query_timeout") <= 0: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Refusing to run a regexp query filter without a configured " + "[api] regexp_query_timeout.", + ) + + def apply_filters_to_select( *, statement: Select, filters: Sequence[OrmClause | None] | None = None ) -> Select: @@ -103,6 +127,7 @@ async def paginated_select_async( session: AsyncSession, return_total_entries: bool = True, ) -> tuple[Select, int | None]: + _reject_unbounded_regex_filters(filters) statement = apply_filters_to_select( statement=statement, filters=filters, @@ -157,6 +182,7 @@ def paginated_select( session: Session = NEW_SESSION, return_total_entries: bool = True, ) -> tuple[Select, int | None]: + _reject_unbounded_regex_filters(filters) statement = apply_filters_to_select( statement=statement, filters=filters, diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py b/airflow-core/src/airflow/api_fastapi/common/parameters.py index c06cf0f7727c3..e858bfee51752 100644 --- a/airflow-core/src/airflow/api_fastapi/common/parameters.py +++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py @@ -525,6 +525,11 @@ class _RegexParam(BaseParam[str]): ``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). + + IMPORTANT: any endpoint that uses this filter must time-bound its query by wrapping the + execution in :func:`airflow.utils.sqlalchemy.apply_regex_query_timeout`, otherwise a + pathological pattern can pin a database backend. Filters passed through + :func:`airflow.api_fastapi.common.db.common.paginated_select` are additionally guarded there. """ def __init__(self, attribute: ColumnElement, value: str | None = None, skip_none: bool = True) -> None: diff --git a/airflow-core/tests/unit/api_fastapi/common/db/test_common.py b/airflow-core/tests/unit/api_fastapi/common/db/test_common.py new file mode 100644 index 0000000000000..176d29a60ff6a --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/common/db/test_common.py @@ -0,0 +1,66 @@ +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import pytest +from fastapi import HTTPException + +from airflow.api_fastapi.common.db.common import _reject_unbounded_regex_filters +from airflow.api_fastapi.common.parameters import FilterParam, _RegexParam +from airflow.models.asset import AssetEvent + +from tests_common.test_utils.config import conf_vars + +if TYPE_CHECKING: + from sqlalchemy.sql import ColumnElement + + +class TestRejectUnboundedRegexFilters: + """Defensive guard: a regexp filter must not run without a configured timeout.""" + + @staticmethod + def _regex_filter_with_value(value: str) -> _RegexParam: + # Build with no value (so the constructor gate does not fire), then set the value to + # simulate a caller wiring up a regexp filter regardless of the current config. + attribute = cast("ColumnElement", AssetEvent.partition_key) + return _RegexParam(attribute).set_value(value) + + def test_raises_when_regex_filter_active_and_timeout_disabled(self): + param = self._regex_filter_with_value("^us") + with conf_vars({("api", "regexp_query_timeout"): "0"}): + with pytest.raises(HTTPException) as exc_info: + _reject_unbounded_regex_filters([param]) + assert exc_info.value.status_code == 500 + assert "regexp_query_timeout" in exc_info.value.detail + + def test_allows_when_timeout_configured(self): + param = self._regex_filter_with_value("^us") + with conf_vars({("api", "regexp_query_timeout"): "30"}): + _reject_unbounded_regex_filters([param]) + + def test_ignores_regex_filter_without_value(self): + param = _RegexParam(cast("ColumnElement", AssetEvent.partition_key)) + with conf_vars({("api", "regexp_query_timeout"): "0"}): + _reject_unbounded_regex_filters([param]) + + def test_ignores_non_regex_filters(self): + non_regex = FilterParam(AssetEvent.partition_key, "2024-01-01") + with conf_vars({("api", "regexp_query_timeout"): "0"}): + _reject_unbounded_regex_filters([non_regex, None])