feat(otel): Switch CloudTraceSpanExporter to telemetry.googleapis.com

PiperOrigin-RevId: 815675872

Author: RKest

pyproject.toml

@@ -52,8 +52,8 @@ dependencies = [
   "python-dateutil>=2.9.0.post0, <3.0.0",                   # For Vertext AI Session Service
   "python-dotenv>=1.0.0, <2.0.0",                           # To manage environment variables
   "requests>=2.32.4, <3.0.0",
-  "sqlalchemy>=2.0, <3.0.0",                                # SQL database ORM
   "sqlalchemy-spanner>=1.14.0",                             # Spanner database session service
+  "sqlalchemy>=2.0, <3.0.0",                                # SQL database ORM
   "starlette>=0.46.2, <1.0.0",                              # For FastAPI CLI
   "tenacity>=8.0.0, <9.0.0",                                # For Retry management
   "typing-extensions>=4.5, <5",

src/google/adk/cli/adk_web_server.py

@@ -292,8 +292,9 @@ def _setup_telemetry(
   else:
     # Old logic - to be removed when above leaves experimental.
     tracer_provider = TracerProvider()
-    for exporter in internal_exporters:
-      tracer_provider.add_span_processor(exporter)
+    if internal_exporters is not None:
+      for exporter in internal_exporters:
+        tracer_provider.add_span_processor(exporter)
     trace.set_tracer_provider(tracer_provider=tracer_provider)
 
 
@@ -312,19 +313,24 @@ def _otel_env_vars_enabled() -> bool:
 def _setup_gcp_telemetry_experimental(
     internal_exporters: list[SpanProcessor] = None,
 ):
-  from ..telemetry.setup import maybe_set_otel_providers
+  if typing.TYPE_CHECKING:
+    from ..telemetry.setup import OTelHooks
 
-  otel_hooks_to_add = []
-  otel_resource = None
+  otel_hooks_to_add: list[OTelHooks] = []
 
   if internal_exporters:
     from ..telemetry.setup import OTelHooks
 
     # Register ADK-specific exporters in trace provider.
     otel_hooks_to_add.append(OTelHooks(span_processors=internal_exporters))
 
+  import google.auth
+
   from ..telemetry.google_cloud import get_gcp_exporters
   from ..telemetry.google_cloud import get_gcp_resource
+  from ..telemetry.setup import maybe_set_otel_providers
+
+  credentials, project_id = google.auth.default()
 
   otel_hooks_to_add.append(
       get_gcp_exporters(
@@ -334,12 +340,14 @@ def _setup_gcp_telemetry_experimental(
           # TODO - reenable metrics once errors during shutdown are fixed.
           enable_cloud_metrics=False,
           enable_cloud_logging=True,
+          google_auth=(credentials, project_id),
       )
   )
-  otel_resource = get_gcp_resource()
+  otel_resource = get_gcp_resource(project_id)
 
   maybe_set_otel_providers(
-      otel_hooks_to_setup=otel_hooks_to_add, otel_resource=otel_resource
+      otel_hooks_to_setup=otel_hooks_to_add,
+      otel_resource=otel_resource,
   )
   _setup_instrumentation_lib_if_installed()
 

src/google/adk/telemetry/google_cloud.py

@@ -15,6 +15,9 @@
 from __future__ import annotations
 
 import logging
+from typing import cast
+from typing import Optional
+from typing import TYPE_CHECKING
 
 import google.auth
 from opentelemetry.sdk._logs import LogRecordProcessor
@@ -29,6 +32,9 @@
 from ..utils.feature_decorator import experimental
 from .setup import OTelHooks
 
+if TYPE_CHECKING:
+  from google.auth.credentials import Credentials
+
 logger = logging.getLogger('google_adk.' + __name__)
 
 
@@ -37,34 +43,43 @@ def get_gcp_exporters(
     enable_cloud_tracing: bool = False,
     enable_cloud_metrics: bool = False,
     enable_cloud_logging: bool = False,
+    google_auth: Optional[tuple[Credentials, str]] = None,
 ) -> OTelHooks:
   """Returns GCP OTel exporters to be used in the app.
 
   Args:
     enable_tracing: whether to enable tracing to Cloud Trace.
     enable_metrics: whether to enable raporting metrics to Cloud Monitoring.
     enable_logging: whether to enable sending logs to Cloud Logging.
+    google_auth: optional custom credentials and project_id. google.auth.default() used when this is omitted.
   """
-  _, project_id = google.auth.default()
+
+  credentials, project_id = (
+      google_auth if google_auth is not None else google.auth.default()
+  )
+  if TYPE_CHECKING:
+    credentials = cast(Credentials, credentials)
+    project_id = cast(str, project_id)
+
   if not project_id:
     logger.warning(
         'Cannot determine GCP Project. OTel GCP Exporters cannot be set up.'
         ' Please make sure to log into correct GCP Project.'
     )
     return OTelHooks()
 
-  span_processors = []
+  span_processors: list[SpanProcessor] = []
   if enable_cloud_tracing:
-    exporter = _get_gcp_span_exporter(project_id)
+    exporter = _get_gcp_span_exporter(credentials)
     span_processors.append(exporter)
 
-  metric_readers = []
+  metric_readers: list[MetricReader] = []
   if enable_cloud_metrics:
     exporter = _get_gcp_metrics_exporter(project_id)
     if exporter:
       metric_readers.append(exporter)
 
-  log_record_processors = []
+  log_record_processors: list[LogRecordProcessor] = []
   if enable_cloud_logging:
     exporter = _get_gcp_logs_exporter(project_id)
     if exporter:
@@ -77,10 +92,18 @@ def get_gcp_exporters(
   )
 
 
-def _get_gcp_span_exporter(project_id: str) -> SpanProcessor:
-  from opentelemetry.exporter.cloud_trace import CloudTraceSpanExporter
+def _get_gcp_span_exporter(credentials: Credentials) -> SpanProcessor:
+  """Adds OTEL span exporter to telemetry.googleapis.com"""
+
+  from google.auth.transport.requests import AuthorizedSession
+  from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
-  return BatchSpanProcessor(CloudTraceSpanExporter(project_id=project_id))
+  return BatchSpanProcessor(
+      OTLPSpanExporter(
+          session=AuthorizedSession(credentials=credentials),
+          endpoint='https://telemetry.googleapis.com/v1/traces',
+      )
+  )
 
 
 def _get_gcp_metrics_exporter(project_id: str) -> MetricReader:
@@ -101,15 +124,22 @@ def _get_gcp_logs_exporter(project_id: str) -> LogRecordProcessor:
   )
 
 
-def get_gcp_resource() -> Resource:
-  # The OTELResourceDetector populates resource labels from
-  # environment variables like OTEL_SERVICE_NAME and OTEL_RESOURCE_ATTRIBUTES.
-  # Then the GCP detector adds attributes corresponding to a correct
-  # monitored resource if ADK runs on one of supported platforms
-  # (e.g. GCE, GKE, CloudRun).
-
-  resource = OTELResourceDetector().detect()
+def get_gcp_resource(project_id: Optional[str] = None) -> Resource:
+  """Returns OTEL with attributes specified in the following order (attributes specified later, overwrite those specified earlier):
+  1. Populates gcp.project_id attribute from the project_id argument if present.
+  2. OTELResourceDetector populates resource labels from environment variables like OTEL_SERVICE_NAME and OTEL_RESOURCE_ATTRIBUTES.
+  3. GCP detector adds attributes corresponding to a correct monitored resource if ADK runs on one of supported platforms (e.g. GCE, GKE, CloudRun).
 
+  Args:
+    project_id: project id to fill out as `gcp.project_id` on the OTEL resource.
+    This may be overwritten by OTELResourceDetector, if `gcp.project_id` is present in `OTEL_RESOURCE_ATTRIBUTES` env var.
+  """
+  resource = Resource(
+      attributes={'gcp.project_id': project_id}
+      if project_id is not None
+      else {}
+  )
+  resource = resource.merge(OTELResourceDetector().detect())
   try:
     from opentelemetry.resourcedetector.gcp_resource_detector import GoogleCloudResourceDetector
 
@@ -121,5 +151,4 @@ def get_gcp_resource() -> Resource:
         'Cloud not import opentelemetry.resourcedetector.gcp_resource_detector'
         ' GCE, GKE or CloudRun related resource attributes may be missing'
     )
-
   return resource

src/google/adk/telemetry/setup.py

@@ -73,10 +73,8 @@ def maybe_set_otel_providers(
     otel_resource: OTel resource to use in providers.
     If empty - default OTel resource detection will be used.
   """
-  if otel_hooks_to_setup is None:
-    otel_hooks_to_setup = []
-  if otel_resource is None:
-    otel_resource = _get_otel_resource()
+  otel_hooks_to_setup = otel_hooks_to_setup or []
+  otel_resource = otel_resource or _get_otel_resource()
 
   # Add generic OTel exporters based on OTel env variables.
   otel_hooks_to_setup.append(_get_otel_exporters())

tests/unittests/telemetry/test_google_cloud.py

@@ -12,9 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+from typing import Optional
 from unittest import mock
 
 from google.adk.telemetry.google_cloud import get_gcp_exporters
+from google.adk.telemetry.google_cloud import get_gcp_resource
 import pytest
 
 
@@ -55,3 +58,34 @@ def test_get_gcp_exporters(
   assert len(otel_hooks.log_record_processors) == (
       1 if enable_cloud_logging else 0
   )
+
+
+@pytest.mark.parametrize("project_id_in_arg", ["project_id_in_arg", None])
+@pytest.mark.parametrize("project_id_on_env", ["project_id_on_env", None])
+def test_get_gcp_resource(
+    project_id_in_arg: Optional[str],
+    project_id_on_env: Optional[str],
+    monkeypatch: pytest.MonkeyPatch,
+):
+  # Arrange.
+  if project_id_on_env is not None:
+    monkeypatch.setenv(
+        "OTEL_RESOURCE_ATTRIBUTES", f"gcp.project_id={project_id_on_env}"
+    )
+
+  # Act.
+  otel_resource = get_gcp_resource(project_id_in_arg)
+
+  # Assert.
+  expected_project_id = (
+      project_id_on_env
+      if project_id_on_env is not None
+      else project_id_in_arg
+      if project_id_in_arg is not None
+      else None
+  )
+  assert otel_resource is not None
+  assert (
+      otel_resource.attributes.get("gcp.project_id", None)
+      == expected_project_id
+  )