find_orm_dags: multiple joinedload() collections cause row-explosion, slow DagModel sync (commonly >10s)
- Dominant language
- Python
- Stars
- 46.9k
- Forks
- 17.8k
- Avg merge
- 2d 9h
- Merged PRs (30d)
- 472
Description
## Summary
`DagModelOperation.find_orm_dags` (`airflow-core/src/airflow/dag_processing/collection.py`) eagerly loads five separate one-to-many collections via `joinedload()` in a single query:
```python
def find_orm_dags(self, *, session: Session) -> dict[str, DagModel]:
stmt: Select[Unpack[tuple[DagModel]]] = with_row_locks(
(
select(DagModel)
.options(joinedload(DagModel.tags, innerjoin=False))
.where(DagModel.dag_id.in_(self.dags))
.options(joinedload(DagModel.schedule_asset_references))
.options(joinedload(DagModel.schedule_asset_alias_references))
.options(joinedload(DagModel.task_outlet_asset_references))
.options(joinedload(DagModel.dag_owner_links))
),
of=DagModel,
session=session,
)
return {dm.dag_id: dm for dm in session.scalars(stmt).unique()}
```
Joining more than one one-to-many collection in the same query via `joinedload` is a well-known SQLAlchemy anti-pattern: it produces a cartesian-product row explosion. A DAG with, say, 3 tags × 2 asset references × 2 owner links returns 12 duplicate rows, each carrying the full wide `dag.*` column set repeated. Called across however many `dag_id`s are in a given `update_dag_parsing_results_in_db` call (which scales with DAG count per file/sweep), this multiplies fast.
## Evidence (production)
Caught live via `pg_stat_activity` on a customer deployment with a large dynamically-generated DAG set (1,400+ DAGs across several files):
```
runtime: 00:00:10.241513, state: active, wait_event: ClientWrite, wait_event_type: Client
application_name: astro-agent 1.13.1 [task:...]
query: SELECT dag.dag_id, dag.is_paused, ... dag_tag_1.name, dag_tag_1.dag_id AS dag_id_1,
dag_owner_attributes_1.dag_id AS dag_id_2, dag_owner_attributes_1.o...
```
`wait_event: ClientWrite` is the key detail — Postgres has already computed the result and is blocked *sending* it, because the client (asyncpg connection) isn't draining the socket fast enough. That points at result-set size, not query planning/execution cost, as the bottleneck. This customer reports these commonly exceeding 10 seconds. Response latency downstream (this call sits in the request path of the DAG-processor's parse-result heartbeat) has been observed causing client-side request timeouts/disconnects (HTTP 499s) on the receiving end in the wild.
## Proposed fix
Replace the multiple `joinedload()` calls on one-to-many collections with `selectinload()`. `selectinload` issues one follow-up `WHERE dag_id IN (...)` query per collection instead of one giant multi-way join, giving the same eager-loading outcome without the multiplicative row blowup. Since there are five separate collections here, this trades one huge query for five small ones — plausibly still a large net win, but worth benchmarking against a DAG set with many tags/owners/asset-refs per DAG to confirm before merging.
Contributor guide
Research direction
Start in airflow-core/src/airflow/dag_processing/collection.py at DagModelOperation.find_orm_dags and inspect the joined eager loads and their use in update_dag_parsing_results_in_db. Benchmark a DAG set with many tags, owners, and asset references, comparing result size and latency before and after the loading change. Done means the collections remain eagerly available without the multi-way row explosion and the measured sync time improves.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, sqlalchemy
- Domain
- backend, data-engineering, databases
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 72/100