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".