From 9c5e8202f6cce476e8ec65f38b6954782cd679a6 Mon Sep 17 00:00:00 2001
From: Daniel Standish <15932138+dstandish@users.noreply.github.com>
Date: Sat, 16 Dec 2023 00:10:44 -0800
Subject: [PATCH] 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 5bd35811da6a9294ba3ee42ec83a9103eab89f04)
---
 airflow/datasets/manager.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/airflow/datasets/manager.py b/airflow/datasets/manager.py
index 08871c9f6591f..047f8494a481d 100644
--- a/airflow/datasets/manager.py
+++ b/airflow/datasets/manager.py
@@ -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
@@ -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