Skip to content

Commit

Permalink
Load consuming_dags attr eagerly before dataset listener (#36247)
Browse files Browse the repository at this point in the history
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)
  • Loading branch information
dstandish authored and ephraimbuddy committed Dec 16, 2023
1 parent a0f4249 commit 9c5e820
Showing 1 changed file with 6 additions and 1 deletion.
7 changes: 6 additions & 1 deletion airflow/datasets/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from typing import TYPE_CHECKING

from sqlalchemy import exc, select
from sqlalchemy.orm import joinedload

from airflow.configuration import conf
from airflow.datasets import Dataset
Expand Down Expand Up @@ -63,7 +64,11 @@ def register_dataset_change(
For local datasets, look them up, record the dataset event, queue dagruns, and broadcast
the dataset event
"""
dataset_model = session.scalar(select(DatasetModel).where(DatasetModel.uri == dataset.uri))
dataset_model = session.scalar(
select(DatasetModel)
.where(DatasetModel.uri == dataset.uri)
.options(joinedload(DatasetModel.consuming_dags))
)
if not dataset_model:
self.log.warning("DatasetModel %s not found", dataset)
return
Expand Down

0 comments on commit 9c5e820

Please sign in to comment.