This is an automated email from the ASF dual-hosted git repository.
hussein-awala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 40311ffe4d7 Add partition_key and partition_key_regexp_pattern filters
for asset events (#64610)
40311ffe4d7 is described below
commit 40311ffe4d7fb19da91115959233dc218b3dbb15
Author: Hussein Awala <[email protected]>
AuthorDate: Mon Jul 20 15:57:32 2026 +0200
Add partition_key and partition_key_regexp_pattern filters for asset events
(#64610)
* 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
* Regenerate supervisor schema snapshot for partition_key filters
* 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.
* 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.
* 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.
* 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).
* 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.
* 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.
* 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.
* 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.
* 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.
* 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.
* Apply regexp query timeout automatically via the filter dependency
Address review feedback: instead of relying on each view to wrap its query
in
apply_regex_query_timeout (easy to forget -> ReDoS), the regexp filter's
dependency (regex_param_factory) now applies the timeout to the request's
session itself. Any endpoint using QueryAssetEventPartitionKeyRegex is
bounded
automatically. The dependency is function-scoped so it can depend on the
session it bounds; the previous timeout value is restored on teardown.
Removes
the manual wrappers from the core and execution asset-event routes and adds
tests for the dependency behavior.
---
.../docs/authoring-and-scheduling/assets.rst | 53 +++-
airflow-core/docs/migrations-ref.rst | 4 +-
.../src/airflow/api_fastapi/common/parameters.py | 84 ++++-
.../core_api/openapi/v2-rest-api-generated.yaml | 20 ++
.../api_fastapi/core_api/routes/public/assets.py | 12 +-
.../execution_api/routes/asset_events.py | 27 +-
.../src/airflow/config_templates/config.yml | 23 ++
...0127_3_4_0_add_idx_asset_event_partition_key.py | 51 +++
airflow-core/src/airflow/models/asset.py | 1 +
.../src/airflow/ui/openapi-gen/queries/common.ts | 6 +-
.../ui/openapi-gen/queries/ensureQueryData.ts | 8 +-
.../src/airflow/ui/openapi-gen/queries/prefetch.ts | 8 +-
.../src/airflow/ui/openapi-gen/queries/queries.ts | 8 +-
.../src/airflow/ui/openapi-gen/queries/suspense.ts | 8 +-
.../ui/openapi-gen/requests/services.gen.ts | 4 +
.../airflow/ui/openapi-gen/requests/types.gen.ts | 5 +
airflow-core/src/airflow/utils/db.py | 2 +-
airflow-core/src/airflow/utils/sqlalchemy.py | 52 +++
.../unit/api_fastapi/common/test_parameters.py | 53 ++++
.../core_api/routes/public/test_assets.py | 156 +++++++++
.../versions/head/test_asset_events.py | 351 +++++++++++++++++++++
airflow-core/tests/unit/utils/test_sqlalchemy.py | 63 ++++
task-sdk/src/airflow/sdk/api/client.py | 6 +
task-sdk/src/airflow/sdk/execution_time/comms.py | 4 +
task-sdk/src/airflow/sdk/execution_time/context.py | 18 ++
.../airflow/sdk/execution_time/schema/schema.json | 48 +++
.../src/airflow/sdk/execution_time/supervisor.py | 4 +
task-sdk/tests/task_sdk/api/test_client.py | 89 ++++++
.../task_sdk/execution_time/test_supervisor.py | 104 ++++++
ts-sdk/src/generated/supervisor.ts | 12 +-
30 files changed, 1263 insertions(+), 21 deletions(-)
diff --git a/airflow-core/docs/authoring-and-scheduling/assets.rst
b/airflow-core/docs/authoring-and-scheduling/assets.rst
index 2a7990d233b..3bd84a054d0 100644
--- a/airflow-core/docs/authoring-and-scheduling/assets.rst
+++ b/airflow-core/docs/authoring-and-scheduling/assets.rst
@@ -247,7 +247,19 @@ Inlet asset events can be read with the ``inlet_events``
accessor in the executi
Each value in the ``inlet_events`` mapping is a sequence-like object that
orders past events of a given asset by ``timestamp``, earliest to latest. It
supports most of Python's list interface, so you can use ``[-1]`` to access the
last event, ``[-2:]`` for the last two, etc. The accessor is lazy and only hits
the database when you access items inside it.
-The accessor also supports chaining methods to filter events before fetching
them. For example, to retrieve only events where specific ``extra`` keys match
given values:
+The accessor also supports chaining methods to filter events before fetching
them. For example, to retrieve only events matching a specific partition key
regular expression:
+
+.. code-block:: python
+
+ @task(inlets=[regional_sales])
+ def process_us_sales(*, inlet_events):
+ us_events =
inlet_events[regional_sales].partition_key_regexp_pattern(r"^us\|")
+ for event in us_events:
+ print(event.extra, event.partition_key)
+
+For an exact partition key match, use ``.partition_key(value)`` instead.
Regexp filtering is opt-in: it is enabled only by setting :ref:`[api]
regexp_query_timeout <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
~~~~~~~~~~~~~~~
diff --git a/airflow-core/docs/migrations-ref.rst
b/airflow-core/docs/migrations-ref.rst
index 3dd6e96a7d0..de64f8123f0 100644
--- a/airflow-core/docs/migrations-ref.rst
+++ b/airflow-core/docs/migrations-ref.rst
@@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are
executed via when you ru
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| Revision ID | Revises ID | Airflow Version | Description
|
+=========================+==================+===================+==============================================================+
-| ``c4e7a1f9b2d0`` (head) | ``436dc127462c`` | ``3.4.0`` | Add index
on asset.uri. |
+| ``7a98f1b7dbd3`` (head) | ``c4e7a1f9b2d0`` | ``3.4.0`` | Add index
on asset_event (asset_id, partition_key). |
++-------------------------+------------------+-------------------+--------------------------------------------------------------+
+| ``c4e7a1f9b2d0`` | ``436dc127462c`` | ``3.4.0`` | Add index
on asset.uri. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
| ``436dc127462c`` | ``5a5d3253e946`` | ``3.4.0`` | Drop
span_status column. |
+-------------------------+------------------+-------------------+--------------------------------------------------------------+
diff --git a/airflow-core/src/airflow/api_fastapi/common/parameters.py
b/airflow-core/src/airflow/api_fastapi/common/parameters.py
index 060d9478a39..31471d9e26a 100644
--- a/airflow-core/src/airflow/api_fastapi/common/parameters.py
+++ b/airflow-core/src/airflow/api_fastapi/common/parameters.py
@@ -17,8 +17,9 @@
from __future__ import annotations
+import re
from abc import ABC, abstractmethod
-from collections.abc import Callable, Iterable, Sequence
+from collections.abc import Callable, Generator, Iterable, Sequence
from datetime import datetime
from enum import Enum
from typing import (
@@ -41,6 +42,7 @@ from sqlalchemy.inspection import inspect
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.taskinstance import TaskInstance
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 search_param_factory(
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 @@ QueryAssetAliasNamePatternSearch = Annotated[
QueryAssetAliasNamePrefixPatternSearch = Annotated[
_PrefixSearchParam,
Depends(prefix_search_param_factory(AssetAliasModel.name,
"name_prefix_pattern"))
]
+QueryAssetEventPartitionKeyFilter = Annotated[
+ FilterParam[str | None],
+ Depends(filter_param_factory(AssetEvent.partition_key, str | None,
filter_name="partition_key")),
+]
+QueryAssetEventPartitionKeyRegex = Annotated[
+ _RegexParam,
+ # ``function`` scope so the dependency can depend on the (function-scoped)
session it bounds.
+ Depends(regex_param_factory(AssetEvent.partition_key,
"partition_key_regexp_pattern"), scope="function"),
+]
QueryAssetDagIdPatternSearch = Annotated[
_DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter.depends)
]
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
index 46860d9d1cb..4d6be859a13 100644
---
a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
+++
b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml
@@ -447,6 +447,26 @@ paths:
- type: integer
- type: 'null'
title: Source Map Index
+ - name: partition_key
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ - type: 'null'
+ title: Partition Key
+ - name: partition_key_regexp_pattern
+ in: query
+ required: false
+ schema:
+ anyOf:
+ - type: string
+ - type: 'null'
+ description: Filter results by matching this regular expression
against
+ the field value.
+ title: Partition Key Regexp Pattern
+ description: Filter results by matching this regular expression
against the
+ field value.
- name: name_pattern
in: query
required: false
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py
index 424938af86d..434ab8f1fdd 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py
@@ -38,6 +38,8 @@ from airflow.api_fastapi.common.parameters import (
QueryAssetAliasNamePrefixPatternSearch,
QueryAssetDagIdPatternSearch,
QueryAssetEventExtraFilter,
+ QueryAssetEventPartitionKeyFilter,
+ QueryAssetEventPartitionKeyRegex,
QueryAssetNamePatternSearch,
QueryAssetNamePrefixPatternSearch,
QueryLimit,
@@ -331,6 +333,8 @@ def get_asset_events(
source_map_index: Annotated[
FilterParam[int | None],
Depends(filter_param_factory(AssetEvent.source_map_index, int | None))
],
+ partition_key: QueryAssetEventPartitionKeyFilter,
+ partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex,
name_pattern: QueryAssetNamePatternSearch,
name_prefix_pattern: QueryAssetNamePrefixPatternSearch,
extra_filter: QueryAssetEventExtraFilter,
@@ -338,6 +342,8 @@ def get_asset_events(
session: SessionDep,
) -> AssetEventCollectionResponse:
"""Get asset events."""
+ # The regexp partition-key filter bounds the query runtime automatically
(its dependency applies
+ # ``apply_regex_query_timeout`` to this request's session), so no explicit
wrapping is needed here.
base_statement = select(AssetEvent)
if name_pattern.value or name_prefix_pattern.value:
base_statement = base_statement.join(AssetModel, AssetEvent.asset_id
== AssetModel.id)
@@ -350,6 +356,8 @@ def get_asset_events(
source_task_id,
source_run_id,
source_map_index,
+ partition_key,
+ partition_key_regexp_pattern,
name_pattern,
name_prefix_pattern,
extra_filter,
@@ -364,7 +372,9 @@ def get_asset_events(
assets_event_select = assets_event_select.options(
subqueryload(AssetEvent.created_dagruns), joinedload(AssetEvent.asset)
)
- assets_events = session.scalars(assets_event_select)
+ # Materialize here (not lazily during response serialization) so the
regexp query runs while the
+ # dependency-applied timeout is still active.
+ assets_events = session.scalars(assets_event_select).all()
return AssetEventCollectionResponse(
asset_events=assets_events,
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py
index b2a98a9de60..4fd73b30672 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py
@@ -23,6 +23,11 @@ from 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 (
@@ -41,13 +46,23 @@ router = APIRouter(
def _get_asset_events_through_sql_clauses(
- *, join_clause, where_clause, session: SessionDep, ascending: bool = True,
limit: int | None = None
+ *,
+ join_clause,
+ where_clause,
+ session: SessionDep,
+ ascending: bool = True,
+ limit: int | None = None,
+ filters: list[BaseParam] | None = None,
) -> AssetEventsResponse:
order_by_clause = AssetEvent.timestamp.asc() if ascending else
AssetEvent.timestamp.desc()
asset_events_query =
select(AssetEvent).join(join_clause).where(where_clause).order_by(order_by_clause)
- if limit is not None:
+ for filter_ in filters or []:
+ asset_events_query = filter_.to_orm(asset_events_query)
+ if limit:
asset_events_query = asset_events_query.limit(limit)
- asset_events = session.scalars(asset_events_query)
+ # A regexp partition-key filter bounds the query runtime automatically
(its dependency applies
+ # apply_regex_query_timeout to this request's session), so no explicit
wrapping is needed here.
+ asset_events = session.scalars(asset_events_query).all()
return AssetEventsResponse.model_validate(
{
"asset_events": [
@@ -96,6 +111,8 @@ def get_asset_event_by_asset_name_uri(
name: Annotated[str | None, Query(description="The name of the Asset")],
uri: Annotated[str | None, Query(description="The URI of the Asset")],
session: SessionDep,
+ partition_key: QueryAssetEventPartitionKeyFilter,
+ partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex,
after: Annotated[UtcDateTime | None, Query(description="The start of the
time range")] = None,
before: Annotated[UtcDateTime | None, Query(description="The end of the
time range")] = None,
ascending: Annotated[bool, Query(description="Whether to sort results in
ascending order")] = True,
@@ -136,6 +153,7 @@ def get_asset_event_by_asset_name_uri(
session=session,
ascending=ascending,
limit=limit,
+ filters=[partition_key, partition_key_regexp_pattern],
)
@@ -143,6 +161,8 @@ def get_asset_event_by_asset_name_uri(
def get_asset_event_by_asset_alias(
name: Annotated[str, Query(description="The name of the Asset Alias")],
session: SessionDep,
+ partition_key: QueryAssetEventPartitionKeyFilter,
+ partition_key_regexp_pattern: QueryAssetEventPartitionKeyRegex,
after: Annotated[UtcDateTime | None, Query(description="The start of the
time range")] = None,
before: Annotated[UtcDateTime | None, Query(description="The end of the
time range")] = None,
ascending: Annotated[bool, Query(description="Whether to sort results in
ascending order")] = True,
@@ -169,4 +189,5 @@ def get_asset_event_by_asset_alias(
session=session,
ascending=ascending,
limit=limit,
+ filters=[partition_key, partition_key_regexp_pattern],
)
diff --git a/airflow-core/src/airflow/config_templates/config.yml
b/airflow-core/src/airflow/config_templates/config.yml
index af11f9fe701..7a65cedb2da 100644
--- a/airflow-core/src/airflow/config_templates/config.yml
+++ b/airflow-core/src/airflow/config_templates/config.yml
@@ -1635,6 +1635,29 @@ api:
type: boolean
example: ~
default: "True"
+ regexp_query_timeout:
+ description: |
+ Timeout, in seconds, for regular-expression based query filters in the
API. Fractional
+ values are allowed (e.g. ``0.5``). This single value both enables the
feature and bounds its
+ runtime: a value of ``0`` (the default) disables regexp filtering
entirely, and any positive
+ value enables it while capping how long such a query may run.
+
+ When enabled, it allows additional filtering features that require
database-side regexp
+ matching, such as the ``partition_key_regexp_pattern`` filter on ``GET
/assets/events``
+ and on the Execution API asset-event endpoints. It is disabled by
default for security
+ reasons: a user-supplied regular expression is passed to the
database's own regex engine,
+ which is a Regular expression Denial of Service (ReDoS) vector: a
crafted pattern can
+ force the engine into pathological backtracking and consume database
backend CPU.
+
+ The timeout is the primary runtime mitigation for that risk: it is
enforced as a
+ transaction-local ``statement_timeout`` on PostgreSQL and as the
session ``max_execution_time``
+ on MySQL, so a pathological pattern is aborted instead of pinning a
database backend. Keep it
+ as low as your legitimate queries allow. Coupling the on/off switch
with the timeout
+ intentionally makes it impossible to enable regexp filtering without a
runtime bound in place.
+ version_added: 3.4.0
+ type: float
+ example: ~
+ default: "0"
secret_key:
description: |
Secret key used to run your api server. It should be as random as
possible. However, when running
diff --git
a/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py
b/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py
new file mode 100644
index 00000000000..d8763debbb9
--- /dev/null
+++
b/airflow-core/src/airflow/migrations/versions/0127_3_4_0_add_idx_asset_event_partition_key.py
@@ -0,0 +1,51 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Add index on asset_event (asset_id, partition_key).
+
+Revision ID: 7a98f1b7dbd3
+Revises: c4e7a1f9b2d0
+Create Date: 2026-04-01 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "7a98f1b7dbd3"
+down_revision = "c4e7a1f9b2d0"
+branch_labels = None
+depends_on = None
+airflow_version = "3.4.0"
+
+
+def upgrade():
+ """Apply Add index on asset_event (asset_id, partition_key)."""
+ op.create_index(
+ "idx_asset_event_asset_id_partition_key",
+ "asset_event",
+ ["asset_id", "partition_key"],
+ )
+
+
+def downgrade():
+ """Unapply Add index on asset_event (asset_id, partition_key)."""
+ op.drop_index("idx_asset_event_asset_id_partition_key",
table_name="asset_event")
diff --git a/airflow-core/src/airflow/models/asset.py
b/airflow-core/src/airflow/models/asset.py
index 0788611d8a0..7aafddb3b94 100644
--- a/airflow-core/src/airflow/models/asset.py
+++ b/airflow-core/src/airflow/models/asset.py
@@ -828,6 +828,7 @@ class AssetEvent(Base):
__tablename__ = "asset_event"
__table_args__ = (
Index("idx_asset_id_timestamp", asset_id, timestamp),
+ Index("idx_asset_event_asset_id_partition_key", asset_id,
partition_key),
{"sqlite_autoincrement": True}, # ensures PK values not reused
)
diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
index 8ecba7c6df5..fa4853f6050 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts
@@ -37,7 +37,7 @@ export const UseAssetServiceGetAssetAliasKeyFn = ({
assetAliasId }: {
export type AssetServiceGetAssetEventsDefaultResponse =
Awaited<ReturnType<typeof AssetService.getAssetEvents>>;
export type AssetServiceGetAssetEventsQueryResult<TData =
AssetServiceGetAssetEventsDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
export const useAssetServiceGetAssetEventsKey = "AssetServiceGetAssetEvents";
-export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit,
namePattern, namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex,
sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte
}: {
+export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId, extra, limit,
namePattern, namePrefixPattern, offset, orderBy, partitionKey,
partitionKeyRegexpPattern, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: {
assetId?: number;
extra?: string[];
limit?: number;
@@ -45,6 +45,8 @@ export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId,
extra, limit, name
namePrefixPattern?: string;
offset?: number;
orderBy?: string[];
+ partitionKey?: string;
+ partitionKeyRegexpPattern?: string;
sourceDagId?: string;
sourceMapIndex?: number;
sourceRunId?: string;
@@ -53,7 +55,7 @@ export const UseAssetServiceGetAssetEventsKeyFn = ({ assetId,
extra, limit, name
timestampGte?: string;
timestampLt?: string;
timestampLte?: string;
-} = {}, queryKey?: Array<unknown>) => [useAssetServiceGetAssetEventsKey,
...(queryKey ?? [{ assetId, extra, limit, namePattern, namePrefixPattern,
offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId,
timestampGt, timestampGte, timestampLt, timestampLte }])];
+} = {}, queryKey?: Array<unknown>) => [useAssetServiceGetAssetEventsKey,
...(queryKey ?? [{ assetId, extra, limit, namePattern, namePrefixPattern,
offset, orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId,
sourceMapIndex, sourceRunId, sourceTaskId, timestampGt, timestampGte,
timestampLt, timestampLte }])];
export type AssetServiceGetAssetQueuedEventsDefaultResponse =
Awaited<ReturnType<typeof AssetService.getAssetQueuedEvents>>;
export type AssetServiceGetAssetQueuedEventsQueryResult<TData =
AssetServiceGetAssetQueuedEventsDefaultResponse, TError = unknown> =
UseQueryResult<TData, TError>;
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 1664d510fc5..f83894f1bef 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts
@@ -81,6 +81,8 @@ export const ensureUseAssetServiceGetAssetAliasData =
(queryClient: QueryClient,
* @param data.sourceTaskId
* @param data.sourceRunId
* @param data.sourceMapIndex
+* @param data.partitionKey
+* @param data.partitionKeyRegexpPattern Filter results by matching this
regular expression against the field value.
* @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g.
`%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`).
Regular expressions are **not** supported.
*
* **Performance note:** this full-match pattern is evaluated as ``ILIKE
'%term%'`` and most of the time prevents the database from using B-tree
indexes, which can be very slow on large tables. Prefer the equivalent
``name_prefix_pattern`` parameter when possible.
@@ -93,7 +95,7 @@ export const ensureUseAssetServiceGetAssetAliasData =
(queryClient: QueryClient,
* @returns AssetEventCollectionResponse Successful Response
* @throws ApiError
*/
-export const ensureUseAssetServiceGetAssetEventsData = (queryClient:
QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset,
orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt,
timestampGte, timestampLt, timestampLte }: {
+export const ensureUseAssetServiceGetAssetEventsData = (queryClient:
QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset,
orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex,
sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte
}: {
assetId?: number;
extra?: string[];
limit?: number;
@@ -101,6 +103,8 @@ export const ensureUseAssetServiceGetAssetEventsData =
(queryClient: QueryClient
namePrefixPattern?: string;
offset?: number;
orderBy?: string[];
+ partitionKey?: string;
+ partitionKeyRegexpPattern?: string;
sourceDagId?: string;
sourceMapIndex?: number;
sourceRunId?: string;
@@ -109,7 +113,7 @@ export const ensureUseAssetServiceGetAssetEventsData =
(queryClient: QueryClient
timestampGte?: string;
timestampLt?: string;
timestampLte?: string;
-} = {}) => queryClient.ensureQueryData({ queryKey:
Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn:
() => AssetService.getAssetEvents({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, [...]
+} = {}) => 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, [...]
/**
* 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 d0cecdeb3d5..1b4b2d46aac 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -81,6 +81,8 @@ export const prefetchUseAssetServiceGetAssetAlias =
(queryClient: QueryClient, {
* @param data.sourceTaskId
* @param data.sourceRunId
* @param data.sourceMapIndex
+* @param data.partitionKey
+* @param data.partitionKeyRegexpPattern Filter results by matching this
regular expression against the field value.
* @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g.
`%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`).
Regular expressions are **not** supported.
*
* **Performance note:** this full-match pattern is evaluated as ``ILIKE
'%term%'`` and most of the time prevents the database from using B-tree
indexes, which can be very slow on large tables. Prefer the equivalent
``name_prefix_pattern`` parameter when possible.
@@ -93,7 +95,7 @@ export const prefetchUseAssetServiceGetAssetAlias =
(queryClient: QueryClient, {
* @returns AssetEventCollectionResponse Successful Response
* @throws ApiError
*/
-export const prefetchUseAssetServiceGetAssetEvents = (queryClient:
QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset,
orderBy, sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt,
timestampGte, timestampLt, timestampLte }: {
+export const prefetchUseAssetServiceGetAssetEvents = (queryClient:
QueryClient, { assetId, extra, limit, namePattern, namePrefixPattern, offset,
orderBy, partitionKey, partitionKeyRegexpPattern, sourceDagId, sourceMapIndex,
sourceRunId, sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte
}: {
assetId?: number;
extra?: string[];
limit?: number;
@@ -101,6 +103,8 @@ export const prefetchUseAssetServiceGetAssetEvents =
(queryClient: QueryClient,
namePrefixPattern?: string;
offset?: number;
orderBy?: string[];
+ partitionKey?: string;
+ partitionKeyRegexpPattern?: string;
sourceDagId?: string;
sourceMapIndex?: number;
sourceRunId?: string;
@@ -109,7 +113,7 @@ export const prefetchUseAssetServiceGetAssetEvents =
(queryClient: QueryClient,
timestampGte?: string;
timestampLt?: string;
timestampLte?: string;
-} = {}) => queryClient.prefetchQuery({ queryKey:
Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }), queryFn:
() => AssetService.getAssetEvents({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, t [...]
+} = {}) => 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, so [...]
/**
* 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 c859fb7260a..3d334e51a96 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts
@@ -81,6 +81,8 @@ export const useAssetServiceGetAssetAlias = <TData =
Common.AssetServiceGetAsset
* @param data.sourceTaskId
* @param data.sourceRunId
* @param data.sourceMapIndex
+* @param data.partitionKey
+* @param data.partitionKeyRegexpPattern Filter results by matching this
regular expression against the field value.
* @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g.
`%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`).
Regular expressions are **not** supported.
*
* **Performance note:** this full-match pattern is evaluated as ``ILIKE
'%term%'`` and most of the time prevents the database from using B-tree
indexes, which can be very slow on large tables. Prefer the equivalent
``name_prefix_pattern`` parameter when possible.
@@ -93,7 +95,7 @@ export const useAssetServiceGetAssetAlias = <TData =
Common.AssetServiceGetAsset
* @returns AssetEventCollectionResponse Successful Response
* @throws ApiError
*/
-export const useAssetServiceGetAssetEvents = <TData =
Common.AssetServiceGetAssetEventsDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: {
+export const useAssetServiceGetAssetEvents = <TData =
Common.AssetServiceGetAssetEventsDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern,
sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt,
timestampGte, timestampLt, timestampLte }: {
assetId?: number;
extra?: string[];
limit?: number;
@@ -101,6 +103,8 @@ export const useAssetServiceGetAssetEvents = <TData =
Common.AssetServiceGetAsse
namePrefixPattern?: string;
offset?: number;
orderBy?: string[];
+ partitionKey?: string;
+ partitionKeyRegexpPattern?: string;
sourceDagId?: string;
sourceMapIndex?: number;
sourceRunId?: string;
@@ -109,7 +113,7 @@ export const useAssetServiceGetAssetEvents = <TData =
Common.AssetServiceGetAsse
timestampGte?: string;
timestampLt?: string;
timestampLte?: string;
-} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ queryKey:
Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte },
queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit,
namePattern, namePrefixPattern, offset, orde [...]
+} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useQuery<TData, TError>({ 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, nam [...]
/**
* 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 4ca06b9e87c..a83ac5cb0ae 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts
@@ -81,6 +81,8 @@ export const useAssetServiceGetAssetAliasSuspense = <TData =
Common.AssetService
* @param data.sourceTaskId
* @param data.sourceRunId
* @param data.sourceMapIndex
+* @param data.partitionKey
+* @param data.partitionKeyRegexpPattern Filter results by matching this
regular expression against the field value.
* @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards (e.g.
`%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`).
Regular expressions are **not** supported.
*
* **Performance note:** this full-match pattern is evaluated as ``ILIKE
'%term%'`` and most of the time prevents the database from using B-tree
indexes, which can be very slow on large tables. Prefer the equivalent
``name_prefix_pattern`` parameter when possible.
@@ -93,7 +95,7 @@ export const useAssetServiceGetAssetAliasSuspense = <TData =
Common.AssetService
* @returns AssetEventCollectionResponse Successful Response
* @throws ApiError
*/
-export const useAssetServiceGetAssetEventsSuspense = <TData =
Common.AssetServiceGetAssetEventsDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte }: {
+export const useAssetServiceGetAssetEventsSuspense = <TData =
Common.AssetServiceGetAssetEventsDefaultResponse, TError = unknown, TQueryKey
extends Array<unknown> = unknown[]>({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, partitionKey, partitionKeyRegexpPattern,
sourceDagId, sourceMapIndex, sourceRunId, sourceTaskId, timestampGt,
timestampGte, timestampLt, timestampLte }: {
assetId?: number;
extra?: string[];
limit?: number;
@@ -101,6 +103,8 @@ export const useAssetServiceGetAssetEventsSuspense = <TData
= Common.AssetServic
namePrefixPattern?: string;
offset?: number;
orderBy?: string[];
+ partitionKey?: string;
+ partitionKeyRegexpPattern?: string;
sourceDagId?: string;
sourceMapIndex?: number;
sourceRunId?: string;
@@ -109,7 +113,7 @@ export const useAssetServiceGetAssetEventsSuspense = <TData
= Common.AssetServic
timestampGte?: string;
timestampLt?: string;
timestampLte?: string;
-} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ queryKey:
Common.UseAssetServiceGetAssetEventsKeyFn({ assetId, extra, limit, namePattern,
namePrefixPattern, offset, orderBy, sourceDagId, sourceMapIndex, sourceRunId,
sourceTaskId, timestampGt, timestampGte, timestampLt, timestampLte },
queryKey), queryFn: () => AssetService.getAssetEvents({ assetId, extra, limit,
namePattern, namePrefixPattern, offs [...]
+} = {}, queryKey?: TQueryKey, options?: Omit<UseQueryOptions<TData, TError>,
"queryKey" | "queryFn">) => useSuspenseQuery<TData, TError>({ 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, li [...]
/**
* 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 7008d465426..0d225056ac2 100644
--- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
+++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -122,6 +122,8 @@ export class AssetService {
* @param data.sourceTaskId
* @param data.sourceRunId
* @param data.sourceMapIndex
+ * @param data.partitionKey
+ * @param data.partitionKeyRegexpPattern Filter results by matching this
regular expression against the field value.
* @param data.namePattern SQL LIKE expression — use `%` / `_` wildcards
(e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 |
dag2`). Regular expressions are **not** supported.
*
* **Performance note:** this full-match pattern is evaluated as ``ILIKE
'%term%'`` and most of the time prevents the database from using B-tree
indexes, which can be very slow on large tables. Prefer the equivalent
``name_prefix_pattern`` parameter when possible.
@@ -147,6 +149,8 @@ export class AssetService {
source_task_id: data.sourceTaskId,
source_run_id: data.sourceRunId,
source_map_index: data.sourceMapIndex,
+ partition_key: data.partitionKey,
+ partition_key_regexp_pattern: data.partitionKeyRegexpPattern,
name_pattern: data.namePattern,
name_prefix_pattern: data.namePrefixPattern,
extra: data.extra,
diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts
index 7e8597168b8..f43c5708a08 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
@@ -2904,6 +2904,11 @@ export type GetAssetEventsData = {
* Attributes to order by, multi criteria sort is supported. Prefix with
`-` for descending order. Supported attributes: `source_task_id, source_dag_id,
source_run_id, source_map_index, timestamp`
*/
orderBy?: Array<(string)>;
+ partitionKey?: string | null;
+ /**
+ * Filter results by matching this regular expression against the field
value.
+ */
+ partitionKeyRegexpPattern?: string | null;
sourceDagId?: string | null;
sourceMapIndex?: number | null;
sourceRunId?: string | null;
diff --git a/airflow-core/src/airflow/utils/db.py
b/airflow-core/src/airflow/utils/db.py
index ce4ef036d20..1076715ee71 100644
--- a/airflow-core/src/airflow/utils/db.py
+++ b/airflow-core/src/airflow/utils/db.py
@@ -117,7 +117,7 @@ _REVISION_HEADS_MAP: dict[str, str] = {
"3.1.8": "509b94a1042d",
"3.2.0": "1d6611b6ab7c",
"3.3.0": "d2f4e1b3c5a7",
- "3.4.0": "c4e7a1f9b2d0",
+ "3.4.0": "7a98f1b7dbd3",
}
# Prefix used to identify tables holding data moved during migration.
diff --git a/airflow-core/src/airflow/utils/sqlalchemy.py
b/airflow-core/src/airflow/utils/sqlalchemy.py
index 7ca788ffe12..0064c0ec3c8 100644
--- a/airflow-core/src/airflow/utils/sqlalchemy.py
+++ b/airflow-core/src/airflow/utils/sqlalchemy.py
@@ -62,6 +62,58 @@ def get_dialect_name(session: Session) -> str | None:
return getattr(bind.dialect, "name", None)
[email protected]
+def apply_regex_query_timeout(session: Session) -> Generator[None, None, None]:
+ """
+ Bound the runtime of a user-supplied regex filter, scoped to the wrapped
query.
+
+ Reads the ``[api] regexp_query_timeout`` config (in seconds, fractional
values allowed) and sets
+ a database-side timeout for the duration of the ``with`` block, then
restores the previous value
+ so it does not affect any other statement running later in the same
transaction/session:
+
+ * PostgreSQL: transaction-local ``statement_timeout`` (via
``set_config(..., is_local=True)``).
+ * MySQL: session ``max_execution_time`` (applies to read-only ``SELECT``
statements).
+
+ The previous value is captured and restored (rather than reset to ``0``)
so a server- or
+ role-level global timeout is preserved instead of being cleared. This is a
ReDoS safeguard: a
+ malicious pattern is aborted instead of pinning a database backend. No-op
on other backends
+ (e.g. SQLite) and when the configured timeout is ``0`` (which also means
regexp filtering is
+ disabled).
+ """
+ timeout_seconds = conf.getfloat("api", "regexp_query_timeout")
+ if timeout_seconds <= 0:
+ yield
+ return
+ # ``timeout_ms`` is derived from the (operator-controlled) config, not
from user input.
+ timeout_ms = int(timeout_seconds * 1000)
+ dialect_name = get_dialect_name(session)
+ if dialect_name == "postgresql":
+ previous = session.execute(text("SELECT
current_setting('statement_timeout')")).scalar()
+ session.execute(
+ text("SELECT set_config('statement_timeout', :timeout,
true)").bindparams(timeout=str(timeout_ms))
+ )
+ try:
+ yield
+ finally:
+ # Restore the previous value (e.g. a global statement_timeout)
instead of clearing it, so
+ # the bound only applies to the regexp query and not later
statements in the transaction.
+ session.execute(
+ text("SELECT set_config('statement_timeout', :timeout,
true)").bindparams(
+ timeout=str(previous) if previous is not None else "0"
+ )
+ )
+ elif dialect_name == "mysql":
+ previous = session.execute(text("SELECT
@@SESSION.max_execution_time")).scalar()
+ session.execute(text(f"SET SESSION max_execution_time = {timeout_ms}"))
+ try:
+ yield
+ finally:
+ # Restore the previous value so a global max_execution_time is
preserved.
+ session.execute(text(f"SET SESSION max_execution_time =
{int(previous or 0)}"))
+ else:
+ yield
+
+
def build_upsert_stmt(
dialect: str | None,
model: Any,
diff --git a/airflow-core/tests/unit/api_fastapi/common/test_parameters.py
b/airflow-core/tests/unit/api_fastapi/common/test_parameters.py
index 4d2f473f4f8..68df51a3f72 100644
--- a/airflow-core/tests/unit/api_fastapi/common/test_parameters.py
+++ b/airflow-core/tests/unit/api_fastapi/common/test_parameters.py
@@ -21,6 +21,7 @@ import re
from datetime import datetime, timezone
from types import SimpleNamespace
from typing import Annotated
+from unittest import mock
import pytest
from fastapi import Depends, FastAPI, HTTPException
@@ -41,11 +42,15 @@ from airflow.api_fastapi.common.parameters import (
_TaskDisplayNamePrefixPatternParam,
datetime_range_filter_factory,
filter_param_factory,
+ regex_param_factory,
)
from airflow.models import DagModel, DagRun, Log
+from airflow.models.asset import AssetEvent
from airflow.models.errors import ParseImportError
from airflow.models.taskinstance import TaskInstance
+from tests_common.test_utils.config import conf_vars
+
class TestFilterParam:
def test_filter_param_factory_description(self):
@@ -537,3 +542,51 @@ class TestDatetimeRangeFilterFactory:
rf = _make_datetime_filter("start_date", upper_bound_lte=bound)
sql = _compile(rf.to_orm(select(TaskInstance)))
assert "coalesce" not in sql
+
+
+class TestRegexParamFactory:
+ """The regexp filter dependency must apply the query timeout itself
(callers can't forget)."""
+
+ @staticmethod
+ def _depends():
+ return regex_param_factory(AssetEvent.partition_key,
"partition_key_regexp_pattern")
+
+ def test_dependency_applies_timeout_when_pattern_provided(self):
+ session = mock.MagicMock()
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ with
mock.patch("airflow.api_fastapi.common.parameters.apply_regex_query_timeout")
as apply_mock:
+ gen = self._depends()(session=session, value="^us")
+ param = next(gen)
+ apply_mock.assert_called_once_with(session)
+ assert param.value == "^us"
+ with pytest.raises(StopIteration):
+ next(gen)
+
+ def test_dependency_skips_timeout_without_pattern(self):
+ session = mock.MagicMock()
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ with
mock.patch("airflow.api_fastapi.common.parameters.apply_regex_query_timeout")
as apply_mock:
+ gen = self._depends()(session=session, value=None)
+ param = next(gen)
+ apply_mock.assert_not_called()
+ assert param.value is None
+ with pytest.raises(StopIteration):
+ next(gen)
+
+ def test_dependency_rejects_pattern_when_disabled(self):
+ session = mock.MagicMock()
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ gen = self._depends()(session=session, value="^us")
+ with pytest.raises(HTTPException) as exc:
+ next(gen)
+ assert exc.value.status_code == 400
+ assert "disabled" in exc.value.detail
+
+ def test_dependency_rejects_invalid_regex(self):
+ session = mock.MagicMock()
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ gen = self._depends()(session=session, value="[invalid(")
+ with pytest.raises(HTTPException) as exc:
+ next(gen)
+ assert exc.value.status_code == 400
+ assert "Invalid regular expression" in exc.value.detail
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
index d5c56d133ad..835a8ccb6dc 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
@@ -1087,6 +1087,162 @@ class TestGetAssetEvents(TestAssets):
}
+class TestGetAssetEventsPartitionKeyRegex(TestAssets):
+ """Tests for partition_key_regexp_pattern regex filter on GET
/assets/events.
+
+ Patterns are written to work consistently across PostgreSQL (~),
+ MySQL (REGEXP), and SQLite (re.match), including both anchored and
+ unanchored expressions where appropriate.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _enable_regexp_query_filters(self):
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ yield
+
+ @pytest.fixture(autouse=True)
+ def _create_partition_key_test_data(self, setup, session):
+ _create_assets(session=session)
+ events = [
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r1",
+ partition_key="2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r2",
+ partition_key="2024-01-02",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r3",
+ partition_key="us|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=2,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r4",
+ partition_key="eu|2024-01-01",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r5",
+ partition_key="apac|2024-03-20",
+ timestamp=DEFAULT_DATE,
+ ),
+ AssetEvent(
+ asset_id=1,
+ extra={},
+ source_task_id="t",
+ source_dag_id="d",
+ source_run_id="r6",
+ partition_key=None,
+ timestamp=DEFAULT_DATE,
+ ),
+ ]
+ session.add_all(events)
+ session.commit()
+
+ @pytest.mark.parametrize(
+ ("partition_key_regexp_pattern", "expected_count"),
+ [
+ ("^2024-01-01$", 1),
+ ("^2024-01-", 2),
+ ("^us\\|", 1),
+ (".*\\|2024-01-01$", 2),
+ ("^(us|eu)\\|", 2),
+ ("^nonexistent", 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_filtering(
+ self, test_client, partition_key_regexp_pattern, expected_count
+ ):
+ response = test_client.get(
+ "/assets/events", params={"partition_key_regexp_pattern":
partition_key_regexp_pattern}
+ )
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ @pytest.mark.parametrize(
+ ("params", "expected_count"),
+ [
+ ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "1"}, 1),
+ ({"partition_key_regexp_pattern": "^us\\|", "asset_id": "2"}, 0),
+ ({"partition_key_regexp_pattern": ".*\\|2024-01-01$",
"source_dag_id": "d"}, 2),
+ ({"partition_key_regexp_pattern": ".*\\|2024-01-01$",
"source_dag_id": "other"}, 0),
+ ],
+ )
+ def test_partition_key_regexp_pattern_combined_filters(self, test_client,
params, expected_count):
+ response = test_client.get("/assets/events", params=params)
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ def test_partition_key_regexp_pattern_invalid_regex_returns_400(self,
test_client):
+ response = test_client.get(
+ "/assets/events", params={"partition_key_regexp_pattern":
"[invalid(regex"}
+ )
+ assert response.status_code == 400
+ assert "Invalid regular expression" in response.json()["detail"]
+
+ def test_partition_key_regexp_pattern_disabled_returns_400(self,
test_client):
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = test_client.get("/assets/events",
params={"partition_key_regexp_pattern": "^2024-"})
+ assert response.status_code == 400
+ assert "disabled" in response.json()["detail"]
+
+ def test_exact_match_works_when_regex_disabled(self, test_client):
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = test_client.get("/assets/events",
params={"partition_key": "2024-01-01"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ def test_partition_key_exact_match_via_regex(self, test_client):
+ response = test_client.get("/assets/events",
params={"partition_key_regexp_pattern": "^2024-01-01$"})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 1
+
+ @pytest.mark.parametrize(
+ ("partition_key", "expected_count"),
+ [
+ ("2024-01-01", 1),
+ ("us|2024-01-01", 1),
+ ("nonexistent", 0),
+ ],
+ )
+ def test_partition_key_exact_match(self, test_client, partition_key,
expected_count):
+ response = test_client.get("/assets/events", params={"partition_key":
partition_key})
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == expected_count
+
+ def test_partition_key_and_pattern_combined(self, test_client):
+ # Both filters are allowed and combine with AND: a disjoint pair
yields no results.
+ response = test_client.get(
+ "/assets/events",
+ params={"partition_key": "2024-01-01",
"partition_key_regexp_pattern": "^2025-"},
+ )
+ assert response.status_code == 200
+ assert response.json()["total_entries"] == 0
+
+
class TestGetAssetEventsExtraFilter(TestAssets):
@pytest.fixture
def _setup(self, session):
diff --git
a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py
b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py
index 6300265472a..da1e3705daf 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 @@ import pytest
from airflow._shared.timezones import timezone
from airflow.models.asset import AssetActive, AssetAliasModel, AssetEvent,
AssetModel
+from tests_common.test_utils.config import conf_vars
+
DEFAULT_DATE = timezone.parse("2021-01-01T00:00:00")
pytestmark = pytest.mark.db_test
@@ -523,6 +525,355 @@ class TestGetAssetEventByAssetAlias:
}
+class TestGetAssetEventByAssetPartitionKey:
+ """Tests for partition_key_regexp_pattern regex filter on execution API.
+
+ Patterns are written to work consistently across PostgreSQL (~),
+ MySQL (REGEXP), and SQLite (re.match). They are not necessarily
+ ^-anchored; some cases use explicit prefixes such as ``.*`` to
+ achieve SQLite ``re.match``-compatible behavior.
+ """
+
+ @pytest.fixture(autouse=True)
+ def _enable_regexp_query_filters(self):
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ yield
+
+ @pytest.fixture
+ def test_partitioned_events(self, session, test_asset):
+ def make_timestamp(day):
+ return datetime(2021, 1, day, tzinfo=timezone.utc)
+
+ common = {
+ "asset_id": 1,
+ "extra": {"foo": "bar"},
+ "source_dag_id": "foo",
+ "source_task_id": "bar",
+ "source_run_id": "custom",
+ "source_map_index": -1,
+ }
+
+ events = [
+ AssetEvent(id=10, timestamp=make_timestamp(1),
partition_key="2024-01-01", **common),
+ AssetEvent(id=11, timestamp=make_timestamp(2),
partition_key="2024-01-02", **common),
+ AssetEvent(id=12, timestamp=make_timestamp(3),
partition_key="us|2024-01-01", **common),
+ AssetEvent(id=13, timestamp=make_timestamp(4),
partition_key="eu|2024-01-01", **common),
+ AssetEvent(id=14, timestamp=make_timestamp(5),
partition_key="apac|2024-03-20", **common),
+ AssetEvent(id=15, timestamp=make_timestamp(6), partition_key=None,
**common),
+ ]
+ session.add_all(events)
+ session.commit()
+ yield events
+
+ for event in events:
+ session.delete(event)
+ session.commit()
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_with_exact_partition_key(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": "^2024-01-01$",
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 1
+ assert data["asset_events"][0]["partition_key"] == "2024-01-01"
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_with_prefix_partition_key(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": "^2024-01-",
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 2
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_with_composite_partition_key_multiple_regions(self,
client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": r"^(us|eu)\|2024-01-.*",
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 2
+ keys = {e["partition_key"] for e in data["asset_events"]}
+ assert keys == {"us|2024-01-01", "eu|2024-01-01"}
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def
test_get_by_asset_with_composite_partition_key_date_across_regions(self,
client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": r".*\|2024-01-01$",
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 2
+ keys = {e["partition_key"] for e in data["asset_events"]}
+ assert keys == {"us|2024-01-01", "eu|2024-01-01"}
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_with_partition_key_no_match(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": "^2025-",
+ },
+ )
+ assert response.status_code == 200
+ assert len(response.json()["asset_events"]) == 0
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_without_partition_key_returns_all(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={"name": "test_get_asset_by_name", "uri": None},
+ )
+ assert response.status_code == 200
+ assert len(response.json()["asset_events"]) == 6
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_with_partition_key_and_time_range(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": "^2024-01-",
+ "after": "2021-01-01T12:00:00Z",
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 1
+ assert data["asset_events"][0]["partition_key"] == "2024-01-02"
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_with_invalid_regex_returns_400(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": "[invalid(regex",
+ },
+ )
+ assert response.status_code == 400
+ assert "Invalid regular expression" in response.json()["detail"]
+
+ def test_get_by_asset_with_pattern_disabled_returns_400(self, client):
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key_regexp_pattern": "^2021-",
+ },
+ )
+ assert response.status_code == 400
+ assert "disabled" in response.json()["detail"]
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ @pytest.mark.parametrize(
+ ("partition_key", "expected_keys"),
+ [
+ ("2024-01-01", ["2024-01-01"]),
+ ("us|2024-01-01", ["us|2024-01-01"]),
+ ("nonexistent", []),
+ ],
+ )
+ def test_get_by_asset_exact_partition_key(self, client, partition_key,
expected_keys):
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key": partition_key,
+ },
+ )
+ assert response.status_code == 200
+ keys = [event["partition_key"] for event in
response.json()["asset_events"]]
+ assert keys == expected_keys
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_events")
+ def test_get_by_asset_partition_key_and_pattern_combined(self, client):
+ # Both filters are allowed and combine with AND: a disjoint pair
yields no results.
+ response = client.get(
+ "/execution/asset-events/by-asset",
+ params={
+ "name": "test_get_asset_by_name",
+ "uri": None,
+ "partition_key": "2024-01-01",
+ "partition_key_regexp_pattern": "^2025-",
+ },
+ )
+ assert response.status_code == 200
+ assert len(response.json()["asset_events"]) == 0
+
+
+class TestGetAssetEventByAssetAliasPartitionKey:
+ """Tests for partition_key_regexp_pattern regex filter on by-asset-alias
endpoint.
+
+ All patterns use ^-anchored regex so they work consistently across
+ PostgreSQL (~), MySQL (REGEXP), and SQLite (re.match).
+ """
+
+ @pytest.fixture(autouse=True)
+ def _enable_regexp_query_filters(self):
+ with conf_vars({("api", "regexp_query_timeout"): "30"}):
+ yield
+
+ @pytest.fixture
+ def test_partitioned_alias_events(self, session, test_asset):
+ def make_timestamp(day):
+ return datetime(2021, 1, day, tzinfo=timezone.utc)
+
+ common = {
+ "asset_id": 1,
+ "extra": {"foo": "bar"},
+ "source_dag_id": "foo",
+ "source_task_id": "bar",
+ "source_run_id": "custom",
+ "source_map_index": -1,
+ }
+
+ events = [
+ AssetEvent(id=20, timestamp=make_timestamp(1),
partition_key="us|2024-01-01", **common),
+ AssetEvent(id=21, timestamp=make_timestamp(2),
partition_key="eu|2024-01-01", **common),
+ AssetEvent(id=22, timestamp=make_timestamp(3), partition_key=None,
**common),
+ ]
+ session.add_all(events)
+ session.commit()
+
+ alias = AssetAliasModel(id=10, name="partitioned_alias")
+ alias.asset_events = events
+ alias.assets.append(test_asset)
+ session.add(alias)
+ session.commit()
+
+ yield events, alias
+
+ session.delete(alias)
+ for event in events:
+ session.delete(event)
+ session.commit()
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_with_partition_key_region(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias",
"partition_key_regexp_pattern": r"^us\|.*"},
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 1
+ assert data["asset_events"][0]["partition_key"] == "us|2024-01-01"
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_with_partition_key_all_regions(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias",
"partition_key_regexp_pattern": r".*\|2024-01-01$"},
+ )
+ assert response.status_code == 200
+ assert len(response.json()["asset_events"]) == 2
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_without_partition_key_returns_all(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias"},
+ )
+ assert response.status_code == 200
+ assert len(response.json()["asset_events"]) == 3
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_with_invalid_regex_returns_400(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias",
"partition_key_regexp_pattern": "[invalid(regex"},
+ )
+ assert response.status_code == 400
+ assert "Invalid regular expression" in response.json()["detail"]
+
+ def test_get_by_alias_with_pattern_disabled_returns_400(self, client):
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias",
"partition_key_regexp_pattern": "^us"},
+ )
+ assert response.status_code == 400
+ assert "disabled" in response.json()["detail"]
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_with_exact_partition_key_regex(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias",
"partition_key_regexp_pattern": r"^us\|2024-01-01$"},
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 1
+ assert data["asset_events"][0]["partition_key"] == "us|2024-01-01"
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_exact_partition_key(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias", "partition_key":
"us|2024-01-01"},
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 1
+ assert data["asset_events"][0]["partition_key"] == "us|2024-01-01"
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_exact_partition_key_no_match(self, client):
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={"name": "partitioned_alias", "partition_key":
"nonexistent"},
+ )
+ assert response.status_code == 200
+ assert len(response.json()["asset_events"]) == 0
+
+ @pytest.mark.usefixtures("test_asset", "test_partitioned_alias_events")
+ def test_get_by_alias_partition_key_and_pattern_combined(self, client):
+ # Both filters are allowed and combine with AND.
+ response = client.get(
+ "/execution/asset-events/by-asset-alias",
+ params={
+ "name": "partitioned_alias",
+ "partition_key": "us|2024-01-01",
+ "partition_key_regexp_pattern": r"^us\|",
+ },
+ )
+ assert response.status_code == 200
+ data = response.json()
+ assert len(data["asset_events"]) == 1
+ assert data["asset_events"][0]["partition_key"] == "us|2024-01-01"
+
+
class TestGetAssetEventByAssetExtraFilter:
@pytest.fixture
def test_events_with_extra(self, session):
diff --git a/airflow-core/tests/unit/utils/test_sqlalchemy.py
b/airflow-core/tests/unit/utils/test_sqlalchemy.py
index 1ef7b42c587..8dbcd413130 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.serialization.serialized_objects import
BaseSerialization
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.sqlalchemy import (
from airflow.utils.state import State
from airflow.utils.types import DagRunTriggeredByType, DagRunType
+from tests_common.test_utils.config import conf_vars
from tests_common.test_utils.dag import sync_dag_to_db
pytestmark = pytest.mark.db_test
@@ -380,3 +382,64 @@ class TestExecutorConfigType:
pickle.dumps(fixed_pod)
assert fixed_pod.local_vars_configuration.refresh_api_key_hook is None
assert
fixed_pod.spec.containers[0].local_vars_configuration.refresh_api_key_hook is
None
+
+
+class TestApplyRegexQueryTimeout:
+ @staticmethod
+ def _mock_session(dialect_name):
+ session = mock.MagicMock()
+ session.get_bind.return_value.dialect.name = dialect_name
+ return session
+
+ def test_sets_and_restores_statement_timeout_on_postgresql(self):
+ session = self._mock_session("postgresql")
+ # Pre-existing (e.g. global) statement_timeout captured before we
override it.
+ session.execute.return_value.scalar.return_value = "10s"
+ with conf_vars({("api", "regexp_query_timeout"): "5"}):
+ with apply_regex_query_timeout(session):
+ # Capture-then-set on enter.
+ assert session.execute.call_count == 2
+ set_stmt = session.execute.call_args_list[1].args[0]
+ # 5 seconds -> 5000 ms, passed as a bound parameter (no SQL
injection).
+ assert set_stmt.compile().params == {"timeout": "5000"}
+ # Restored on exit to the previous value (not reset to 0), so a global
timeout is preserved.
+ assert session.execute.call_count == 3
+ restore_stmt = session.execute.call_args_list[2].args[0]
+ assert restore_stmt.compile().params == {"timeout": "10s"}
+
+ def test_sets_and_restores_max_execution_time_on_mysql(self):
+ session = self._mock_session("mysql")
+ # Pre-existing (e.g. global) max_execution_time captured before we
override it.
+ session.execute.return_value.scalar.return_value = 1000
+ with conf_vars({("api", "regexp_query_timeout"): "5"}):
+ with apply_regex_query_timeout(session):
+ # Capture-then-set on enter (5 seconds -> 5000 ms).
+ assert session.execute.call_count == 2
+ assert "max_execution_time = 5000" in
str(session.execute.call_args.args[0])
+ # Restored on exit to the previous value, so a global limit is
preserved.
+ assert session.execute.call_count == 3
+ assert "max_execution_time = 1000" in
str(session.execute.call_args.args[0])
+
+ def test_fractional_seconds_are_converted_to_milliseconds(self):
+ session = self._mock_session("postgresql")
+ session.execute.return_value.scalar.return_value = "0"
+ with conf_vars({("api", "regexp_query_timeout"): "0.5"}):
+ with apply_regex_query_timeout(session):
+ # Set is the second call (after capturing the previous value).
+ set_stmt = session.execute.call_args_list[1].args[0]
+ # 0.5 seconds -> 500 ms.
+ assert set_stmt.compile().params == {"timeout": "500"}
+
+ def test_noop_on_sqlite(self):
+ session = self._mock_session("sqlite")
+ with conf_vars({("api", "regexp_query_timeout"): "5"}):
+ with apply_regex_query_timeout(session):
+ pass
+ session.execute.assert_not_called()
+
+ def test_noop_when_timeout_disabled(self):
+ session = self._mock_session("postgresql")
+ with conf_vars({("api", "regexp_query_timeout"): "0"}):
+ with apply_regex_query_timeout(session):
+ pass
+ session.execute.assert_not_called()
diff --git a/task-sdk/src/airflow/sdk/api/client.py
b/task-sdk/src/airflow/sdk/api/client.py
index 74a68005371..7e681c11149 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -858,6 +858,8 @@ class AssetEventOperations:
before: datetime | None = None,
ascending: bool = True,
limit: int | None = None,
+ partition_key: str | None = None,
+ partition_key_regexp_pattern: str | None = None,
extra: dict[str, str] | None = None,
) -> AssetEventsResponse:
"""Get Asset event from the API server."""
@@ -869,6 +871,10 @@ class AssetEventOperations:
common_params["ascending"] = ascending
if limit is not None:
common_params["limit"] = limit
+ if partition_key is not None:
+ common_params["partition_key"] = partition_key
+ if partition_key_regexp_pattern is not None:
+ common_params["partition_key_regexp_pattern"] =
partition_key_regexp_pattern
extra_params: list[tuple[str, str]] = []
if extra:
extra_params = [("extra", f"{k}={v}") for k, v in extra.items()]
diff --git a/task-sdk/src/airflow/sdk/execution_time/comms.py
b/task-sdk/src/airflow/sdk/execution_time/comms.py
index 6e9e49b2676..5c4ae14dff1 100644
--- a/task-sdk/src/airflow/sdk/execution_time/comms.py
+++ b/task-sdk/src/airflow/sdk/execution_time/comms.py
@@ -1142,6 +1142,8 @@ class GetAssetEventByAsset(BaseModel):
before: AwareDatetime | None = None
limit: int | None = None
ascending: bool = True
+ partition_key: str | None = None
+ partition_key_regexp_pattern: str | None = None
extra: dict[str, str] | None = None
type: Literal["GetAssetEventByAsset"] = "GetAssetEventByAsset"
@@ -1152,6 +1154,8 @@ class GetAssetEventByAssetAlias(BaseModel):
before: AwareDatetime | None = None
limit: int | None = None
ascending: bool = True
+ partition_key: str | None = None
+ partition_key_regexp_pattern: str | None = None
extra: dict[str, str] | None = None
type: Literal["GetAssetEventByAssetAlias"] = "GetAssetEventByAssetAlias"
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index 2828525a7bb..2c2364a49b2 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -1114,6 +1114,8 @@ class InletEventsAccessor(Sequence["AssetEventResult"]):
_before: str | datetime | None
_ascending: bool
_limit: int | None
+ _partition_key: str | None
+ _partition_key_regexp_pattern: str | None
_extra: dict[str, str]
_asset_name: str | None
_asset_uri: str | None
@@ -1129,6 +1131,8 @@ class InletEventsAccessor(Sequence["AssetEventResult"]):
self._before = None
self._ascending = True
self._limit = None
+ self._partition_key = None
+ self._partition_key_regexp_pattern = None
self._extra: dict[str, str] = {}
def after(self, after: str) -> Self:
@@ -1151,6 +1155,18 @@ class InletEventsAccessor(Sequence["AssetEventResult"]):
self._reset_cache()
return self
+ def partition_key(self, key: str) -> Self:
+ """Filter by exact partition key match."""
+ self._partition_key = key
+ self._reset_cache()
+ return self
+
+ def partition_key_regexp_pattern(self, pattern: str) -> Self:
+ """Filter by partition key regexp pattern."""
+ self._partition_key_regexp_pattern = pattern
+ self._reset_cache()
+ return self
+
def extra(self, key: str, value: str) -> Self:
self._extra[key] = value
self._reset_cache()
@@ -1165,6 +1181,8 @@ class InletEventsAccessor(Sequence["AssetEventResult"]):
"before": self._before,
"ascending": self._ascending,
"limit": self._limit,
+ "partition_key": self._partition_key,
+ "partition_key_regexp_pattern": self._partition_key_regexp_pattern,
"extra": self._extra or None,
}
diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
index 01e3c8970b5..0ec8fe4e49a 100644
--- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
+++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json
@@ -1872,6 +1872,30 @@
"title": "Ascending",
"type": "boolean"
},
+ "partition_key": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Partition Key"
+ },
+ "partition_key_regexp_pattern": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Partition Key Regexp Pattern"
+ },
"extra": {
"anyOf": [
{
@@ -1950,6 +1974,30 @@
"title": "Ascending",
"type": "boolean"
},
+ "partition_key": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Partition Key"
+ },
+ "partition_key_regexp_pattern": {
+ "anyOf": [
+ {
+ "type": "string"
+ },
+ {
+ "type": "null"
+ }
+ ],
+ "default": null,
+ "title": "Partition Key Regexp Pattern"
+ },
"extra": {
"anyOf": [
{
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index 447f7c7594d..87311f02da7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1774,6 +1774,8 @@ class ActivitySubprocess(WatchedSubprocess):
before=msg.before,
ascending=msg.ascending,
limit=msg.limit,
+ partition_key=msg.partition_key,
+ partition_key_regexp_pattern=msg.partition_key_regexp_pattern,
extra=msg.extra,
)
asset_event_result =
AssetEventsResult.from_asset_events_response(asset_event_resp)
@@ -1786,6 +1788,8 @@ class ActivitySubprocess(WatchedSubprocess):
before=msg.before,
ascending=msg.ascending,
limit=msg.limit,
+ partition_key=msg.partition_key,
+ partition_key_regexp_pattern=msg.partition_key_regexp_pattern,
extra=msg.extra,
)
asset_event_result =
AssetEventsResult.from_asset_events_response(asset_event_resp)
diff --git a/task-sdk/tests/task_sdk/api/test_client.py
b/task-sdk/tests/task_sdk/api/test_client.py
index 205b12bd4f2..0f9b8130e26 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 @@ class TestAssetEventOperations:
assert result.asset_events[0].asset.name == "this_asset"
assert result.asset_events[0].asset.uri == "s3://bucket/key"
+ def test_partition_key_exact_match_param_passed(self):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ params = request.url.params
+ assert params.get("partition_key") == "2024-01-01"
+ assert "partition_key_regexp_pattern" not in params
+ return httpx.Response(
+ status_code=200,
+ json={
+ "asset_events": [
+ {
+ "id": 1,
+ "asset": {
+ "name": "this_asset",
+ "uri": "s3://bucket/key",
+ "group": "asset",
+ },
+ "created_dagruns": [],
+ "timestamp": "2023-01-01T00:00:00Z",
+ "partition_key": "2024-01-01",
+ }
+ ]
+ },
+ )
+
+ client = make_client(httpx.MockTransport(handle_request))
+ result = client.asset_events.get(name="this_asset",
partition_key="2024-01-01")
+
+ assert isinstance(result, AssetEventsResponse)
+ assert len(result.asset_events) == 1
+
+ def test_partition_key_regexp_pattern_param_passed(self):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ params = request.url.params
+ assert params.get("partition_key_regexp_pattern") == "^2024-01-"
+ assert "partition_key" not in params
+ return httpx.Response(
+ status_code=200,
+ json={"asset_events": []},
+ )
+
+ client = make_client(httpx.MockTransport(handle_request))
+ result = client.asset_events.get(name="this_asset",
partition_key_regexp_pattern="^2024-01-")
+
+ assert isinstance(result, AssetEventsResponse)
+ assert len(result.asset_events) == 0
+
+ def test_partition_key_regexp_pattern_with_other_params(self):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ params = request.url.params
+ assert params.get("partition_key_regexp_pattern") ==
r"^us\|2024-.*"
+ assert "partition_key" not in params
+ assert params.get("after") == "2023-06-01T00:00:00+00:00"
+ assert params.get("limit") == "5"
+ assert params.get("ascending") == "false"
+ return httpx.Response(
+ status_code=200,
+ json={"asset_events": []},
+ )
+
+ client = make_client(httpx.MockTransport(handle_request))
+ result = client.asset_events.get(
+ name="this_asset",
+ partition_key_regexp_pattern=r"^us\|2024-.*",
+ after=datetime(2023, 6, 1, tzinfo=dt_timezone.utc),
+ limit=5,
+ ascending=False,
+ )
+
+ assert isinstance(result, AssetEventsResponse)
+ assert len(result.asset_events) == 0
+
+ def test_partition_key_with_alias(self):
+ def handle_request(request: httpx.Request) -> httpx.Response:
+ assert request.url.path == "/asset-events/by-asset-alias"
+ params = request.url.params
+ assert params.get("name") == "my_alias"
+ assert params.get("partition_key") == "us-east|2024-01-01"
+ assert "partition_key_regexp_pattern" not in params
+ return httpx.Response(
+ status_code=200,
+ json={"asset_events": []},
+ )
+
+ client = make_client(httpx.MockTransport(handle_request))
+ result = client.asset_events.get(alias_name="my_alias",
partition_key="us-east|2024-01-01")
+
+ assert isinstance(result, AssetEventsResponse)
+ assert len(result.asset_events) == 0
+
def test_extra_dict_param_passed(self):
def handle_request(request: httpx.Request) -> httpx.Response:
params = request.url.params
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 870b2e6deed..c39ceb92b0e 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -1940,6 +1940,8 @@ REQUEST_TEST_CASES = [
"before": None,
"limit": None,
"ascending": True,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -1984,6 +1986,8 @@ REQUEST_TEST_CASES = [
"before": timezone.parse("2024-10-15T12:00:00Z"),
"limit": 5,
"ascending": False,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2021,6 +2025,8 @@ REQUEST_TEST_CASES = [
"before": None,
"limit": None,
"ascending": True,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2065,6 +2071,8 @@ REQUEST_TEST_CASES = [
"before": timezone.parse("2024-10-15T12:00:00Z"),
"limit": 5,
"ascending": False,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2102,6 +2110,8 @@ REQUEST_TEST_CASES = [
"before": None,
"limit": None,
"ascending": True,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2146,6 +2156,8 @@ REQUEST_TEST_CASES = [
"before": timezone.parse("2024-10-15T12:00:00Z"),
"limit": 5,
"ascending": False,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2182,6 +2194,8 @@ REQUEST_TEST_CASES = [
"before": None,
"limit": None,
"ascending": True,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2224,6 +2238,8 @@ REQUEST_TEST_CASES = [
"before": timezone.parse("2024-10-15T12:00:00Z"),
"limit": 5,
"ascending": False,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": None,
},
response=AssetEventsResult(
@@ -2239,6 +2255,49 @@ REQUEST_TEST_CASES = [
),
test_id="get_asset_events_by_asset_alias_with_filters",
),
+ RequestTestCase(
+ message=GetAssetEventByAsset(
+ uri="s3://bucket/obj",
+ name="test",
+ partition_key="us|2024-01-15",
+ ),
+ expected_body={
+ "asset_events": [
+ {
+ "id": 1,
+ "timestamp": timezone.parse("2024-10-31T12:00:00Z"),
+ "asset": {"name": "asset", "uri": "s3://bucket/obj",
"group": "asset"},
+ "created_dagruns": [],
+ }
+ ],
+ "type": "AssetEventsResult",
+ },
+ client_mock=ClientMock(
+ method_path="asset_events.get",
+ kwargs={
+ "uri": "s3://bucket/obj",
+ "name": "test",
+ "after": None,
+ "before": None,
+ "limit": None,
+ "ascending": True,
+ "partition_key": "us|2024-01-15",
+ "partition_key_regexp_pattern": None,
+ "extra": None,
+ },
+ response=AssetEventsResult(
+ asset_events=[
+ AssetEventResponse(
+ id=1,
+ asset=AssetResponse(name="asset",
uri="s3://bucket/obj", group="asset"),
+ created_dagruns=[],
+ timestamp=timezone.parse("2024-10-31T12:00:00Z"),
+ )
+ ]
+ ),
+ ),
+ test_id="get_asset_events_by_name_with_partition_key",
+ ),
RequestTestCase(
message=GetAssetEventByAsset(
uri="s3://bucket/obj",
@@ -2265,6 +2324,8 @@ REQUEST_TEST_CASES = [
"before": None,
"limit": None,
"ascending": True,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": {"region": "us"},
},
response=AssetEventsResult(
@@ -2280,6 +2341,47 @@ REQUEST_TEST_CASES = [
),
test_id="get_asset_events_with_extra_filter",
),
+ RequestTestCase(
+ message=GetAssetEventByAssetAlias(
+ alias_name="test_alias",
+ partition_key="eu|2024-03",
+ ),
+ expected_body={
+ "asset_events": [
+ {
+ "id": 1,
+ "timestamp": timezone.parse("2024-10-31T12:00:00Z"),
+ "asset": {"name": "asset", "uri": "s3://bucket/obj",
"group": "asset"},
+ "created_dagruns": [],
+ }
+ ],
+ "type": "AssetEventsResult",
+ },
+ client_mock=ClientMock(
+ method_path="asset_events.get",
+ kwargs={
+ "alias_name": "test_alias",
+ "after": None,
+ "before": None,
+ "limit": None,
+ "ascending": True,
+ "partition_key": "eu|2024-03",
+ "partition_key_regexp_pattern": None,
+ "extra": None,
+ },
+ response=AssetEventsResult(
+ asset_events=[
+ AssetEventResponse(
+ id=1,
+ asset=AssetResponse(name="asset",
uri="s3://bucket/obj", group="asset"),
+ created_dagruns=[],
+ timestamp=timezone.parse("2024-10-31T12:00:00Z"),
+ )
+ ]
+ ),
+ ),
+ test_id="get_asset_events_by_alias_with_partition_key",
+ ),
RequestTestCase(
message=GetAssetEventByAssetAlias(
alias_name="test_alias",
@@ -2304,6 +2406,8 @@ REQUEST_TEST_CASES = [
"before": None,
"limit": None,
"ascending": True,
+ "partition_key": None,
+ "partition_key_regexp_pattern": None,
"extra": {"env": "prod"},
},
response=AssetEventsResult(
diff --git a/ts-sdk/src/generated/supervisor.ts
b/ts-sdk/src/generated/supervisor.ts
index 4766c0e8a79..ab2632831ab 100644
--- a/ts-sdk/src/generated/supervisor.ts
+++ b/ts-sdk/src/generated/supervisor.ts
@@ -372,6 +372,8 @@ export type After = string | null;
export type Before = string | null;
export type Limit = number | null;
export type Ascending = boolean;
+export type PartitionKey5 = string | null;
+export type PartitionKeyRegexpPattern = string | null;
export type Extra7 = {
[k: string]: string;
} | null;
@@ -381,6 +383,8 @@ export type After1 = string | null;
export type Before1 = string | null;
export type Limit1 = number | null;
export type Ascending1 = boolean;
+export type PartitionKey6 = string | null;
+export type PartitionKeyRegexpPattern1 = string | null;
export type Extra8 = {
[k: string]: string;
} | null;
@@ -589,7 +593,7 @@ export type Conf2 = {
[k: string]: unknown;
} | null;
export type ResetDagRun = boolean | null;
-export type PartitionKey5 = string | null;
+export type PartitionKey7 = string | null;
export type Note2 = string | null;
export type DagId22 = string;
export type DagRunId = string;
@@ -1223,6 +1227,8 @@ export interface GetAssetEventByAsset {
before?: Before;
limit?: Limit;
ascending?: Ascending;
+ partition_key?: PartitionKey5;
+ partition_key_regexp_pattern?: PartitionKeyRegexpPattern;
extra?: Extra7;
type?: Type29;
}
@@ -1236,6 +1242,8 @@ export interface GetAssetEventByAssetAlias {
before?: Before1;
limit?: Limit1;
ascending?: Ascending1;
+ partition_key?: PartitionKey6;
+ partition_key_regexp_pattern?: PartitionKeyRegexpPattern1;
extra?: Extra8;
type?: Type30;
}
@@ -1803,7 +1811,7 @@ export interface TriggerDagRun {
run_after?: RunAfter2;
conf?: Conf2;
reset_dag_run?: ResetDagRun;
- partition_key?: PartitionKey5;
+ partition_key?: PartitionKey7;
note?: Note2;
dag_id: DagId22;
run_id: DagRunId;
