From 2a5613ef39a3d20887055fc92b1305b4deef190b Mon Sep 17 00:00:00 2001 From: Ankush Malaker <43288948+AnkushMalaker@users.noreply.github.com> Date: Sun, 22 Feb 2026 22:19:44 +0530 Subject: [PATCH] refactor: improve handling of project and logstream attributes in GalileoOTLPExporter and GalileoSpanProcessor - Updated the exporter to conditionally set headers only if project and logstream attributes are present. - Enhanced the span processor to fallback to environment variables for project and logstream if not provided in context. - Adjusted tests to verify the new behavior, ensuring that None values are skipped and environment variables are utilized correctly. --- src/galileo/otel.py | 21 ++++++++++------ tests/test_otel.py | 60 ++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 70 insertions(+), 11 deletions(-) diff --git a/src/galileo/otel.py b/src/galileo/otel.py index a968788c..4a5115cb 100644 --- a/src/galileo/otel.py +++ b/src/galileo/otel.py @@ -185,12 +185,12 @@ def export(self, spans: typing.Sequence[Any]) -> "Any": # for the last span update the headers if spans: last_span = spans[-1] - self._session.headers.update( - { - "project": last_span.attributes.get("galileo.project.name"), - "logstream": last_span.attributes.get("galileo.logstream.name"), - } - ) + project = last_span.attributes.get("galileo.project.name") + logstream = last_span.attributes.get("galileo.logstream.name") + if project: + self._session.headers["project"] = project + if logstream: + self._session.headers["logstream"] = logstream return super().export(spans) @@ -246,7 +246,12 @@ def __init__( _log_stream_context.set(logstream) self._project = _project_context.get() + if self._project is None and "GALILEO_PROJECT" in os.environ: + self._project = os.environ["GALILEO_PROJECT"] + self._logstream = _log_stream_context.get() + if self._logstream is None and "GALILEO_LOG_STREAM" in os.environ: + self._logstream = os.environ["GALILEO_LOG_STREAM"] # Create the exporter using the config-based approach self._exporter = GalileoOTLPExporter(**kwargs) @@ -259,8 +264,8 @@ def __init__( def on_start(self, span: Span, parent_context: Optional[context.Context] = None) -> None: """Handle span start events by delegating to the underlying processor.""" # Set Galileo context attributes on the span - project = _project_context.get(self._project) - log_stream = _log_stream_context.get(self._logstream) + project = _project_context.get(None) or self._project + log_stream = _log_stream_context.get(None) or self._logstream experiment_id = _experiment_id_context.get(None) session_id = _session_id_context.get(None) diff --git a/tests/test_otel.py b/tests/test_otel.py index 6584a399..f59fa089 100644 --- a/tests/test_otel.py +++ b/tests/test_otel.py @@ -375,7 +375,7 @@ def test_processor_on_start_sets_span_attributes(self, mock_processor_deps, rese assert ("galileo.project.name", "test-project") in actual_calls assert ("galileo.session.id", "test-session") in actual_calls - # Test with only project set (None values should be skipped) + # Test with only project set in context (None values should fall back to env vars or be skipped) _log_stream_context.set(None) _experiment_id_context.set(None) _session_id_context.set(None) @@ -384,8 +384,11 @@ def test_processor_on_start_sets_span_attributes(self, mock_processor_deps, rese mock_span2 = Mock() processor2.on_start(mock_span2, None) - assert mock_span2.set_attribute.call_count == 1 - mock_span2.set_attribute.assert_called_with("galileo.project.name", "test-project") + # project from context, logstream from env var fallback, experiment/session skipped + actual_calls2 = {(args[0], args[1]) for args, _ in mock_span2.set_attribute.call_args_list} + assert ("galileo.project.name", "test-project") in actual_calls2 + assert "galileo.experiment.id" not in {c[0] for c in actual_calls2} + assert "galileo.session.id" not in {c[0] for c in actual_calls2} @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available") @patch("galileo.otel.OTLPSpanExporter.export") @@ -453,3 +456,54 @@ def test_exporter_export_merges_resource_attributes(self, mock_resource_class, m mock_resource_class.assert_not_called() mock_span2.resource.merge.assert_not_called() mock_parent_export.assert_called_once_with([mock_span2]) + + @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available") + def test_processor_env_var_fallback(self, mock_processor_deps, reset_decorator_context): + """Test processor reads env vars when no args or context are provided.""" + # Given: context vars are cleared and env vars are set + with patch.dict(os.environ, {"GALILEO_PROJECT": "env-project", "GALILEO_LOG_STREAM": "env-logstream"}): + # When: creating a processor with no explicit args + processor = GalileoSpanProcessor() + + # Then: processor picks up values from env vars + assert processor._project == "env-project" + assert processor._logstream == "env-logstream" + + # And: on_start sets the correct span attributes + mock_span = Mock() + processor.on_start(mock_span, None) + + actual_calls = {(args[0], args[1]) for args, _ in mock_span.set_attribute.call_args_list} + assert ("galileo.project.name", "env-project") in actual_calls + assert ("galileo.logstream.name", "env-logstream") in actual_calls + + @pytest.mark.skipif(not OTEL_AVAILABLE, reason="OpenTelemetry not available") + @patch("galileo.otel.OTLPSpanExporter.export") + def test_export_does_not_overwrite_headers_with_none(self, mock_parent_export, reset_decorator_context): + """Test export preserves original headers when spans lack galileo attributes.""" + # Given: an exporter with explicit project/logstream headers + with ( + patch("galileo.otel.OTLPSpanExporter.__init__", return_value=None), + patch("galileo.otel.GalileoPythonConfig.get") as mock_config_get, + ): + config = Mock() + config.api_url = "https://api.galileo.ai" + config.api_key = SecretStr("test-key") + mock_config_get.return_value = config + + exporter = GalileoOTLPExporter(project="original-project", logstream="original-logstream") + exporter._session = Mock() + exporter._session.headers = { + "project": "original-project", + "logstream": "original-logstream", + } + + # When: exporting spans that have no galileo.* attributes + mock_span = Mock() + mock_span.attributes = {"some.other.attribute": "value"} + mock_span.resource = Mock() + exporter.export([mock_span]) + + # Then: the original headers are preserved, not overwritten with None + assert exporter._session.headers["project"] == "original-project" + assert exporter._session.headers["logstream"] == "original-logstream"