(airflow) branch main updated: Fix Breeze unit test (#50395)

[kaxilnaik]

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

kaxilnaik 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 3cf457f89cc Fix Breeze unit test (#50395)
3cf457f89cc is described below

commit 3cf457f89ccae3e557bfddb5446d7a0b5946562c
Author: Kaxil Naik <kaxiln...@apache.org>
AuthorDate: Fri May 9 18:35:29 2025 +0530

    Fix Breeze unit test (#50395)
---
 .../src/airflow/providers/common/compat/standard/operators.py  |  6 +++++-
 .../src/airflow/providers/snowflake/operators/snowpark.py      | 10 +---------
 2 files changed, 6 insertions(+), 10 deletions(-)

diff --git 
a/providers/common/compat/src/airflow/providers/common/compat/standard/operators.py
 
b/providers/common/compat/src/airflow/providers/common/compat/standard/operators.py
index c53e8b53e92..c8838fe3c21 100644
--- 
a/providers/common/compat/src/airflow/providers/common/compat/standard/operators.py
+++ 
b/providers/common/compat/src/airflow/providers/common/compat/standard/operators.py
@@ -39,8 +39,12 @@ else:
             _SERIALIZERS,
             PythonOperator,
             ShortCircuitOperator,
-            get_current_context,
         )
 
+    try:
+        from airflow.sdk import get_current_context
+    except (ImportError, ModuleNotFoundError):
+        from airflow.providers.standard.operators.python import 
get_current_context
+
 
 __all__ = ["PythonOperator", "_SERIALIZERS", "ShortCircuitOperator", 
"get_current_context"]
diff --git 
a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark.py 
b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark.py
index 939fffef19c..98c96923ddb 100644
--- a/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark.py
+++ b/providers/snowflake/src/airflow/providers/snowflake/operators/snowpark.py
@@ -20,7 +20,7 @@ from __future__ import annotations
 from collections.abc import Collection, Mapping, Sequence
 from typing import Any, Callable
 
-from airflow.providers.common.compat.standard.operators import PythonOperator
+from airflow.providers.common.compat.standard.operators import PythonOperator, 
get_current_context
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 from airflow.providers.snowflake.utils.snowpark import 
inject_session_into_op_kwargs
 
@@ -115,14 +115,6 @@ class SnowparkOperator(PythonOperator):
             session_parameters=self.session_parameters,
         )
         session = hook.get_snowpark_session()
-
-        from airflow.providers.snowflake.version_compat import 
AIRFLOW_V_3_0_PLUS
-
-        if AIRFLOW_V_3_0_PLUS:
-            from airflow.sdk import get_current_context
-        else:
-            from airflow.providers.standard.operators.python import 
get_current_context
-
         context = get_current_context()
         session.update_query_tag(
             {