refactor(plugins)!: use OpenTelemetry for BigQuery plugin tracing

This refactors the BigQueryAgentAnalyticsPlugin to use the standard OpenTelemetry API for trace and span ID generation and propagation, replacing the custom ContextVar implementation.

Key changes:
- Utilizes `opentelemetry.trace` for starting/ending spans.
- Correctly uses `opentelemetry.context` for context attachment and detachment.
- Span information is now derived from the OpenTelemetry context when available.
- Added a fallback mechanism to ensure span_id and parent_span_id are still populated if the OpenTelemetry SDK is not initialized.

To get standard OpenTelemetry trace information in BigQuery logs, users should install `opentelemetry-sdk` and initialize a global `TracerProvider` in their application *before* initializing ADK components.

Example minimal initialization:
```python
# Install: pip install opentelemetry-sdk
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
trace.set_tracer_provider(TracerProvider())
```

PiperOrigin-RevId: 858965562
This commit is contained in:
Google Team Member
2026-01-21 01:05:53 -08:00
committed by Copybara-Service
parent 645bd0ed10
commit ab89d12834
2 changed files with 383 additions and 117 deletions
@@ -37,6 +37,7 @@ import google.auth.credentials
from google.cloud import bigquery
from google.cloud import exceptions as cloud_exceptions
from google.genai import types
from opentelemetry import trace
import pyarrow as pa
import pytest
@@ -502,28 +503,34 @@ class TestBigQueryAgentAnalyticsPlugin:
bigquery_agent_analytics_plugin.TraceManager.init_trace(callback_context)
async def branch_1():
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, span_id="span-1"
s_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, span_name="span-1"
)
await asyncio.sleep(0.02)
s_id = bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
current_s_id = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
assert s_id == current_s_id
bigquery_agent_analytics_plugin.TraceManager.pop_span()
return s_id
async def branch_2():
bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, span_id="span-2"
s_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, span_name="span-2"
)
await asyncio.sleep(0.02)
s_id = bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
current_s_id = (
bigquery_agent_analytics_plugin.TraceManager.get_current_span_id()
)
assert s_id == current_s_id
bigquery_agent_analytics_plugin.TraceManager.pop_span()
return s_id
# Run concurrently
results = await asyncio.gather(branch_1(), branch_2())
# If they shared the same list/dict, they would interfere.
assert "span-1" in results
assert "span-2" in results
assert results[0] is not None
assert results[1] is not None
assert results[0] != results[1]
@pytest.mark.asyncio
@@ -1953,3 +1960,104 @@ class TestBigQueryAgentAnalyticsPlugin:
content_json = json.loads(log_entry["content"])
assert content_json["result"]["id"] == "inc-123"
assert content_json["result"]["kpi_missed"][0]["kpi"] == "latency"
@pytest.mark.asyncio
async def test_otel_integration(
self,
callback_context,
):
"""Verifies OpenTelemetry integration in TraceManager."""
# Mock the tracer and span
mock_tracer = mock.Mock()
mock_span = mock.Mock()
mock_context = mock.Mock()
# Setup mock IDs (128-bit trace_id, 64-bit span_id)
trace_id_int = 0x12345678123456781234567812345678
span_id_int = 0x1234567812345678
mock_context.trace_id = trace_id_int
mock_context.span_id = span_id_int
mock_context.is_valid = True
mock_span.get_span_context.return_value = mock_context
mock_span.start_time = 1234567890000000000 # Mock start time in ns
mock_tracer.start_span.return_value = mock_span
# Patch the global tracer in the plugin module
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.tracer", mock_tracer
):
# Test push_span
span_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "test_span"
)
mock_tracer.start_span.assert_called_with("test_span")
assert span_id == format(span_id_int, "016x")
# Test get_trace_id
# We need to mock trace.get_current_span() to return our mock span
# because push_span calls trace.attach(), which affects the global context
with mock.patch(
"opentelemetry.trace.get_current_span", return_value=mock_span
):
trace_id = bigquery_agent_analytics_plugin.TraceManager.get_trace_id(
callback_context
)
assert trace_id == format(trace_id_int, "032x")
# Test pop_span
# pop_span calls span.end()
bigquery_agent_analytics_plugin.TraceManager.pop_span()
mock_span.end.assert_called_once()
@pytest.mark.asyncio
async def test_otel_integration_real_provider(self, callback_context):
"""Verifies TraceManager with a real OpenTelemetry TracerProvider."""
# Setup OTEL with in-memory exporter
# pylint: disable=g-import-not-at-top
from opentelemetry.sdk import trace as trace_sdk
from opentelemetry.sdk.trace import export as trace_export
from opentelemetry.sdk.trace.export import in_memory_span_exporter
# pylint: enable=g-import-not-at-top
provider = trace_sdk.TracerProvider()
exporter = in_memory_span_exporter.InMemorySpanExporter()
processor = trace_export.SimpleSpanProcessor(exporter)
provider.add_span_processor(processor)
tracer = provider.get_tracer("test_tracer")
# Patch the global tracer in the plugin module
with mock.patch(
"google.adk.plugins.bigquery_agent_analytics_plugin.tracer", tracer
):
# 1. Start a span
span_id = bigquery_agent_analytics_plugin.TraceManager.push_span(
callback_context, "test_span"
)
# Verify a span was started but not ended
current_spans = exporter.get_finished_spans()
assert not current_spans
# Verify we can retrieve the trace ID
trace_id = bigquery_agent_analytics_plugin.TraceManager.get_trace_id(
callback_context
)
assert trace_id is not None
# 2. End the span
popped_span_id, _ = (
bigquery_agent_analytics_plugin.TraceManager.pop_span()
)
assert popped_span_id == span_id
# Verify span is now finished and exported
finished_spans = exporter.get_finished_spans()
assert len(finished_spans) == 1
assert finished_spans[0].name == "test_span"
assert format(finished_spans[0].context.span_id, "016x") == span_id
assert format(finished_spans[0].context.trace_id, "032x") == trace_id