(airflow) branch main updated: Handle optional session in _refresh_from_db (#38572)

[ds...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dstandish pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 5c7b3e9fa7 Handle optional session in _refresh_from_db (#38572)
5c7b3e9fa7 is described below

commit 5c7b3e9fa7e9a46044f02ef7a31ebc0344cfb816
Author: Daniel Standish <15...@users.noreply.github.com>
AuthorDate: Wed Mar 27 17:59:44 2024 -0700

    Handle optional session in _refresh_from_db (#38572)
    
    Needed for AIP-44
---
 airflow/models/taskinstance.py | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py
index a254130e00..ca680d6e24 100644
--- a/airflow/models/taskinstance.py
+++ b/airflow/models/taskinstance.py
@@ -491,7 +491,10 @@ def _execute_task(task_instance: TaskInstance | TaskInstancePydantic, context: C
 
 
 def _refresh_from_db(
-    *, task_instance: TaskInstance | TaskInstancePydantic, session: Session, lock_for_update: bool = False
+    *,
+    task_instance: TaskInstance | TaskInstancePydantic,
+    session: Session | None = None,
+    lock_for_update: bool = False,
 ) -> None:
     """
     Refresh the task instance from the database based on the primary key.
@@ -504,7 +507,7 @@ def _refresh_from_db(
 
     :meta private:
     """
-    if task_instance in session:
+    if session and task_instance in session:
         session.refresh(task_instance, TaskInstance.__mapper__.column_attrs.keys())
 
     ti = TaskInstance.get_task_instance(