Add partition_key and partition_key_regexp_pattern filters for asset events

airflow-core/docs/authoring-and-scheduling/assets.rst

@@ -247,7 +247,19 @@ Inlet asset events can be read with the ``inlet_events`` accessor in the executi
 
 Each value in the ``inlet_events`` mapping is a sequence-like object that orders past events of a given asset by ``timestamp``, earliest to latest. It supports most of Python's list interface, so you can use ``[-1]`` to access the last event, ``[-2:]`` for the last two, etc. The accessor is lazy and only hits the database when you access items inside it.
 
-The accessor also supports chaining methods to filter events before fetching them. For example, to retrieve only events where specific ``extra`` keys match given values:
+The accessor also supports chaining methods to filter events before fetching them. For example, to retrieve only events matching a specific partition key regular expression:
+
+.. code-block:: python
+
+    @task(inlets=[regional_sales])
+    def process_us_sales(*, inlet_events):
+        us_events = inlet_events[regional_sales].partition_key_regexp_pattern(r"^us\|")
+        for event in us_events:
+            print(event.extra, event.partition_key)
+
+For an exact partition key match, use ``.partition_key(value)`` instead. Regexp filtering is opt-in: it is enabled only by setting :ref:`[api] regexp_query_timeout <config:api__regexp_query_timeout>` to a positive number of seconds, which also bounds the query runtime; see the config for the security trade-off.
+
+You can also filter events by their ``extra`` key-value pairs:
 
 .. code-block:: python
 
@@ -975,6 +987,45 @@ When a runtime run emits exactly one partition key, the producing
 these events the same way as timetable-produced partitions, through
 ``PartitionedAssetTimetable``.
 
+You can also query asset events filtered by partition key using the REST API.
+Two parameters are available:
+
+- ``partition_key`` for **exact match** — uses the B-tree index for fast lookups:
+
+.. code-block:: bash
+
+    curl -G "http://<airflow-host>/api/v2/assets/events" \
+      --data-urlencode "partition_key=us|2026-03-10"
+
+- ``partition_key_regexp_pattern`` for **regular-expression filtering**:
+
+.. code-block:: bash
+
+    curl -G "http://<airflow-host>/api/v2/assets/events" \
+      --data-urlencode "partition_key_regexp_pattern=^us"
+
+Both parameters can be combined; the conditions are applied with AND logic.
+
+.. note::
+
+    ``partition_key_regexp_pattern`` is evaluated by the database's own regular-expression engine,
+    which is a Regular expression Denial of Service (ReDoS) surface. For that reason it is **disabled
+    by default**: it is enabled only by setting :ref:`[api] regexp_query_timeout <config:api__regexp_query_timeout>`
+    to a positive number of seconds (fractional values allowed), which simultaneously bounds the query
+    runtime (enforced as a ``statement_timeout`` on PostgreSQL and as ``max_execution_time`` on MySQL).
+    Prefer the exact-match ``partition_key`` (which uses the B-tree index and is always enabled)
+    whenever a full key is known.
+
+The same filters are available in the ``InletEventsAccessor``:
+
+.. code-block:: python
+
+    # Exact match
+    events = inlet_events[Asset("my_asset")].partition_key("us|2026-03-10").limit(1)
+
+    # Regular-expression pattern
+    events = inlet_events[Asset("my_asset")].partition_key_regexp_pattern(r"^us\|2026-03-").limit(10)
+
 Fan-out mappers
 ~~~~~~~~~~~~~~~
 

airflow-core/docs/migrations-ref.rst

@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | Revision ID             | Revises ID       | Airflow Version   | Description                                                  |
 +=========================+==================+===================+==============================================================+
-| ``c4e7a1f9b2d0`` (head) | ``436dc127462c`` | ``3.4.0``         | Add index on asset.uri.                                      |
+| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0``         | Add index on asset_event (asset_id, partition_key).          |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``c4e7a1f9b2d0``        | ``436dc127462c`` | ``3.4.0``         | Add index on asset.uri.                                      |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+
 | ``436dc127462c``        | ``5a5d3253e946`` | ``3.4.0``         | Drop span_status column.                                     |
 +-------------------------+------------------+-------------------+--------------------------------------------------------------+

airflow-core/src/airflow/api_fastapi/common/parameters.py

@@ -17,8 +17,9 @@
 
 from __future__ import annotations
 
+import re
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Generator, Iterable, Sequence
 from datetime import datetime
 from enum import Enum
 from typing import (
@@ -41,6 +42,7 @@
 from sqlalchemy.sql.functions import FunctionElement
 
 from airflow._shared.timezones import timezone
+from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT
 from airflow.api_fastapi.core_api.base import OrmClause
 from airflow.api_fastapi.core_api.security import GetUserDep
@@ -68,7 +70,7 @@
 from airflow.models.variable import Variable
 from airflow.models.xcom import XComModel
 from airflow.typing_compat import Self
-from airflow.utils.sqlalchemy import JsonContains
+from airflow.utils.sqlalchemy import JsonContains, apply_regex_query_timeout
 from airflow.utils.state import DagRunState, TaskInstanceState
 from airflow.utils.types import DagRunType
 
@@ -516,6 +518,75 @@ def depends_search(
     return depends_search
 
 
+class _RegexParam(BaseParam[str]):
+    """
+    Filter using database-level regex matching (regexp_match).
+
+    The pattern is handed to the database's own regex engine (via SQLAlchemy's
+    ``regexp_match``), so this filter is gated behind the ``[api] regexp_query_timeout``
+    setting to contain the ReDoS attack surface: it cannot be instantiated with a value
+    unless a positive timeout is configured (which both enables the feature and bounds it).
+
+    Use :func:`regex_param_factory` to build the FastAPI dependency for this filter. That
+    dependency also applies :func:`airflow.utils.sqlalchemy.apply_regex_query_timeout` to the
+    request's session, so the query runtime is bounded automatically and callers never need to
+    remember to do it in the view.
+    """
+
+    def __init__(self, attribute: ColumnElement, value: str | None = None, skip_none: bool = True) -> None:
+        super().__init__(value=value, skip_none=skip_none)
+        self.attribute: ColumnElement = attribute
+        if value is not None and conf.getfloat("api", "regexp_query_timeout") <= 0:
+            raise HTTPException(
+                status_code=status.HTTP_400_BAD_REQUEST,
+                detail="Regexp query filters are disabled. "
+                "Set [api] regexp_query_timeout to a positive number of seconds to enable them.",
+            )
+
+    def to_orm(self, select: Select) -> Select:
+        if self.value is None and self.skip_none:
+            return select
+        return select.where(self.attribute.regexp_match(self.value))
+
+    @classmethod
+    def depends(cls, *args: Any, **kwargs: Any) -> Self:
+        raise NotImplementedError("Use regex_param_factory instead, depends is not implemented.")
+
+
+_DEFAULT_REGEX_DESCRIPTION = "Filter results by matching this regular expression against the field value."
+
+
+def regex_param_factory(
+    attribute: ColumnElement,
+    pattern_name: str,
+    skip_none: bool = True,
+    description: str = _DEFAULT_REGEX_DESCRIPTION,
+) -> Callable[..., Generator[_RegexParam, None, None]]:
+    def depends_regex(
+        session: SessionDep,
+        value: str | None = Query(alias=pattern_name, default=None, description=description),
+    ) -> Generator[_RegexParam, None, None]:
+        if value is not None:
+            try:
+                re.compile(value)
+            except re.error as e:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail=f"Invalid regular expression: {e}",
+                )
+        # ``__init__`` rejects the request (400) when a pattern is supplied while the feature is off.
+        param = _RegexParam(attribute, value, skip_none)
+        if value is None:
+            yield param
+            return
+        # Bound the query runtime for the whole request so a view can never forget to do it; the
+        # previous timeout is restored on teardown (after the response) before the session closes.
+        with apply_regex_query_timeout(session):
+            yield param
+
+    return depends_regex
+
+
 def prefix_search_param_factory(
     attribute: ColumnElement,
     prefix_pattern_name: str,
@@ -1687,6 +1758,15 @@ def _transform_ti_states(states: list[str] | None) -> list[TaskInstanceState | N
 QueryAssetAliasNamePrefixPatternSearch = Annotated[
     _PrefixSearchParam, Depends(prefix_search_param_factory(AssetAliasModel.name, "name_prefix_pattern"))
 ]
+QueryAssetEventPartitionKeyFilter = Annotated[
+    FilterParam[str | None],
+    Depends(filter_param_factory(AssetEvent.partition_key, str | None, filter_name="partition_key")),
+]
+QueryAssetEventPartitionKeyRegex = Annotated[
+    _RegexParam,
+    # ``function`` scope so the dependency can depend on the (function-scoped) session it bounds.
+    Depends(regex_param_factory(AssetEvent.partition_key, "partition_key_regexp_pattern"), scope="function"),
+]
 QueryAssetDagIdPatternSearch = Annotated[
     _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter.depends)
 ]

airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml

@@ -447,6 +447,26 @@ paths:
           - type: integer
           - type: 'null'
           title: Source Map Index
+      - name: partition_key
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          title: Partition Key
+      - name: partition_key_regexp_pattern
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: string
+          - type: 'null'
+          description: Filter results by matching this regular expression against
+            the field value.
+          title: Partition Key Regexp Pattern
+        description: Filter results by matching this regular expression against the
+          field value.
       - name: name_pattern
         in: query
         required: false

airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py

@@ -38,6 +38,8 @@
     QueryAssetAliasNamePrefixPatternSearch,
     QueryAssetDagIdPatternSearch,
     QueryAssetEventExtraFilter,
+    QueryAssetEventPartitionKeyFilter,
+    QueryAssetEventPartitionKeyRegex,
     QueryAssetNamePatternSearch,
     QueryAssetNamePrefixPatternSearch,
     QueryLimit,
@@ -331,13 +333,17 @@ def get_asset_events(
     source_map_index: Annotated[
         FilterParam[int | None], Depends(filter_param_factory(AssetEvent.source_map_index, int | None))
     ],
+    partition_key: QueryAssetEventPartitionKeyFilter,
+    partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex,
     name_pattern: QueryAssetNamePatternSearch,
     name_prefix_pattern: QueryAssetNamePrefixPatternSearch,
     extra_filter: QueryAssetEventExtraFilter,
     timestamp_range: Annotated[RangeFilter, Depends(datetime_range_filter_factory("timestamp", AssetEvent))],
     session: SessionDep,
 ) -> AssetEventCollectionResponse:
     """Get asset events."""
+    # The regexp partition-key filter bounds the query runtime automatically (its dependency applies
+    # ``apply_regex_query_timeout`` to this request's session), so no explicit wrapping is needed here.
     base_statement = select(AssetEvent)
     if name_pattern.value or name_prefix_pattern.value:
         base_statement = base_statement.join(AssetModel, AssetEvent.asset_id == AssetModel.id)
@@ -350,6 +356,8 @@ def get_asset_events(
             source_task_id,
             source_run_id,
             source_map_index,
+            partition_key,
+            partition_key_regexp_pattern,
             name_pattern,
             name_prefix_pattern,
             extra_filter,
@@ -364,7 +372,9 @@ def get_asset_events(
     assets_event_select = assets_event_select.options(
         subqueryload(AssetEvent.created_dagruns), joinedload(AssetEvent.asset)
     )
-    assets_events = session.scalars(assets_event_select)
+    # Materialize here (not lazily during response serialization) so the regexp query runs while the
+    # dependency-applied timeout is still active.
+    assets_events = session.scalars(assets_event_select).all()
 
     return AssetEventCollectionResponse(
         asset_events=assets_events,

airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py

@@ -23,6 +23,11 @@
 from sqlalchemy import and_, select
 
 from airflow.api_fastapi.common.db.common import SessionDep
+from airflow.api_fastapi.common.parameters import (
+    BaseParam,
+    QueryAssetEventPartitionKeyFilter,
+    QueryAssetEventPartitionKeyRegex,
+)
 from airflow.api_fastapi.common.types import UtcDateTime
 from airflow.api_fastapi.execution_api.datamodels.asset import AssetResponse
 from airflow.api_fastapi.execution_api.datamodels.asset_event import (
@@ -41,13 +46,23 @@
 
 
 def _get_asset_events_through_sql_clauses(
-    *, join_clause, where_clause, session: SessionDep, ascending: bool = True, limit: int | None = None
+    *,
+    join_clause,
+    where_clause,
+    session: SessionDep,
+    ascending: bool = True,
+    limit: int | None = None,
+    filters: list[BaseParam] | None = None,
 ) -> AssetEventsResponse:
     order_by_clause = AssetEvent.timestamp.asc() if ascending else AssetEvent.timestamp.desc()
     asset_events_query = select(AssetEvent).join(join_clause).where(where_clause).order_by(order_by_clause)
-    if limit is not None:
+    for filter_ in filters or []:
+        asset_events_query = filter_.to_orm(asset_events_query)
+    if limit:
         asset_events_query = asset_events_query.limit(limit)
-    asset_events = session.scalars(asset_events_query)
+    # A regexp partition-key filter bounds the query runtime automatically (its dependency applies
+    # apply_regex_query_timeout to this request's session), so no explicit wrapping is needed here.
+    asset_events = session.scalars(asset_events_query).all()
     return AssetEventsResponse.model_validate(
         {
             "asset_events": [
@@ -96,6 +111,8 @@ def get_asset_event_by_asset_name_uri(
     name: Annotated[str | None, Query(description="The name of the Asset")],
     uri: Annotated[str | None, Query(description="The URI of the Asset")],
     session: SessionDep,
+    partition_key: QueryAssetEventPartitionKeyFilter,
+    partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex,
     after: Annotated[UtcDateTime | None, Query(description="The start of the time range")] = None,
     before: Annotated[UtcDateTime | None, Query(description="The end of the time range")] = None,
     ascending: Annotated[bool, Query(description="Whether to sort results in ascending order")] = True,
@@ -136,13 +153,16 @@ def get_asset_event_by_asset_name_uri(
         session=session,
         ascending=ascending,
         limit=limit,
+        filters=[partition_key, partition_key_regexp_pattern],
     )
 
 
 @router.get("/by-asset-alias")
 def get_asset_event_by_asset_alias(
     name: Annotated[str, Query(description="The name of the Asset Alias")],
     session: SessionDep,
+    partition_key: QueryAssetEventPartitionKeyFilter,
+    partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex,
     after: Annotated[UtcDateTime | None, Query(description="The start of the time range")] = None,
     before: Annotated[UtcDateTime | None, Query(description="The end of the time range")] = None,
     ascending: Annotated[bool, Query(description="Whether to sort results in ascending order")] = True,
@@ -169,4 +189,5 @@ def get_asset_event_by_asset_alias(
         session=session,
         ascending=ascending,
         limit=limit,
+        filters=[partition_key, partition_key_regexp_pattern],
     )

airflow-core/src/airflow/config_templates/config.yml

@@ -1635,6 +1635,29 @@ api:
       type: boolean
       example: ~
       default: "True"
+    regexp_query_timeout:
+      description: |
+        Timeout, in seconds, for regular-expression based query filters in the API. Fractional
+        values are allowed (e.g. ``0.5``). This single value both enables the feature and bounds its
+        runtime: a value of ``0`` (the default) disables regexp filtering entirely, and any positive
+        value enables it while capping how long such a query may run.
+
+        When enabled, it allows additional filtering features that require database-side regexp
+        matching, such as the ``partition_key_regexp_pattern`` filter on ``GET /assets/events``
+        and on the Execution API asset-event endpoints. It is disabled by default for security
+        reasons: a user-supplied regular expression is passed to the database's own regex engine,
+        which is a Regular expression Denial of Service (ReDoS) vector: a crafted pattern can
+        force the engine into pathological backtracking and consume database backend CPU.
+
+        The timeout is the primary runtime mitigation for that risk: it is enforced as a
+        transaction-local ``statement_timeout`` on PostgreSQL and as the session ``max_execution_time``
+        on MySQL, so a pathological pattern is aborted instead of pinning a database backend. Keep it
+        as low as your legitimate queries allow. Coupling the on/off switch with the timeout
+        intentionally makes it impossible to enable regexp filtering without a runtime bound in place.
+      version_added: 3.4.0
+      type: float
+      example: ~
+      default: "0"
     secret_key:
       description: |
         Secret key used to run your api server. It should be as random as possible. However, when running