Skip to content

Commit 9c5e820

Browse files
dstandishephraimbuddy
authored andcommitted
Load consuming_dags attr eagerly before dataset listener (#36247)
Previously, we could get detached instance error if the dataset listener code closes the session (which happens e.g. if BaseHook.get_connection is called). By eager loading the attribute (which we need anyway), we can avoid this problem. (cherry picked from commit 5bd3581)
1 parent a0f4249 commit 9c5e820

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

airflow/datasets/manager.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
from typing import TYPE_CHECKING
2121

2222
from sqlalchemy import exc, select
23+
from sqlalchemy.orm import joinedload
2324

2425
from airflow.configuration import conf
2526
from airflow.datasets import Dataset
@@ -63,7 +64,11 @@ def register_dataset_change(
6364
For local datasets, look them up, record the dataset event, queue dagruns, and broadcast
6465
the dataset event
6566
"""
66-
dataset_model = session.scalar(select(DatasetModel).where(DatasetModel.uri == dataset.uri))
67+
dataset_model = session.scalar(
68+
select(DatasetModel)
69+
.where(DatasetModel.uri == dataset.uri)
70+
.options(joinedload(DatasetModel.consuming_dags))
71+
)
6772
if not dataset_model:
6873
self.log.warning("DatasetModel %s not found", dataset)
6974
return

0 commit comments

Comments
 (0)