feat(observability): [WIP] implement universal 4-path OpenTelemetry tracing

.pre-commit-config.yaml

@@ -14,6 +14,7 @@
 #
 # See https://pre-commit.com for more information
 # See https://pre-commit.com/hooks.html for more hooks
+exclude: '^packages/gapic-generator/tests/integration/goldens/'
 repos:
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v6.0.0

packages/gapic-generator/gapic/schema/mixins.py

@@ -19,50 +19,60 @@
         "DeleteOperation",
         request_type="operations_pb2.DeleteOperationRequest",
         response_type="None",
+        rpc_name="google.longrunning.Operations/DeleteOperation",
     ),
     "WaitOperation": wrappers.MixinMethod(
         "WaitOperation",
         request_type="operations_pb2.WaitOperationRequest",
         response_type="operations_pb2.Operation",
+        rpc_name="google.longrunning.Operations/WaitOperation",
     ),
     "ListOperations": wrappers.MixinMethod(
         "ListOperations",
         request_type="operations_pb2.ListOperationsRequest",
         response_type="operations_pb2.ListOperationsResponse",
+        rpc_name="google.longrunning.Operations/ListOperations",
     ),
     "CancelOperation": wrappers.MixinMethod(
         "CancelOperation",
         request_type="operations_pb2.CancelOperationRequest",
         response_type="None",
+        rpc_name="google.longrunning.Operations/CancelOperation",
     ),
     "GetOperation": wrappers.MixinMethod(
         "GetOperation",
         request_type="operations_pb2.GetOperationRequest",
         response_type="operations_pb2.Operation",
+        rpc_name="google.longrunning.Operations/GetOperation",
     ),
     "TestIamPermissions": wrappers.MixinMethod(
         "TestIamPermissions",
         request_type="iam_policy_pb2.TestIamPermissionsRequest",
         response_type="iam_policy_pb2.TestIamPermissionsResponse",
+        rpc_name="google.iam.v1.IAMPolicy/TestIamPermissions",
     ),
     "GetIamPolicy": wrappers.MixinMethod(
         "GetIamPolicy",
         request_type="iam_policy_pb2.GetIamPolicyRequest",
         response_type="policy_pb2.Policy",
+        rpc_name="google.iam.v1.IAMPolicy/GetIamPolicy",
     ),
     "SetIamPolicy": wrappers.MixinMethod(
         "SetIamPolicy",
         request_type="iam_policy_pb2.SetIamPolicyRequest",
         response_type="policy_pb2.Policy",
+        rpc_name="google.iam.v1.IAMPolicy/SetIamPolicy",
     ),
     "ListLocations": wrappers.MixinMethod(
         "ListLocations",
         request_type="locations_pb2.ListLocationsRequest",
         response_type="locations_pb2.ListLocationsResponse",
+        rpc_name="google.cloud.location.Locations/ListLocations",
     ),
     "GetLocation": wrappers.MixinMethod(
         "GetLocation",
         request_type="locations_pb2.GetLocationRequest",
         response_type="locations_pb2.Location",
+        rpc_name="google.cloud.location.Locations/GetLocation",
     ),
 }

packages/gapic-generator/gapic/schema/wrappers.py

@@ -1463,6 +1463,8 @@ class MixinMethod:
     name: str
     request_type: str
     response_type: str
+    rpc_name: str = ""
+
 
 
 @dataclasses.dataclass(frozen=True)

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/_compat.py.j2

@@ -5,14 +5,15 @@
 """A compatibility module for older versions of google-api-core."""
 {% set has_auto_populated_fields = api.all_method_settings.values()|map(attribute="auto_populated_fields", default=[])|select|list %}
 {# TODO(https://github.com/googleapis/google-cloud-python/issues/17813):
-Clean up this file/functions when the minimum supported version of 
+Clean up this file/functions when the minimum supported version of
 google-api-core has the functions in `_compat.py.j2`. #}
-{# TODO(https://github.com/googleapis/google-cloud-python/issues/17884): 
-Add conditional logic to check if static code exists in google-api-core and use it from there, 
+{# TODO(https://github.com/googleapis/google-cloud-python/issues/17884):
+Add conditional logic to check if static code exists in google-api-core and use it from there,
 falling back to the local implementation if not present. #}
-{# TODO(https://github.com/googleapis/google-cloud-python/issues/17883): 
+{# TODO(https://github.com/googleapis/google-cloud-python/issues/17883):
 Backfill compatibility functions being removed from the client layer. #}
 
+import contextlib
 import os
 import json
 {% if has_auto_populated_fields %}
@@ -34,9 +35,30 @@ from google.auth.exceptions import MutualTLSChannelError
 from google.protobuf import json_format
 from urllib.parse import urlparse, urlunparse
 
+# The _observability module was introduced in google-api-core 2.36.0+.
+# On older versions of google-api-core or when type-checking against them,
+# mypy may flag attr-defined or assignment errors when fallback to None occurs.
+try:
+    from google.api_core import _observability  # type: ignore[attr-defined]
+except ImportError:
+    _observability = None  # type: ignore[assignment]
+
+if _observability is not None and hasattr(_observability, "trace_http_request"):
+    trace_http_request = _observability.trace_http_request
+else:
+    @contextlib.contextmanager
+    def trace_http_request(*args: Any, **kwargs: Any):
+        yield None
+
+if _observability is not None and hasattr(_observability, "record_http_response"):
+    record_http_response = _observability.record_http_response
+else:
+    def record_http_response(span: Any, response: Any) -> None:
+        pass
+
 try:
     # note: `#type: ignore` is added because the return type for `should_use_client_cert`
-    # is different than that of the fallback implementation below. This will be removed once 
+    # is different than that of the fallback implementation below. This will be removed once
     # we bump the minimum supported version of google-auth.
     from google.auth.transport.mtls import should_use_client_cert  # type: ignore
 except ImportError:  # pragma: NO COVER

packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/_shared_macros.j2

@@ -159,30 +159,39 @@ def _get_http_options():
     session,
     timeout,
     transcoded_request,
-    body=None):
-
+    body=None,
+    client_options=None):
+    """Execute the HTTP request over the transport session with
+    OpenTelemetry tracing and metadata propagation."""
+
     uri = transcoded_request['uri']
     method = transcoded_request['method']
     headers = dict(metadata)
     headers['Content-Type'] = 'application/json'
-    response = {{ await_prefix }}getattr(session, method)(
-        "{host}{uri}".format(host=host, uri=uri),
-        timeout=timeout,
+    url = "{host}{uri}".format(host=host, uri=uri)
+
+    with trace_http_request(
+        client_options=client_options,
+        method=method,
+        url=url,
+        url_template=uri,
         headers=headers,
-        params=rest_helpers.flatten_query_params(query_params, strict=True),
-        {% if body_spec %}
-        data=body,
-        {% endif %}
-        {% if not is_async and is_streaming_method %}
-        {# NOTE: The underlying `requests` library used for making a sync request
-         # requires us to set `stream=True` to avoid loading the entire response
-         # into memory at once. For an async request, given its nature where it
-         # reads data chunk by chunk, this is not required.
-        #}
-        stream=True,
-        {% endif %}
+        body=body,
+    ) as span:
+        response = {{ await_prefix }}getattr(session, method)(
+            url,
+            timeout=timeout,
+            headers=headers,
+            params=rest_helpers.flatten_query_params(query_params, strict=True),
+            {% if body_spec %}
+            data=body,
+            {% endif %}
+            {% if not is_async and is_streaming_method %}
+            stream=True,
+            {% endif %}
         )
-    return response
+        record_http_response(span, response)
+        return response
 {% endmacro %}
 
 {# rest_call_method_common includes the common code for a rest __call__ method to be
@@ -240,13 +249,24 @@ def _get_http_options():
         )
 
     # Send the request
-    response = {{ await_prefix }}{{ async_class_prefix }}{{ service_name }}RestTransport._{{method_name}}._get_response(self._host, metadata, query_params, self._session, timeout, transcoded_request{% if body_spec %}, body{% endif %})
+    response = {{ await_prefix }}{{ async_class_prefix }}{{ service_name }}RestTransport._{{method_name}}._get_response(
+        self._host,
+        metadata,
+        query_params,
+        self._session,
+        timeout,
+        transcoded_request,
+        {% if body_spec %}
+        body,
+        {% endif %}
+        client_options=getattr(self, "_client_options", None),
+    )
 
     # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
     # subclass.
     if response.status_code >= 400:
-        {# Note: format_http_response_error takes in more parameters than from_http_response and the 
-        latter only supports a response of type requests.Response. 
+        {# Note: format_http_response_error takes in more parameters than from_http_response and the
+        latter only supports a response of type requests.Response.
         TODO: Clean up the sync response error handling and use format_http_response_error.
         See issue: https://github.com/googleapis/gapic-generator-python/issues/2116. #}
         {% if is_async %}
@@ -327,6 +347,10 @@ def _prep_wrapped_messages(self, client_info):
             {% endif %}
             default_timeout={{ method.timeout }},
             client_info=client_info,
+            method_name="{{ '.'.join(method.meta.address.package) }}.{{ service.name }}/{{ method.name }}",
+            {% if method.client_streaming or method.server_streaming %}
+            is_streaming=True,
+            {% endif %}
         ),
         {% endfor %}{# service.methods.values() #}
         {% for method_name in api.mixin_api_methods.keys() %}
@@ -337,26 +361,12 @@ def _prep_wrapped_messages(self, client_info):
             self.{{ method_name|snake_case }},
             default_timeout=None,
             client_info=client_info,
+            method_name="{{ api.mixin_api_signatures[method_name].rpc_name }}",
         ),
         {% endfor %}{# method_name in api.mixin_api_methods.keys() #}
     }
 {% endmacro %}
 
-{# TODO: This helper logic to check whether `kind` needs to be configured in wrap_method
-can be removed once we require the correct version of the google-api-core dependency to
-avoid having a gRPC code path in an async REST call.
-See related issue: https://github.com/googleapis/python-api-core/issues/661.
-In the meantime, if an older version of the dependency is installed (which has a wrap_method with
-no kind parameter), then an async gRPC call will work correctly and async REST transport
-will not be available as a transport.
-See related issue: https://github.com/googleapis/gapic-generator-python/issues/2119. #}
-{% macro wrap_async_method_macro() %}
-def _wrap_method(self, func, *args, **kwargs):
-    {# TODO: Remove `pragma: NO COVER` once https://github.com/googleapis/python-api-core/pull/688 is merged. #}
-    if self._wrap_with_kind:  # pragma: NO COVER
-        kwargs["kind"] = self.kind
-    return gapic_v1.method_async.wrap_method(func, *args, **kwargs)
-{% endmacro %}
 
 {# `create_interceptor_class` generates an Interceptor class for
  # synchronous and asynchronous rest transports
@@ -437,11 +447,11 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
 
         Override in a subclass to read or manipulate the response or metadata after it
         is returned by the {{ service.name }} server but before it is returned to user code.
-        
+
         We recommend only using this `post_{{ method.name|snake_case }}_with_metadata`
         interceptor in new development instead of the `post_{{ method.name|snake_case }}` interceptor.
         When both interceptors are used, this `post_{{ method.name|snake_case }}_with_metadata` interceptor runs after the
-        `post_{{ method.name|snake_case }}` interceptor. The (possibly modified) response returned by 
+        `post_{{ method.name|snake_case }}` interceptor. The (possibly modified) response returned by
         `post_{{ method.name|snake_case }}` will be passed to
         `post_{{ method.name|snake_case }}_with_metadata`.
         """
@@ -487,15 +497,15 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
 
 @property
 def {{ name|snake_case }}(self):
-    return self.{{ name|make_private }}(self._session, self._host, self._interceptor) # type: ignore
+    return self.{{ name|make_private }}(self._session, self._host, self._interceptor, getattr(self, "_client_options", None)) # type: ignore
 
 class {{ name|make_private }}(_Base{{ service.name }}RestTransport._Base{{name}}, {{ async_method_name_prefix }}{{service.name}}RestStub):
     def __hash__(self):
         return hash("{{ async_method_name_prefix }}{{ service.name }}RestTransport.{{ name }}")
 
     {% set body_spec = api.mixin_http_options["{}".format(name)][0].body %}
     {{ response_method(body_spec, is_async=is_async, is_streaming_method=None) | indent(4) }}
-    
+
     {{ async_prefix }}def __call__(self,
         request: {{ sig.request_type }}, *,
         retry: OptionalRetry=gapic_v1.method.DEFAULT,
@@ -521,7 +531,7 @@ class {{ name|make_private }}(_Base{{ service.name }}RestTransport._Base{{name}}
         {% endif %}
         """
         {{ rest_call_method_common(body_spec, name, service, is_async)|indent(4) }}
-        
+
         {% if sig.response_type == "None" %}
         return {{ await_prefix }}self._interceptor.post_{{ name|snake_case }}(None)
         {% else %}
@@ -543,7 +553,7 @@ class {{ name|make_private }}(_Base{{ service.name }}RestTransport._Base{{name}}
                 "payload": response_payload,
                 "headers":  dict(response.headers),
                 "status": response.status_code,
-            }        
+            }
             _LOGGER.debug(
                 "Received response for {{ service.meta.address.proto_package_versioned }}.{{ service.async_client_name }}.{{ name }}",
                 extra = {
@@ -565,7 +575,7 @@ class {{ name|make_private }}(_Base{{ service.name }}RestTransport._Base{{name}}
 
 {% macro client_method_metadata_default_value() %}(){% endmacro %}
 
-{% macro client_method_metadata_argument_doc() %}metadata ({{ client_method_metadata_type() }}): Key/value pairs which should be 
+{% macro client_method_metadata_argument_doc() %}metadata ({{ client_method_metadata_type() }}): Key/value pairs which should be
             sent along with the request as metadata. Normally, each value must be of type `str`,
             but for metadata keys ending with the suffix `-bin`, the corresponding values must
             be of type `bytes`.{% endmacro %}