Add partition_key and partition_key_pattern filters for asset events

[hussein-awala]

Summary

Add two query parameters for filtering asset events by partition key across the full stack:

• partition_key — exact-match filter using == equality, leverages the B-tree
  composite index (asset_id, partition_key) for fast lookups.
• partition_key_pattern — regex filter using SQLAlchemy's regexp_match, which maps
  to database-native regex on all backends (PostgreSQL ~, MySQL REGEXP, SQLite re.match).

The two parameters are mutually exclusive — providing both returns a 400 error.

Also includes:

• Composite database index (asset_id, partition_key) on asset_event table with Alembic migration.
• Documentation in assets.rst with usage examples for both parameters.
• InletEventsAccessor exposes .partition_key(value) for exact match and
  .partition_key_pattern(pattern) for regex within tasks.

Details

Exact match (partition_key)

Uses a simple WHERE partition_key = ? equality check. This is the preferred option
when you know the full partition key — it leverages the B-tree index directly with
zero regex overhead.

Regex filter (partition_key_pattern)

Uses regexp_match which maps to database-native regex on all supported backends:
PostgreSQL (~ operator), MySQL (REGEXP), and SQLite (via Python re.match).
Note that SQLite's re.match is anchored at the start of the string, so patterns
should use ^-anchored regex for consistent cross-backend behavior.

Invalid regex patterns are validated with re.compile() before hitting the database,
returning a 400 error with a descriptive message.

A _RegexParam class and regex_param_factory were added to the common FastAPI
parameters module for reuse (the factory accepts a customizable description).

Performance note

Currently, partition_key_pattern performs a full table scan (filtered by any other
WHERE clauses) since standard B-tree indexes cannot accelerate regex matching. For
exact matches, the composite B-tree index (asset_id, partition_key) is used directly.

For high-volume deployments where regex filtering on partition_key becomes a bottleneck,
a trigram (n-gram) index can significantly improve performance on PostgreSQL. This
requires the pg_trgm extension and can be created manually by a database administrator
without any application-level changes:

CREATE EXTENSION IF NOT EXISTS pg_trgm;

CREATE INDEX CONCURRENTLY idx_asset_event_partition_key_trgm
ON asset_event USING gin (partition_key gin_trgm_ops);

PostgreSQL's query planner will automatically use the trigram index for ~ (regex)
queries emitted by regexp_match — no Airflow configuration changes needed.

When it helps:

• Non-prefix patterns where the B-tree index cannot assist: .*sales$, .*\|2024-.*, us.*west
• Substring or suffix matching across large tables

When it doesn't help (and what to use instead):

• Simple prefix patterns like ^us\|2024- — for these, use partition_key (exact match)
  which leverages the B-tree index directly with zero overhead
• Tables with very high write throughput — trigram GIN indexes are larger and slower to
  update than B-tree, adding write-time cost on every INSERT into asset_event

Shipping this as an Alembic migration would require CREATE EXTENSION privileges
(typically superuser) and a cross-backend strategy, since MySQL and SQLite have no
equivalent. This can be discussed separately if there is demand.

Files changed

Core API layer:

• parameters.py — new _RegexParam, regex_param_factory, QueryAssetEventPartitionKeyFilter
  (exact match), QueryAssetEventPartitionKeyRegex (regex)
• assets.py — accept both partition_key and partition_key_pattern filters in
  GET /api/v2/assets/events with mutual-exclusion validation

Execution API layer:

• asset_events.py — accept both params in /by-asset and /by-asset-alias endpoints,
  with mutual-exclusion and regex validation

Task SDK:

• client.py — send partition_key and partition_key_pattern as separate query params,
  with client-side ValueError if both are provided
• comms.py — add both fields to request models
• supervisor.py — forward both fields to API client
• context.py — InletEventsAccessor exposes .partition_key() and .partition_key_pattern()
  (mutually exclusive — setting one clears the other)

Database:

• asset.py — composite index idx_asset_event_asset_id_partition_key
• New migration 0112_3_3_0_add_idx_asset_event_partition_key.py

Auto-generated:

• OpenAPI specs (core API, private UI) — new query parameters
• UI TypeScript clients (openapi-gen/queries/*, openapi-gen/requests/*)

Tests:

• Core API and Execution API tests for exact match, regex, mutual-exclusion 400 error
• Task SDK client tests for param forwarding and mutual exclusivity
• Supervisor tests for forwarding both params

Docs:

• assets.rst — usage examples for both parameters via REST API and InletEventsAccessor

Test plan

• Unit tests for Core REST API partition key exact match and regex filtering
• Unit tests for Execution API partition key filtering (both by-asset and by-alias)
• Mutual-exclusion validation returns HTTP 400 when both params are provided
• Invalid regex patterns return HTTP 400 with descriptive message
• Task SDK client raises ValueError when both params are provided
• InletEventsAccessor clears the other field when setting one
• All tests run on all backends (PostgreSQL, MySQL, SQLite)

[Copilot]

Pull request overview

Adds a regex-based partition key filter for asset events end-to-end (public REST API, execution API, Task SDK/supervisor), plus a supporting DB index/migration and documentation updates.

Changes:

• Introduces partition_key_pattern (REST) / partition_key (execution API + SDK) regex filtering and wires it through API ↔ supervisor ↔ Task SDK accessor.
• Adds (asset_id, partition_key) composite index on asset_event with Alembic migration + updates migration head mapping/docs.
• Regenerates OpenAPI artifacts (YAML + UI TS clients) and adds tests/docs for the new filter.

Reviewed changes

Copilot reviewed 29 out of 29 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
task-sdk/tests/task_sdk/execution_time/test_supervisor.py	Supervisor request forwarding test updates + new partition_key cases
task-sdk/tests/task_sdk/api/test_client.py	SDK client query param forwarding tests for partition_key
task-sdk/src/airflow/sdk/execution_time/supervisor.py	Forwards partition_key from supervisor messages to API client
task-sdk/src/airflow/sdk/execution_time/context.py	Adds .partition_key(pattern) chaining to InletEventsAccessor
task-sdk/src/airflow/sdk/execution_time/comms.py	Adds partition_key to supervisor request message models
task-sdk/src/airflow/sdk/api/client.py	Adds partition_key query param forwarding in asset_events.get()
airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_asset_events.py	Adds execution API regex filter tests (postgres-only)
airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py	Adds REST API partition_key_pattern filtering tests (postgres-only)
airflow-core/src/airflow/utils/db.py	Updates migration head mapping for 3.3.0
airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts	Regenerates UI request types incl. partitionKeyPattern
airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts	Regenerates request service to pass partition_key_pattern
airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts	Regenerates request schemas (incl. ValidationError schema changes)
airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts	Regenerates suspense query hook with partitionKeyPattern
airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts	Regenerates query hook with partitionKeyPattern
airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts	Regenerates prefetch helper with partitionKeyPattern
airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts	Regenerates ensureQueryData helper with partitionKeyPattern
airflow-core/src/airflow/ui/openapi-gen/queries/common.ts	Regenerates query key fn to include partitionKeyPattern
airflow-core/src/airflow/models/asset.py	Adds ORM index idx_asset_event_asset_id_partition_key
airflow-core/src/airflow/migrations/versions/0111_3_3_0_add_idx_asset_event_partition_key.py	New migration to create/drop the composite index
airflow-core/src/airflow/api_fastapi/execution_api/routes/asset_events.py	Adds execution API regex filter param + query clause
airflow-core/src/airflow/api_fastapi/core_api/routes/public/assets.py	Adds REST partition_key_pattern filter dependency to /assets/events
airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml	Regenerates REST OpenAPI (adds param + changes ValidationError schema)
airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml	Regenerates private UI OpenAPI (ValidationError schema changes)
airflow-core/src/airflow/api_fastapi/common/parameters.py	Adds reusable _RegexParam + regex_param_factory + query alias type
airflow-core/src/airflow/api_fastapi/auth/managers/simple/ui/openapi-gen/requests/types.gen.ts	Regenerates simple auth manager UI request types
airflow-core/src/airflow/api_fastapi/auth/managers/simple/ui/openapi-gen/requests/schemas.gen.ts	Regenerates simple auth manager UI request schemas
airflow-core/src/airflow/api_fastapi/auth/managers/simple/openapi/v2-simple-auth-manager-generated.yaml	Regenerates simple auth manager OpenAPI (ValidationError schema changes)
airflow-core/docs/migrations-ref.rst	Updates auto-generated migration reference table (new head)
airflow-core/docs/authoring-and-scheduling/assets.rst	Documents Task SDK accessor + REST query usage for partition key regex

Comments suppressed due to low confidence (1)

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

• If partition_key_pattern contains an invalid regex, the underlying DB will raise during query execution (e.g. PostgreSQL error), which currently propagates as a 500. Consider catching the relevant SQLAlchemy/DBAPI exception around the query execution and returning a 400 with a clear message that the regex pattern is invalid for the current backend.

    assets_event_select, total_entries = paginated_select(
        statement=base_statement,
        filters=[
            asset_id,
            source_dag_id,
            source_task_id,
            source_run_id,
            source_map_index,
            partition_key_pattern,
            name_pattern,
            timestamp_range,
        ],
        order_by=order_by,
        offset=offset,
        limit=limit,
        session=session,
    )

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 5 comments.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 5 comments.

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 7 comments.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 3 comments.

[pierrejeambrun]

Is regexp required here? We have a lot of pattern and prefix_pattern filters doing both substring match and prefix match. (supporting logical OR, some special character too "~" etc....)

That introduce an inconsistency for that specific pattern param as well as introducing complexity handling regexp.

Would it be possible to implement that search similarly to what we already have on other APIs and discuss on a separate issue/PR specifically the introduction of Regexp matching.

[hussein-awala]

Sorry for the late reply @pierrejeambrun, and thanks for the review.

You raise a fair point about consistency, and I agree the regex introduction deserves a dedicated discussion. Let me explain the concrete gap so we can decide together.

The existing pattern (substring ILIKE '%term%') and prefix_pattern (prefix range scan) filters are great for free-text columns like dag_id, but they fall short for partition keys, which are typically structured composite strings (e.g. us|2024-01-15, region=us/date=2024-01-15). A few queries that are common for partition keys and that pattern/prefix_pattern cannot express:

1. The | delimiter collides with the OR operator

Both filters treat | as logical OR (val_str.split("|")). But partition keys very frequently use | as a field separator (see this PR's own examples like us|2026-03-10). So:

• prefix_pattern=us|2024-01-15 is parsed as OR(prefix "us", prefix "2024-01-15") — not a match on the literal key.
• With regex you escape it: partition_key_pattern=^us\|2024-01-15$.

2. Suffix / "ends with" matching

Find all partitions for a given date regardless of the region prefix:

• regex: partition_key_pattern=\|2024-01-15$
• pattern can only do unanchored substring (%2024-01-15%), which also matches 2024-01-15T09:00; prefix_pattern only matches from the start. Neither can anchor to the end.

3. Character-class / format constraints

Match only well-formed date partitions for a region:

• regex: partition_key_pattern=^us\|\d{4}-\d{2}-\d{2}$
• LIKE's _ matches any character, so us|2024-__-__ would also match us|2024-ab-cd. There's no way to require digits.

4. Structured mid-string alternation

Match either region, but only for a specific month:

• regex: partition_key_pattern=^(us|eu)\|2024-03
• The | OR in pattern/prefix_pattern splits the whole term, so it can't express "us OR eu, each followed by |2024-03".

[hussein-awala]

Hey @pierrejeambrun,

I did the security deep-dive on the regex filtering. Summary and what I've done:

The risk is real, and it's DB-side

partition_key_pattern is only syntax-checked with Python re.compile(), then handed to the database's own regex engine via SQLAlchemy regexp_match (~ on PostgreSQL, REGEXP on MySQL). So a malicious pattern burns database CPU, not app CPU, same impact profile as our prior ReDoS, CVE-2023-36543 (authenticated user hangs a request), which was fixed by switching to google-re2.

Per-engine picture (only the backends we support)

• PostgreSQL (primary): Henry Spencer's engine -> more resistant to classic "evil regex" than PCRE, but not immune; PostgreSQL's own guidance is to bound it with statement_timeout.
• MySQL 8.0+: ICU engine with built-in regexp_time_limit/regexp_stack_limit -> protected by default.
• MariaDB (PCRE, the classic backtracking engine): explicitly not supported by Airflow -> so out of scope.
• SQLite: dev-only, and Airflow doesn't even register a REGEXP function.

Important subtlety about re2

re2 fixes ReDoS when we do the matching (like CVE-2023-36543's Python-side fix). Here the database matches, so a Python-side re2 can't protect the DB engine. The SQL-world equivalent of "an engine that stops backtracking" is a statement timeout.

What I implemented

1. Opt-in feature flag: [api] enable_regexp_query_filters, default False. Regex filtering is off out of the box; a request using partition_key_pattern while disabled gets a 400. So there's zero attack surface unless an operator knowingly turns it on. Exact-match partition_key (B-tree indexed) is always available and unaffected.
2. Bounded query timeout: [api] regexp_query_timeout (seconds, default 30), enforced as a transaction-local statement_timeout on PostgreSQL so a pathological pattern is aborted instead of pinning a backend. MySQL relies on its built-in limit. 0 disables it.
3. Docs explain the security rationale for both settings; tests cover the disabled path and the timeout helper.

(I initially also added a pattern-length cap but dropped it, it doesn't stop short catastrophic patterns like (a+)+$ and risks rejecting legitimate ones, so it added little over the flag + timeout.)

Net: the feature is safe-by-default (off), and when enabled it's bounded by a configurable timeout on the one backend that needs it. Happy to adjust the default timeout or add anything else you'd like.

[pierrejeambrun]

Nice, thanks a lot for the investigation. This makes sense.

A few suggestions, feel free to disregards if this doesn't apply.

Would be great to have another pair of eyes especially for the task sdk part.

[pierrejeambrun]

Why not use QueryAssetEventPartitionKeyRegex and the other one?

This create duplicated validation code in the view

[hussein-awala]

Good catch, they were different when I started the implementation but it is not the case now.

The execution routes now depend on the shared QueryAssetEventPartitionKeyFilter / QueryAssetEventPartitionKeyRegex, applied via a filters argument on the SQL helper. The inline Query(...) params and the _validate_partition_key_params helper are gone, so validation (gate + regex compile) lives in one place.

[hussein-awala]

Thanks @pierrejeambrun! I've applied all of these in the last two commits and I updated the PR after #64611 is merged.

I will try to test all the new changes before merging

[pierrejeambrun]

LGTM, we might want a second pair of eyes as this is a significant change.

[pierrejeambrun]

Nit: This restores 0, but if there was a global timeout on the db side, this will override it. We should restore to the db side default.

[pierrejeambrun]

Maybe capture the value before and restore it after?

[pierrejeambrun]

The thing I wanted to prevent by moving the timeout inside the regex_param_factory or the filter class directly, is because if someone add a filter in here using the factory, and forget to call the timeout in the view -> CVE.

So I was trying to be defensive for a future bug.

[pierrejeambrun]

Maybe in paginated_select, add a check. If one of the filters is an instance of _RegexpParam and there is no timeout set for the request -> big server failure, refusing to process.

That doesn't cover 'all cases' someone could still use a bare _RegexpParam without using the paginated_select helper but that will cover 'most cases' of using filters.

And maybe something in the _RegexParam docstring, explaining that the caller query using this should time bound calling apply_regex_query_timeout

[ashb]

Mostly looks okay - not having a db query timeout on MySQL feels like something we need to fix

[ashb]

I think we can do a direct link to the config via

:conf:`api__regexp_query__timeout`

Something like this exists

[ashb]

Float maybe? To allow 0.5 etc?

[ashb]

So this timeout config doesn't apply to mysql? That feels risky?

[ashb]

Nit: it seems wasteful to create the same rows for each parameter set? Class scoped fixture perhaps? That could also delete/clean on teardown that way

[ashb]

These could be 4 cases in a single parameteized test i think?

[ashb]

I would have expected providing both to return a 400?

[ashb]

More Parameteization here too?