From 4596929a7b999ff760a9c3027e8ea24252adccc2 Mon Sep 17 00:00:00 2001 From: Pankaj Koti Date: Mon, 21 Oct 2024 17:59:09 +0530 Subject: [PATCH] Try fixing NEW_SESSION import in Airflow 2.2 --- tests/utils.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/utils.py b/tests/utils.py index 1f73b69..5b00d72 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -12,11 +12,21 @@ from airflow.models.taskinstance import TaskInstance from airflow.secrets.local_filesystem import LocalFilesystemBackend from airflow.utils import timezone -from airflow.utils.session import NEW_SESSION, provide_session +from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, State from airflow.utils.types import DagRunType from sqlalchemy.orm.session import Session +try: + from airflow.utils.session import NEW_SESSION +except ImportError: + # Airflow < 2.3 did not have NEW_SESSION in airflow.utils.session + from typing import cast + + from airflow import settings + + NEW_SESSION: settings.SASession = cast(settings.SASession, None) + log = logging.getLogger(__name__)