SNOW-1473726: Dataframe tracing should record exceptions during actio…

…n functions (#1755)

src/snowflake/snowpark/_internal/open_telemetry.py

@@ -56,7 +56,6 @@ def open_telemetry_context_manager(func, dataframe):
                 logger.warning(f"Error when acquiring span attributes. {e}")
             finally:
                 yield
-
     else:
         yield
 
@@ -101,11 +100,9 @@ def open_telemetry_udf_context_manager(
             except Exception as e:
                 logger.warning(f"Error when acquiring span attributes. {e}")
             finally:
-                try:
-                    yield
-                except Exception as e:
-                    cur_span.set_status(trace.Status(trace.StatusCode.ERROR, str(e)))
-                    raise e
+                yield
+    else:
+        yield
 
 
 def decorator_count(func):

tests/integ/test_open_telemetry.py

@@ -13,6 +13,7 @@
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
 
 import snowflake
+from snowflake.snowpark.exceptions import SnowparkSQLException
 from snowflake.snowpark.functions import sproc, udaf, udf, udtf
 from snowflake.snowpark.types import (
     BinaryType,
@@ -54,6 +55,37 @@ def dict_exporter():
     dict_exporter.shutdown()
 
 
+def test_catch_error_during_action_function(session, dict_exporter):
+    df = session.sql("select 1/0")
+    with pytest.raises(SnowparkSQLException, match="Division by zero"):
+        df.collect()
+    spans = spans_to_dict(dict_exporter.get_finished_spans())
+    assert "collect" in spans
+    span = spans["collect"]
+    assert span.status.status_code == trace.status.StatusCode.ERROR
+    assert "Division by zero" in span.status.description
+    dict_exporter.clear()
+
+
+def test_catch_error_during_registration_function(session, dict_exporter):
+    with pytest.raises(ValueError, match="does not exist"):
+        session.udf.register_from_file(
+            "empty file",
+            "mod5",
+            name="mod5_function",
+            return_type=IntegerType(),
+            input_types=[IntegerType()],
+            replace=True,
+            immutable=True,
+        )
+    spans = spans_to_dict(dict_exporter.get_finished_spans())
+    span = spans["register_from_file"]
+    assert span.status.status_code == trace.status.StatusCode.ERROR
+    assert "does not exist" in span.status.description
+
+    dict_exporter.clear()
+
+
 @pytest.mark.skipif(
     IS_IN_STORED_PROC,
     reason="Cannot create session in SP",

tests/unit/test_open_telemetry.py

@@ -84,6 +84,30 @@ def dict_exporter():
     dict_exporter.shutdown()
 
 
+def test_without_open_telemetry(monkeypatch, dict_exporter):
+    from snowflake.snowpark._internal import open_telemetry
+
+    monkeypatch.setattr(open_telemetry, "open_telemetry_found", False)
+    mock_connection = mock.create_autospec(ServerConnection)
+    mock_connection._conn = mock.MagicMock()
+    session = snowflake.snowpark.session.Session(mock_connection)
+    _add_session(session)
+    session._conn._telemetry_client = mock.MagicMock()
+    session.create_dataframe([1, 2, 3, 4]).to_df("a").collect()
+
+    spans = spans_to_dict(dict_exporter.get_finished_spans())
+    assert len(spans) == 0
+    dict_exporter.clear()
+
+    @udf(name="minus", session=session)
+    def minus_udf(x: int, y: int) -> int:
+        return x - y
+
+    spans = spans_to_dict(dict_exporter.get_finished_spans())
+    assert len(spans) == 0
+    dict_exporter.clear()
+
+
 def test_register_udaf_from_file(dict_exporter):
     test_file = os.path.normpath(
         os.path.join(