fix(storage): add sync context manager fallback and prevent trace state leakage across unit tests

Author: shradhakatyal

packages/google-cloud-storage/google/cloud/storage/_bucket_metadata_cache.py

@@ -61,7 +61,7 @@ def get_or_queue_fetch(self, bucket_name):
             else:
                 # fire a background fetch and get bucket metadata.
                 self._inflight_fetches.add(bucket_name)
-                if hasattr(self._client, "grpc_client"):
+                if getattr(self._client, "_is_async_grpc_client", False) is True:
                     try:
                         loop = asyncio.get_running_loop()
                         loop.create_task(self._fetch_background_async(bucket_name))
@@ -81,7 +81,7 @@ def check_and_evict(self, bucket_name):
             if bucket_name in self._inflight_checks:
                 return
             self._inflight_checks.add(bucket_name)
-            if hasattr(self._client, "grpc_client"):
+            if getattr(self._client, "_is_async_grpc_client", False) is True:
                 try:
                     loop = asyncio.get_running_loop()
                     loop.create_task(
@@ -178,25 +178,31 @@ def update_from_bucket(self, bucket, bucket_name=None):
         if not bucket:
             return
         name = bucket_name or getattr(bucket, "name", None)
-        if not name:
+        if not name or not isinstance(name, str):
             return
         if name.startswith("projects/") and "/buckets/" in name:
             name = name.split("/buckets/", 1)[1]
 
         project_number = getattr(bucket, "project_number", None)
         if not project_number:
-            project = getattr(bucket, "project", None)
-            if project:
-                proj_str = str(project)
+            proj_attr = getattr(bucket, "project", None)
+            if isinstance(proj_attr, (str, int)):
+                proj_str = str(proj_attr)
                 if proj_str.startswith("projects/"):
                     project_number = proj_str.split("projects/", 1)[1]
-                else:
+                elif proj_str:
                     project_number = proj_str
 
-        location = getattr(bucket, "location", None) or "global"
-        location = location.lower()
-        location_type = getattr(bucket, "location_type", None) or "region"
-        location_type = location_type.lower()
+        loc_attr = getattr(bucket, "location", None)
+        location = (
+            loc_attr.lower() if isinstance(loc_attr, str) and loc_attr else "global"
+        )
+        loc_type_attr = getattr(bucket, "location_type", None)
+        location_type = (
+            loc_type_attr.lower()
+            if isinstance(loc_type_attr, str) and loc_type_attr
+            else "region"
+        )
 
         if location_type in ("multi-region", "dual-region"):
             location = "global"

packages/google-cloud-storage/google/cloud/storage/_helpers.py

@@ -19,6 +19,7 @@
 
 import base64
 import datetime
+import inspect
 import logging
 import os
 import secrets
@@ -32,18 +33,15 @@
 from google.auth import environment_vars
 from google.cloud.exceptions import NotFound
 
-from google.cloud.storage._opentelemetry_tracing import (
-    _is_bucket_metadata_disabled,
-)
-from google.cloud.storage._opentelemetry_tracing import (
-    create_trace_span as _base_create_trace_span,
-)
+from google.cloud.storage import _opentelemetry_tracing
 from google.cloud.storage.constants import _DEFAULT_TIMEOUT
 from google.cloud.storage.retry import (
     DEFAULT_RETRY,
     DEFAULT_RETRY_IF_METAGENERATION_SPECIFIED,
 )
 
+_base_create_trace_span = None
+
 _logger = logging.getLogger(__name__)
 
 STORAGE_EMULATOR_ENV_VAR = "STORAGE_EMULATOR_HOST"  # Despite name, includes scheme.
@@ -172,7 +170,8 @@ def _prepare_base_cm(self):
             and client
             and hasattr(client, "_bucket_metadata_cache")
             and client._bucket_metadata_cache
-            and not _is_bucket_metadata_disabled()
+            and _opentelemetry_tracing._is_otel_traces_enabled()
+            and not _opentelemetry_tracing._is_bucket_metadata_disabled()
         ):
             try:
                 if name in (
@@ -202,7 +201,12 @@ def _prepare_base_cm(self):
         if "client" not in kwargs and client:
             kwargs["client"] = client
 
-        self._base_cm = _base_create_trace_span(name, attributes=span_attrs, **kwargs)
+        create_span_fn = (
+            _base_create_trace_span
+            if _base_create_trace_span is not None
+            else _opentelemetry_tracing.create_trace_span
+        )
+        self._base_cm = create_span_fn(name, attributes=span_attrs, **kwargs)
         return self._base_cm
 
     def _handle_not_found(self):
@@ -212,6 +216,7 @@ def _handle_not_found(self):
             and self.client
             and hasattr(self.client, "_bucket_metadata_cache")
             and self.client._bucket_metadata_cache
+            and _opentelemetry_tracing._is_otel_traces_enabled()
         ):
             try:
                 self.client._bucket_metadata_cache.check_and_evict(self.bucket_name)
@@ -235,15 +240,21 @@ def __exit__(self, exc_type, exc_val, exc_tb):
 
     async def __aenter__(self):
         self._prepare_base_cm()
-        return await self._base_cm.__aenter__()
+        if hasattr(self._base_cm, "__aenter__"):
+            res = self._base_cm.__aenter__()
+            return await res if inspect.isawaitable(res) else res
+        return self._base_cm.__enter__()
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         if exc_val is not None and isinstance(
             exc_val, (NotFound, api_exceptions.NotFound)
         ):
             self._handle_not_found()
         if self._base_cm is not None:
-            return await self._base_cm.__aexit__(exc_type, exc_val, exc_tb)
+            if hasattr(self._base_cm, "__aexit__"):
+                res = self._base_cm.__aexit__(exc_type, exc_val, exc_tb)
+                return await res if inspect.isawaitable(res) else res
+            return self._base_cm.__exit__(exc_type, exc_val, exc_tb)
         return False
 
 

packages/google-cloud-storage/google/cloud/storage/_opentelemetry_tracing.py

@@ -43,6 +43,16 @@ def _is_bucket_metadata_disabled() -> bool:
 enable_otel_traces = _parse_bool_env(
     ENABLE_OTEL_TRACES_ENV_VAR, _DEFAULT_ENABLE_OTEL_TRACES_VALUE
 )
+
+
+def _is_otel_traces_enabled() -> bool:
+    if not HAS_OPENTELEMETRY or not enable_otel_traces:
+        return False
+    return _parse_bool_env(
+        ENABLE_OTEL_TRACES_ENV_VAR, _DEFAULT_ENABLE_OTEL_TRACES_VALUE
+    )
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -96,7 +106,7 @@ def __init__(
         self._span = None
 
     def __enter__(self):
-        if not HAS_OPENTELEMETRY or not enable_otel_traces:
+        if not _is_otel_traces_enabled():
             return None
 
         tracer = trace.get_tracer(__name__)

packages/google-cloud-storage/google/cloud/storage/asyncio/_utils.py

@@ -45,10 +45,7 @@ def update_write_handle_if_exists(obj, response):
 def inject_traceparent_to_metadata(metadata=None):
     """Inject W3C traceparent (and tracestate) into gRPC metadata tuple or list."""
     meta_list = list(metadata) if metadata else []
-    if (
-        not _opentelemetry_tracing.HAS_OPENTELEMETRY
-        or not _opentelemetry_tracing.enable_otel_traces
-    ):
+    if not _opentelemetry_tracing._is_otel_traces_enabled():
         return (
             tuple(meta_list)
             if isinstance(metadata, tuple) or metadata is None

packages/google-cloud-storage/google/cloud/storage/asyncio/async_grpc_client.py

@@ -76,6 +76,7 @@ def __init__(
         *,
         attempt_direct_path=True,
     ):
+        self._is_async_grpc_client = True
         self._bucket_metadata_cache = BucketMetadataCache(self)
         if isinstance(credentials, auth_credentials.AnonymousCredentials):
             if client_options is None or client_options.api_endpoint is None:

packages/google-cloud-storage/tests/unit/asyncio/test_zonal_observability.py

@@ -47,6 +47,10 @@ def exporter(monkeypatch):
     monkeypatch.delenv("DISABLE_GCS_PYTHON_CLIENT_OTEL_BUCKET_METADATA", raising=False)
     importlib.reload(_opentelemetry_tracing)
 
+    if hasattr(trace_api, "_TRACER_PROVIDER_SET_ONCE"):
+        trace_api._TRACER_PROVIDER_SET_ONCE._done = False
+    trace_api._TRACER_PROVIDER = None
+
     tracer_provider = TracerProvider()
     memory_exporter = InMemorySpanExporter()
     span_processor = export.SimpleSpanProcessor(memory_exporter)
@@ -55,6 +59,8 @@ def exporter(monkeypatch):
 
     yield memory_exporter
     memory_exporter.clear()
+    monkeypatch.setenv("ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES", "false")
+    importlib.reload(_opentelemetry_tracing)
 
 
 @pytest.fixture
@@ -75,8 +81,7 @@ def mock_client():
 
 
 def test_inject_traceparent_to_metadata_when_disabled(monkeypatch):
-    monkeypatch.setenv("ENABLE_GCS_PYTHON_CLIENT_OTEL_TRACES", "false")
-    importlib.reload(_opentelemetry_tracing)
+    monkeypatch.setattr(_opentelemetry_tracing, "enable_otel_traces", False)
 
     orig = (("x-goog-request-params", "bucket=foo"),)
     result = _utils.inject_traceparent_to_metadata(orig)
@@ -195,7 +200,8 @@ async def test_async_multi_range_downloader_spans(exporter, mock_client):
         mock_stream_cls.return_value = mock_stream
 
         mock_mux = mock.AsyncMock()
-        mock_mux.register.return_value = mock.AsyncMock()
+        mock_mux.register = mock.Mock(return_value=mock.AsyncMock())
+        mock_mux.unregister = mock.Mock()
         mock_mux_cls.return_value = mock_mux
 
         mock_retry_mgr = mock.AsyncMock()
@@ -250,3 +256,23 @@ async def test_bucket_metadata_cache_async_grpc_fetch(mock_client):
         "projects/987654321/buckets/new-zonal-bucket",
         "us-west1-b",
     )
+
+
+@pytest.mark.asyncio
+async def test_create_trace_span_helper_with_sync_mock_context_manager(mock_client):
+    """Verify fallback when _base_create_trace_span returns a synchronous context manager."""
+    from contextlib import contextmanager
+    from google.cloud.storage import _helpers
+
+    fake_span = object()
+
+    @contextmanager
+    def sync_cm(*args, **kwargs):
+        yield fake_span
+
+    with mock.patch.object(_helpers, "_base_create_trace_span", side_effect=sync_cm):
+        async with _helpers.create_trace_span_helper(
+            mock_client, "my-zonal-bucket", "Test.SyncFallback"
+        ) as span:
+            assert span is fake_span
+