diff --git a/src/snowflake/snowpark/_internal/open_telemetry.py b/src/snowflake/snowpark/_internal/open_telemetry.py index 3cab5710ecc..e8a4f5df03b 100644 --- a/src/snowflake/snowpark/_internal/open_telemetry.py +++ b/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): diff --git a/tests/integ/test_open_telemetry.py b/tests/integ/test_open_telemetry.py index 0beb7572cf0..3a5b5d9d15c 100644 --- a/tests/integ/test_open_telemetry.py +++ b/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", diff --git a/tests/unit/test_open_telemetry.py b/tests/unit/test_open_telemetry.py index 586d9f670c6..af9f0fc7e43 100644 --- a/tests/unit/test_open_telemetry.py +++ b/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(