Merge branch 'main' into add-support-for-protobuf-5-x

parthea · web-flow · commit 583b8ce5b92d · 2024-05-30T13:24:29.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,13 @@
 
 [1]: https://pypi.org/project/google-api-core/#history
 
+## [2.19.0](https://github.com/googleapis/python-api-core/compare/v2.18.0...v2.19.0) (2024-04-29)
+
+
+### Features
+
+* Add google.api_core.version_header ([#638](https://github.com/googleapis/python-api-core/issues/638)) ([a7b53e9](https://github.com/googleapis/python-api-core/commit/a7b53e9e9a7deb88baf92a2827958429e3677069))
+
 ## [2.18.0](https://github.com/googleapis/python-api-core/compare/v2.17.1...v2.18.0) (2024-03-20)
 
 
diff --git a/google/api_core/gapic_v1/method.py b/google/api_core/gapic_v1/method.py
@@ -42,24 +42,6 @@ class _MethodDefault(enum.Enum):
 so the default should be used."""
 
 
-def _is_not_none_or_false(value):
-    return value is not None and value is not False
-
-
-def _apply_decorators(func, decorators):
-    """Apply a list of decorators to a given function.
-
-    ``decorators`` may contain items that are ``None`` or ``False`` which will
-    be ignored.
-    """
-    filtered_decorators = filter(_is_not_none_or_false, reversed(decorators))
-
-    for decorator in filtered_decorators:
-        func = decorator(func)
-
-    return func
-
-
 class _GapicCallable(object):
     """Callable that applies retry, timeout, and metadata logic.
 
@@ -91,6 +73,8 @@ def __init__(
     ):
         self._target = target
         self._retry = retry
+        if isinstance(timeout, (int, float)):
+            timeout = TimeToDeadlineTimeout(timeout=timeout)
         self._timeout = timeout
         self._compression = compression
         self._metadata = metadata
@@ -100,35 +84,42 @@ def __call__(
     ):
         """Invoke the low-level RPC with retry, timeout, compression, and metadata."""
 
-        if retry is DEFAULT:
-            retry = self._retry
-
-        if timeout is DEFAULT:
-            timeout = self._timeout
-
         if compression is DEFAULT:
             compression = self._compression
-
-        if isinstance(timeout, (int, float)):
-            timeout = TimeToDeadlineTimeout(timeout=timeout)
-
-        # Apply all applicable decorators.
-        wrapped_func = _apply_decorators(self._target, [retry, timeout])
+        if compression is not None:
+            kwargs["compression"] = compression
 
         # Add the user agent metadata to the call.
         if self._metadata is not None:
-            metadata = kwargs.get("metadata", [])
-            # Due to the nature of invocation, None should be treated the same
-            # as not specified.
-            if metadata is None:
-                metadata = []
-            metadata = list(metadata)
-            metadata.extend(self._metadata)
-            kwargs["metadata"] = metadata
-        if self._compression is not None:
-            kwargs["compression"] = compression
+            try:
+                # attempt to concatenate default metadata with user-provided metadata
+                kwargs["metadata"] = (*kwargs["metadata"], *self._metadata)
+            except (KeyError, TypeError):
+                # if metadata is not provided, use just the default metadata
+                kwargs["metadata"] = self._metadata
+
+        call = self._build_wrapped_call(timeout, retry)
+        return call(*args, **kwargs)
+
+    @functools.lru_cache(maxsize=4)
+    def _build_wrapped_call(self, timeout, retry):
+        """
+        Build a wrapped callable that applies retry, timeout, and metadata logic.
+        """
+        wrapped_func = self._target
+        if timeout is DEFAULT:
+            timeout = self._timeout
+        elif isinstance(timeout, (int, float)):
+            timeout = TimeToDeadlineTimeout(timeout=timeout)
+        if timeout is not None:
+            wrapped_func = timeout(wrapped_func)
+
+        if retry is DEFAULT:
+            retry = self._retry
+        if retry is not None:
+            wrapped_func = retry(wrapped_func)
 
-        return wrapped_func(*args, **kwargs)
+        return wrapped_func
 
 
 def wrap_method(
diff --git a/google/api_core/grpc_helpers_async.py b/google/api_core/grpc_helpers_async.py
@@ -159,7 +159,6 @@ class _WrappedStreamStreamCall(
 
 def _wrap_unary_errors(callable_):
     """Map errors for Unary-Unary async callables."""
-    grpc_helpers._patch_callable_name(callable_)
 
     @functools.wraps(callable_)
     def error_remapped_callable(*args, **kwargs):
@@ -169,23 +168,13 @@ def error_remapped_callable(*args, **kwargs):
     return error_remapped_callable
 
 
-def _wrap_stream_errors(callable_):
+def _wrap_stream_errors(callable_, wrapper_type):
     """Map errors for streaming RPC async callables."""
-    grpc_helpers._patch_callable_name(callable_)
 
     @functools.wraps(callable_)
     async def error_remapped_callable(*args, **kwargs):
         call = callable_(*args, **kwargs)
-
-        if isinstance(call, aio.UnaryStreamCall):
-            call = _WrappedUnaryStreamCall().with_call(call)
-        elif isinstance(call, aio.StreamUnaryCall):
-            call = _WrappedStreamUnaryCall().with_call(call)
-        elif isinstance(call, aio.StreamStreamCall):
-            call = _WrappedStreamStreamCall().with_call(call)
-        else:
-            raise TypeError("Unexpected type of call %s" % type(call))
-
+        call = wrapper_type().with_call(call)
         await call.wait_for_connection()
         return call
 
@@ -207,10 +196,17 @@ def wrap_errors(callable_):
 
     Returns: Callable: The wrapped gRPC callable.
     """
+    grpc_helpers._patch_callable_name(callable_)
     if isinstance(callable_, aio.UnaryUnaryMultiCallable):
         return _wrap_unary_errors(callable_)
+    elif isinstance(callable_, aio.UnaryStreamMultiCallable):
+        return _wrap_stream_errors(callable_, _WrappedUnaryStreamCall)
+    elif isinstance(callable_, aio.StreamUnaryMultiCallable):
+        return _wrap_stream_errors(callable_, _WrappedStreamUnaryCall)
+    elif isinstance(callable_, aio.StreamStreamMultiCallable):
+        return _wrap_stream_errors(callable_, _WrappedStreamStreamCall)
     else:
-        return _wrap_stream_errors(callable_)
+        raise TypeError("Unexpected type of callable: {}".format(type(callable_)))
 
 
 def create_channel(
diff --git a/google/api_core/version.py b/google/api_core/version.py
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "2.18.0"
+__version__ = "2.19.0"
diff --git a/google/api_core/version_header.py b/google/api_core/version_header.py
@@ -0,0 +1,29 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+API_VERSION_METADATA_KEY = "x-goog-api-version"
+
+
+def to_api_version_header(version_identifier):
+    """Returns data for the API Version header for the given `version_identifier`.
+
+    Args:
+        version_identifier (str): The version identifier to be used in the
+            tuple returned.
+
+    Returns:
+        Tuple(str, str): A tuple containing the API Version metadata key and
+            value.
+    """
+    return (API_VERSION_METADATA_KEY, version_identifier)
diff --git a/tests/asyncio/test_grpc_helpers_async.py b/tests/asyncio/test_grpc_helpers_async.py
@@ -97,12 +97,40 @@ async def test_common_methods_in_wrapped_call():
     assert mock_call.wait_for_connection.call_count == 1
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "callable_type,expected_wrapper_type",
+    [
+        (grpc.aio.UnaryStreamMultiCallable, grpc_helpers_async._WrappedUnaryStreamCall),
+        (grpc.aio.StreamUnaryMultiCallable, grpc_helpers_async._WrappedStreamUnaryCall),
+        (
+            grpc.aio.StreamStreamMultiCallable,
+            grpc_helpers_async._WrappedStreamStreamCall,
+        ),
+    ],
+)
+async def test_wrap_errors_w_stream_type(callable_type, expected_wrapper_type):
+    class ConcreteMulticallable(callable_type):
+        def __call__(self, *args, **kwargs):
+            raise NotImplementedError("Should not be called")
+
+    with mock.patch.object(
+        grpc_helpers_async, "_wrap_stream_errors"
+    ) as wrap_stream_errors:
+        callable_ = ConcreteMulticallable()
+        grpc_helpers_async.wrap_errors(callable_)
+        assert wrap_stream_errors.call_count == 1
+        wrap_stream_errors.assert_called_once_with(callable_, expected_wrapper_type)
+
+
 @pytest.mark.asyncio
 async def test_wrap_stream_errors_unary_stream():
     mock_call = mock.Mock(aio.UnaryStreamCall, autospec=True)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedUnaryStreamCall
+    )
 
     await wrapped_callable(1, 2, three="four")
     multicallable.assert_called_once_with(1, 2, three="four")
@@ -114,7 +142,9 @@ async def test_wrap_stream_errors_stream_unary():
     mock_call = mock.Mock(aio.StreamUnaryCall, autospec=True)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamUnaryCall
+    )
 
     await wrapped_callable(1, 2, three="four")
     multicallable.assert_called_once_with(1, 2, three="four")
@@ -126,22 +156,26 @@ async def test_wrap_stream_errors_stream_stream():
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
 
     await wrapped_callable(1, 2, three="four")
     multicallable.assert_called_once_with(1, 2, three="four")
     assert mock_call.wait_for_connection.call_count == 1
 
 
 @pytest.mark.asyncio
-async def test_wrap_stream_errors_type_error():
+async def test_wrap_errors_type_error():
+    """
+    If wrap_errors is called with an unexpected type, it should raise a TypeError.
+    """
     mock_call = mock.Mock()
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
-
-    with pytest.raises(TypeError):
-        await wrapped_callable()
+    with pytest.raises(TypeError) as exc:
+        grpc_helpers_async.wrap_errors(multicallable)
+    assert "Unexpected type" in str(exc.value)
 
 
 @pytest.mark.asyncio
@@ -151,7 +185,9 @@ async def test_wrap_stream_errors_raised():
     mock_call.wait_for_connection = mock.AsyncMock(side_effect=[grpc_error])
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
 
     with pytest.raises(exceptions.InvalidArgument):
         await wrapped_callable()
@@ -166,7 +202,9 @@ async def test_wrap_stream_errors_read():
     mock_call.read = mock.AsyncMock(side_effect=grpc_error)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
 
     wrapped_call = await wrapped_callable(1, 2, three="four")
     multicallable.assert_called_once_with(1, 2, three="four")
@@ -189,7 +227,9 @@ async def test_wrap_stream_errors_aiter():
     mock_call.__aiter__ = mock.Mock(return_value=mocked_aiter)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
     wrapped_call = await wrapped_callable()
 
     with pytest.raises(exceptions.InvalidArgument) as exc_info:
@@ -210,7 +250,9 @@ async def test_wrap_stream_errors_aiter_non_rpc_error():
     mock_call.__aiter__ = mock.Mock(return_value=mocked_aiter)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
     wrapped_call = await wrapped_callable()
 
     with pytest.raises(TypeError) as exc_info:
@@ -224,7 +266,9 @@ async def test_wrap_stream_errors_aiter_called_multiple_times():
     mock_call = mock.Mock(aio.StreamStreamCall, autospec=True)
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
     wrapped_call = await wrapped_callable()
 
     assert wrapped_call.__aiter__() == wrapped_call.__aiter__()
@@ -239,7 +283,9 @@ async def test_wrap_stream_errors_write():
     mock_call.done_writing = mock.AsyncMock(side_effect=[None, grpc_error])
     multicallable = mock.Mock(return_value=mock_call)
 
-    wrapped_callable = grpc_helpers_async._wrap_stream_errors(multicallable)
+    wrapped_callable = grpc_helpers_async._wrap_stream_errors(
+        multicallable, grpc_helpers_async._WrappedStreamStreamCall
+    )
 
     wrapped_call = await wrapped_callable()
 
@@ -295,7 +341,9 @@ def test_wrap_errors_streaming(wrap_stream_errors):
     result = grpc_helpers_async.wrap_errors(callable_)
 
     assert result == wrap_stream_errors.return_value
-    wrap_stream_errors.assert_called_once_with(callable_)
+    wrap_stream_errors.assert_called_once_with(
+        callable_, grpc_helpers_async._WrappedUnaryStreamCall
+    )
 
 
 @pytest.mark.parametrize(
diff --git a/tests/unit/gapic/test_method.py b/tests/unit/gapic/test_method.py
@@ -222,3 +222,24 @@ def test_wrap_method_with_call_not_supported():
     with pytest.raises(ValueError) as exc_info:
         google.api_core.gapic_v1.method.wrap_method(method, with_call=True)
     assert "with_call=True is only supported for unary calls" in str(exc_info.value)
+
+
+def test_benchmark_gapic_call():
+    """
+    Ensure the __call__ method performance does not regress from expectations
+
+    __call__ builds a new wrapped function using passed-in timeout and retry, but
+    subsequent calls are cached
+
+    Note: The threshold has been tuned for the CI workers. Test may flake on
+    slower hardware
+    """
+    from google.api_core.gapic_v1.method import _GapicCallable
+    from google.api_core.retry import Retry
+    from timeit import timeit
+
+    gapic_callable = _GapicCallable(
+        lambda *a, **k: 1, retry=Retry(), timeout=1010, compression=False
+    )
+    avg_time = timeit(lambda: gapic_callable(), number=10_000)
+    assert avg_time < 0.4
diff --git a/tests/unit/test_version_header.py b/tests/unit/test_version_header.py
