chore: import upstream snapshot with attribution

tests/tracing/processor/test_uc_table_processor.py

@@ -0,0 +1,180 @@
+from unittest import mock
+
+import pytest
+
+import mlflow
+import mlflow.tracking.context.default_context
+from mlflow.entities.span import LiveSpan
+from mlflow.entities.trace_location import TraceLocationType, UCSchemaLocation
+from mlflow.entities.trace_state import TraceState
+from mlflow.environment_variables import MLFLOW_TRACKING_USERNAME
+from mlflow.exceptions import MlflowException
+from mlflow.tracing.constant import TraceMetadataKey
+from mlflow.tracing.processor.uc_table import DatabricksUCTableSpanProcessor
+from mlflow.tracing.provider import _MLFLOW_TRACE_USER_DESTINATION
+from mlflow.tracing.trace_manager import InMemoryTraceManager
+
+from tests.tracing.helper import (
+    create_mock_otel_span,
+    create_test_trace_info,
+)
+
+
+@pytest.fixture
+def active_uc_schema_destination():
+    destination = UCSchemaLocation(catalog_name="catalog1", schema_name="schema1")
+    destination._otel_spans_table_name = "spans_table"
+    _MLFLOW_TRACE_USER_DESTINATION.set(destination)
+    try:
+        yield
+    finally:
+        _MLFLOW_TRACE_USER_DESTINATION.reset()
+
+
+def test_on_start_with_uc_table_name(monkeypatch, active_uc_schema_destination):
+    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")
+    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "alice")
+
+    # Root span should create a new trace on start
+    trace_id = 12345
+    span = create_mock_otel_span(trace_id=trace_id, span_id=1, parent_id=None, start_time=5_000_000)
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock.MagicMock())
+    processor.on_start(span)
+
+    # Check that trace was created in trace manager
+    trace_manager = InMemoryTraceManager.get_instance()
+    traces = trace_manager._traces
+    assert len(traces) == 1
+
+    # Get the created trace
+    created_trace = list(traces.values())[0]
+    trace_info = created_trace.info
+
+    # Verify trace location is UC_SCHEMA type
+    assert trace_info.trace_location.type == TraceLocationType.UC_SCHEMA
+    uc_schema = trace_info.trace_location.uc_schema
+    assert uc_schema.catalog_name == "catalog1"
+    assert uc_schema.schema_name == "schema1"
+
+    # Verify trace state and timing
+    assert trace_info.state == TraceState.IN_PROGRESS
+    assert trace_info.request_time == 5  # 5_000_000 nanoseconds -> 5 milliseconds
+    assert trace_info.execution_duration is None
+
+
+def test_on_start_without_uc_table_name(monkeypatch):
+    monkeypatch.setattr(mlflow.tracking.context.default_context, "_get_source_name", lambda: "test")
+    monkeypatch.setenv(MLFLOW_TRACKING_USERNAME.name, "alice")
+
+    # Root span should create a new trace on start
+    trace_id = 12345
+    span = create_mock_otel_span(trace_id=trace_id, span_id=1, parent_id=None, start_time=5_000_000)
+
+    _MLFLOW_TRACE_USER_DESTINATION.reset()
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock.MagicMock())
+    with pytest.raises(MlflowException, match="Unity Catalog spans table name is not set"):
+        processor.on_start(span)
+
+    # Check that trace was still created in trace manager
+    trace_manager = InMemoryTraceManager.get_instance()
+    traces = trace_manager._traces
+    assert len(traces) == 0
+
+
+def test_constructor_disables_metrics_export():
+    mock_exporter = mock.MagicMock()
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock_exporter)
+
+    # The export_metrics should be False
+    assert not processor._export_metrics
+
+
+def test_trace_id_generation_with_uc_schema(active_uc_schema_destination):
+    trace_id = 12345
+    span = create_mock_otel_span(trace_id=trace_id, span_id=1, parent_id=None, start_time=5_000_000)
+
+    with mock.patch(
+        "mlflow.tracing.processor.uc_table.generate_trace_id_v4",
+        return_value="trace:/catalog1.schema1/12345",
+    ) as mock_generate_trace_id:
+        processor = DatabricksUCTableSpanProcessor(span_exporter=mock.MagicMock())
+        processor.on_start(span)
+
+        # Verify generate_trace_id_v4 was called with correct arguments
+        mock_generate_trace_id.assert_called_once_with(span, "catalog1.schema1")
+
+
+def test_on_end():
+    trace_info = create_test_trace_info("request_id", 0)
+    trace_manager = InMemoryTraceManager.get_instance()
+    trace_manager.register_trace("trace_id", trace_info)
+
+    otel_span = create_mock_otel_span(
+        name="foo",
+        trace_id="trace_id",
+        span_id=1,
+        parent_id=None,
+        start_time=5_000_000,
+        end_time=9_000_000,
+    )
+    span = LiveSpan(otel_span, "request_id")
+    span.set_status("OK")
+    span.set_inputs({"input1": "test input"})
+    span.set_outputs({"output": "test output"})
+
+    mock_exporter = mock.MagicMock()
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock_exporter)
+
+    processor.on_end(otel_span)
+
+    # Verify span was exported
+    mock_exporter.export.assert_called_once_with((otel_span,))
+
+
+def test_on_end_sets_user_session_span_attributes():
+    trace_manager = InMemoryTraceManager.get_instance()
+    with mock.patch.object(trace_manager, "pop_trace", return_value=None):
+        with mlflow.start_span("foo") as live_span:
+            mlflow.update_current_trace(
+                metadata={
+                    TraceMetadataKey.TRACE_USER: "alice",
+                    TraceMetadataKey.TRACE_SESSION: "sess-123",
+                }
+            )
+            otel_span = live_span._span
+
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock.MagicMock())
+    processor.on_end(otel_span)
+
+    assert otel_span.attributes["user.id"] == "alice"
+    assert otel_span.attributes["session.id"] == "sess-123"
+
+
+def test_on_end_does_not_set_user_session_attributes_when_missing():
+    trace_manager = InMemoryTraceManager.get_instance()
+    with mock.patch.object(trace_manager, "pop_trace", return_value=None):
+        with mlflow.start_span("foo") as live_span:
+            otel_span = live_span._span
+
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock.MagicMock())
+    processor.on_end(otel_span)
+
+    assert "user.id" not in otel_span.attributes
+    assert "session.id" not in otel_span.attributes
+
+
+def test_trace_metadata_and_tags(active_uc_schema_destination):
+    trace_id = 12345
+    span = create_mock_otel_span(trace_id=trace_id, span_id=1, parent_id=None, start_time=5_000_000)
+    processor = DatabricksUCTableSpanProcessor(span_exporter=mock.MagicMock())
+    processor.on_start(span)
+
+    # Get the created trace
+    trace_manager = InMemoryTraceManager.get_instance()
+    traces = trace_manager._traces
+    created_trace = list(traces.values())[0]
+    trace_info = created_trace.info
+
+    # Check that metadata and tags are present
+    assert trace_info.trace_metadata is not None
+    assert trace_info.tags is not None