(airflow) branch main updated: fix memray decorator error handling (#60624)

[po...@apache.org]

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

potiuk 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 37d890bdabc fix memray decorator error handling (#60624)
37d890bdabc is described below

commit 37d890bdabce2bdce3eb319ecc8f6628c8fa9afd
Author: Jeongwoo Do <48...@users.noreply.github.com>
AuthorDate: Tue Feb 3 22:44:54 2026 +0900

    fix memray decorator error handling (#60624)
---
 airflow-core/src/airflow/utils/memray_utils.py     | 17 ++++++++---------
 airflow-core/tests/unit/utils/test_memray_utils.py | 12 ------------
 2 files changed, 8 insertions(+), 21 deletions(-)

diff --git a/airflow-core/src/airflow/utils/memray_utils.py b/airflow-core/src/airflow/utils/memray_utils.py
index 8a5132efc03..86d9e4e8493 100644
--- a/airflow-core/src/airflow/utils/memray_utils.py
+++ b/airflow-core/src/airflow/utils/memray_utils.py
@@ -69,15 +69,14 @@ def enable_memray_trace(component: MemrayTraceComponents) -> Callable[[Callable[
                     return func(*args, **kwargs)
             except ImportError as error:
                 # Silently fall back to running without tracking
-                log.warning(
-                    "ImportError memray.Tracker: %s in %s, please check the memray is installed",
-                    error.msg,
-                    component.value,
-                )
-                return func(*args, **kwargs)
-            except Exception as exception:
-                log.warning("Fail to apply memray.Tracker in %s, error: %s", component.value, exception)
-                return func(*args, **kwargs)
+                if "memray" in str(error):
+                    log.warning(
+                        "ImportError memray.Tracker: %s in %s, please check the memray is installed",
+                        error.msg,
+                        component.value,
+                    )
+                    return func(*args, **kwargs)
+                raise error
 
         return wrapper
 
diff --git a/airflow-core/tests/unit/utils/test_memray_utils.py b/airflow-core/tests/unit/utils/test_memray_utils.py
index 4d8eb47c8aa..e04a0cb8013 100644
--- a/airflow-core/tests/unit/utils/test_memray_utils.py
+++ b/airflow-core/tests/unit/utils/test_memray_utils.py
@@ -163,15 +163,3 @@ class TestEnableMemrayTrackErrorHandling:
 
             self.mock_function.assert_called_once_with("arg1")
             assert result == "test_result"
-
-    @conf_vars({("profiling", "memray_trace_components"): "scheduler,api,dag_processor"})
-    def test_graceful_fallback_on_exception(self):
-        """
-        Verify graceful degradation when exception occurs
-        """
-        with patch("memray.Tracker", side_effect=RuntimeError("Failed to initialize tracker")):
-            decorated_function = enable_memray_trace(MemrayTraceComponents.dag_processor)(self.mock_function)
-            result = decorated_function("arg1")
-
-            self.mock_function.assert_called_once_with("arg1")
-            assert result == "test_result"