Skip to content

Commit 4a0bac3

Browse files
chalmerlowechalmerlowe
authored
refactor(observability): simplify AsyncGapicCallable inheritance and forward client_options in resumable stubs (#18596)
1 parent cf1a0e5 commit 4a0bac3

46 files changed

Lines changed: 4709 additions & 549 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/base.py.j2‎

Lines changed: 44 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -179,20 +179,34 @@ class {{ service.name }}Transport(abc.ABC):
179179
stripped for backward compatibility with older `google-api-core`
180180
versions.
181181
"""
182-
if _WRAP_METHOD_SUPPORTS_TRACING: # pragma: NO COVER
182+
if _WRAP_METHOD_SUPPORTS_TRACING:
183183
kwargs["client_options"] = self._client_options
184184
try:
185185
kind = self.kind
186-
except NotImplementedError: # pragma: NO COVER
186+
except NotImplementedError:
187187
kind = None
188-
if kind: # pragma: NO COVER
188+
if kind:
189189
kwargs["kind"] = kind
190190
return gapic_v1.method.wrap_method(func, *args, **kwargs)
191-
# The fallback below strips tracing-specific arguments when an older version
192-
# of google-api-core is installed (which does not accept client_options, etc.).
193-
for k in ["client_options", "method_name", "is_streaming", "kind"]: # pragma: NO COVER
194-
kwargs.pop(k, None) # pragma: NO COVER
195-
return gapic_v1.method.wrap_method(func, *args, **kwargs) # pragma: NO COVER
191+
# Fallback for older runtime versions of google-api-core:
192+
# - Era 1 (< 2.29.0): Neither tracing args nor `kind` are supported.
193+
# - Era 2 (>= 2.29.0, < 2.36.0): `kind` is supported (prevents REST from
194+
# falling back to gRPC wrapping), but tracing args are not.
195+
# (These fallbacks can be removed once google-api-core < 2.36.0 is no longer supported.)
196+
for k in ["client_options", "method_name", "is_streaming"]:
197+
kwargs.pop(k, None)
198+
199+
try:
200+
kind = self.kind
201+
except NotImplementedError:
202+
kind = None
203+
204+
if _WRAP_METHOD_SUPPORTS_KIND and kind:
205+
kwargs["kind"] = kind
206+
else:
207+
kwargs.pop("kind", None)
208+
209+
return gapic_v1.method.wrap_method(func, *args, **kwargs)
196210

197211
def _wrap_async_method(self, func, *args, **kwargs):
198212
"""Wrap an async RPC method with common client-level features.
@@ -204,20 +218,34 @@ class {{ service.name }}Transport(abc.ABC):
204218
stripped for backward compatibility with older `google-api-core`
205219
versions.
206220
"""
207-
if _ASYNC_WRAP_METHOD_SUPPORTS_TRACING: # pragma: NO COVER
221+
if _ASYNC_WRAP_METHOD_SUPPORTS_TRACING:
208222
kwargs["client_options"] = self._client_options
209223
try:
210224
kind = self.kind
211-
except NotImplementedError: # pragma: NO COVER
225+
except NotImplementedError:
212226
kind = None
213-
if kind: # pragma: NO COVER
227+
if kind:
214228
kwargs["kind"] = kind
215229
return gapic_v1.method_async.wrap_method(func, *args, **kwargs)
216-
# The fallback below strips tracing-specific arguments when an older version
217-
# of google-api-core is installed (which does not accept client_options, etc.).
218-
for k in ["client_options", "method_name", "is_streaming", "kind"]: # pragma: NO COVER
219-
kwargs.pop(k, None) # pragma: NO COVER
220-
return gapic_v1.method_async.wrap_method(func, *args, **kwargs) # pragma: NO COVER
230+
# Fallback for older runtime versions of google-api-core:
231+
# - Era 1 (< 2.29.0): Neither tracing args nor `kind` are supported.
232+
# - Era 2 (>= 2.29.0, < 2.36.0): `kind` is supported (prevents REST from
233+
# falling back to gRPC wrapping), but tracing args are not.
234+
# (These fallbacks can be removed once google-api-core < 2.36.0 is no longer supported.)
235+
for k in ["client_options", "method_name", "is_streaming"]:
236+
kwargs.pop(k, None)
237+
238+
try:
239+
kind = self.kind
240+
except NotImplementedError:
241+
kind = None
242+
243+
if _ASYNC_WRAP_METHOD_SUPPORTS_KIND and kind:
244+
kwargs["kind"] = kind
245+
else:
246+
kwargs.pop("kind", None)
247+
248+
return gapic_v1.method_async.wrap_method(func, *args, **kwargs)
221249

222250
def _prep_wrapped_messages(self, client_info):
223251
"""Precompute and cache wrapped methods for RPC dispatch."""

‎packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/grpc.py.j2‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -426,6 +426,7 @@ class {{ service.name }}GrpcTransport({{ service.name }}Transport):
426426
credentials=transport._credentials,
427427
client_info=transport._client_info,
428428
client_cert_source_for_mtls=transport._client_cert_source_for_mtls,
429+
client_options=getattr(transport, "_client_options", None),
429430
)
430431
return transport._rest_transport.{{ method.transport_safe_name|snake_case }}(*args, **kwargs)
431432
self._stubs['{{ method.transport_safe_name|snake_case }}'] = _RestStub()

‎packages/gapic-generator/gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/grpc_asyncio.py.j2‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -435,6 +435,7 @@ class {{ service.grpc_asyncio_transport_name }}({{ service.name }}Transport):
435435
host=transport._host,
436436
credentials=transport._credentials,
437437
client_info=transport._client_info,
438+
client_options=getattr(transport, "_client_options", None),
438439
)
439440
return transport._rest_transport.{{ method.transport_safe_name|snake_case }}(*args, **kwargs)
440441
self._stubs['{{ method.transport_safe_name|snake_case }}'] = _AsyncRestStub()

‎packages/gapic-generator/gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2‎

Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1817,6 +1817,7 @@ async def test_{{ service.name|snake_case }}_{{ method.name|snake_case }}_grpc_a
18171817
host=transport._host,
18181818
credentials=transport._credentials,
18191819
client_info=transport._client_info,
1820+
client_options=getattr(transport, "_client_options", None),
18201821
)
18211822
assert mock_rest_transport.{{ method.transport_safe_name|snake_case }}.call_count == 2
18221823
await transport.close()
@@ -2475,4 +2476,200 @@ def test_api_key_credentials(client_class, transport_class):
24752476
api_audience=None,
24762477
)
24772478

2479+
2480+
def test_{{ service.name|snake_case }}_base_transport_wrap_method_modern_api_core(monkeypatch):
2481+
monkeypatch.setattr(transports.base, "_WRAP_METHOD_SUPPORTS_TRACING", True)
2482+
# 1. Concrete transport: passes client_options=None and kind="{{ opts.transport[0] }}"
2483+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method") as mock_wrap:
2484+
transport = transports.{{ service.name }}{{ opts.transport[0].capitalize() }}Transport(
2485+
credentials=ga_credentials.AnonymousCredentials()
2486+
)
2487+
func = mock.MagicMock()
2488+
transport._wrap_method(func)
2489+
assert mock_wrap.call_args.kwargs["kind"] == "{{ opts.transport[0] }}"
2490+
assert "client_options" in mock_wrap.call_args.kwargs
2491+
2492+
# 2. Base transport: kind raises NotImplementedError and is cleanly omitted
2493+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method") as mock_wrap:
2494+
base_transport = transports.{{ service.name }}Transport(
2495+
credentials=ga_credentials.AnonymousCredentials()
2496+
)
2497+
func = mock.MagicMock()
2498+
base_transport._wrap_method(func)
2499+
assert "kind" not in mock_wrap.call_args.kwargs
2500+
assert "client_options" in mock_wrap.call_args.kwargs
2501+
2502+
2503+
def test_{{ service.name|snake_case }}_base_transport_wrap_method_older_api_core_fallbacks(monkeypatch):
2504+
monkeypatch.setattr(transports.base, "_WRAP_METHOD_SUPPORTS_TRACING", False)
2505+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method"):
2506+
transport = transports.{{ service.name }}{{ opts.transport[0].capitalize() }}Transport(
2507+
credentials=ga_credentials.AnonymousCredentials()
2508+
)
2509+
2510+
# Era 1: Ancient google-api-core (< 2.29.0: can be removed once google-api-core < 2.29.0 is no longer supported)
2511+
monkeypatch.setattr(transports.base, "_WRAP_METHOD_SUPPORTS_KIND", False)
2512+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method") as mock_wrap:
2513+
func = mock.MagicMock()
2514+
transport._wrap_method(
2515+
func,
2516+
client_options="fake_options",
2517+
method_name="fake_method",
2518+
is_streaming=False,
2519+
)
2520+
kwargs = mock_wrap.call_args.kwargs
2521+
assert "kind" not in kwargs
2522+
assert "client_options" not in kwargs
2523+
assert "method_name" not in kwargs
2524+
assert "is_streaming" not in kwargs
2525+
2526+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method") as mock_wrap:
2527+
base_transport = transports.{{ service.name }}Transport(
2528+
credentials=ga_credentials.AnonymousCredentials()
2529+
)
2530+
func = mock.MagicMock()
2531+
base_transport._wrap_method(
2532+
func,
2533+
client_options="fake_options",
2534+
method_name="fake_method",
2535+
is_streaming=False,
2536+
)
2537+
kwargs = mock_wrap.call_args.kwargs
2538+
assert "kind" not in kwargs
2539+
assert "client_options" not in kwargs
2540+
assert "method_name" not in kwargs
2541+
assert "is_streaming" not in kwargs
2542+
2543+
# Era 2: Intermediate google-api-core (>= 2.29.0, < 2.36.0: can be removed once google-api-core < 2.36.0 is no longer supported)
2544+
monkeypatch.setattr(transports.base, "_WRAP_METHOD_SUPPORTS_KIND", True)
2545+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method") as mock_wrap:
2546+
func = mock.MagicMock()
2547+
transport._wrap_method(
2548+
func,
2549+
client_options="fake_options",
2550+
method_name="fake_method",
2551+
is_streaming=False,
2552+
)
2553+
kwargs = mock_wrap.call_args.kwargs
2554+
assert kwargs.get("kind") == "{{ opts.transport[0] }}"
2555+
assert "client_options" not in kwargs
2556+
assert "method_name" not in kwargs
2557+
assert "is_streaming" not in kwargs
2558+
2559+
with mock.patch.object(transports.base.gapic_v1.method, "wrap_method") as mock_wrap:
2560+
base_transport = transports.{{ service.name }}Transport(
2561+
credentials=ga_credentials.AnonymousCredentials()
2562+
)
2563+
func = mock.MagicMock()
2564+
base_transport._wrap_method(
2565+
func,
2566+
client_options="fake_options",
2567+
method_name="fake_method",
2568+
is_streaming=False,
2569+
)
2570+
kwargs = mock_wrap.call_args.kwargs
2571+
assert "kind" not in kwargs
2572+
assert "client_options" not in kwargs
2573+
assert "method_name" not in kwargs
2574+
assert "is_streaming" not in kwargs
2575+
2576+
2577+
{% if 'grpc' in opts.transport %}
2578+
def test_{{ service.name|snake_case }}_base_transport_wrap_async_method_modern_api_core(monkeypatch):
2579+
monkeypatch.setattr(transports.base, "_ASYNC_WRAP_METHOD_SUPPORTS_TRACING", True)
2580+
# 1. Concrete transport: passes client_options=None and kind="grpc_asyncio"
2581+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method") as mock_wrap:
2582+
transport = transports.{{ service.grpc_asyncio_transport_name }}(
2583+
credentials=ga_credentials.AnonymousCredentials()
2584+
)
2585+
func = mock.MagicMock()
2586+
transport._wrap_async_method(func)
2587+
assert mock_wrap.call_args.kwargs["kind"] == "grpc_asyncio"
2588+
assert "client_options" in mock_wrap.call_args.kwargs
2589+
2590+
# 2. Base transport: kind raises NotImplementedError and is cleanly omitted
2591+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method") as mock_wrap:
2592+
base_transport = transports.{{ service.name }}Transport(
2593+
credentials=ga_credentials.AnonymousCredentials()
2594+
)
2595+
func = mock.MagicMock()
2596+
base_transport._wrap_async_method(func)
2597+
assert "kind" not in mock_wrap.call_args.kwargs
2598+
assert "client_options" in mock_wrap.call_args.kwargs
2599+
2600+
2601+
def test_{{ service.name|snake_case }}_base_transport_wrap_async_method_older_api_core_fallbacks(monkeypatch):
2602+
monkeypatch.setattr(transports.base, "_ASYNC_WRAP_METHOD_SUPPORTS_TRACING", False)
2603+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method"):
2604+
transport = transports.{{ service.grpc_asyncio_transport_name }}(
2605+
credentials=ga_credentials.AnonymousCredentials()
2606+
)
2607+
2608+
# Era 1: Ancient google-api-core (< 2.29.0: can be removed once google-api-core < 2.29.0 is no longer supported)
2609+
monkeypatch.setattr(transports.base, "_ASYNC_WRAP_METHOD_SUPPORTS_KIND", False)
2610+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method") as mock_wrap:
2611+
func = mock.MagicMock()
2612+
transport._wrap_async_method(
2613+
func,
2614+
client_options="fake_options",
2615+
method_name="fake_method",
2616+
is_streaming=False,
2617+
)
2618+
kwargs = mock_wrap.call_args.kwargs
2619+
assert "kind" not in kwargs
2620+
assert "client_options" not in kwargs
2621+
assert "method_name" not in kwargs
2622+
assert "is_streaming" not in kwargs
2623+
2624+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method") as mock_wrap:
2625+
base_transport = transports.{{ service.name }}Transport(
2626+
credentials=ga_credentials.AnonymousCredentials()
2627+
)
2628+
func = mock.MagicMock()
2629+
base_transport._wrap_async_method(
2630+
func,
2631+
client_options="fake_options",
2632+
method_name="fake_method",
2633+
is_streaming=False,
2634+
)
2635+
kwargs = mock_wrap.call_args.kwargs
2636+
assert "kind" not in kwargs
2637+
assert "client_options" not in kwargs
2638+
assert "method_name" not in kwargs
2639+
assert "is_streaming" not in kwargs
2640+
2641+
# Era 2: Intermediate google-api-core (>= 2.29.0, < 2.36.0: can be removed once google-api-core < 2.36.0 is no longer supported)
2642+
monkeypatch.setattr(transports.base, "_ASYNC_WRAP_METHOD_SUPPORTS_KIND", True)
2643+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method") as mock_wrap:
2644+
func = mock.MagicMock()
2645+
transport._wrap_async_method(
2646+
func,
2647+
client_options="fake_options",
2648+
method_name="fake_method",
2649+
is_streaming=False,
2650+
)
2651+
kwargs = mock_wrap.call_args.kwargs
2652+
assert kwargs.get("kind") == "grpc_asyncio"
2653+
assert "client_options" not in kwargs
2654+
assert "method_name" not in kwargs
2655+
assert "is_streaming" not in kwargs
2656+
2657+
with mock.patch.object(transports.base.gapic_v1.method_async, "wrap_method") as mock_wrap:
2658+
base_transport = transports.{{ service.name }}Transport(
2659+
credentials=ga_credentials.AnonymousCredentials()
2660+
)
2661+
func = mock.MagicMock()
2662+
base_transport._wrap_async_method(
2663+
func,
2664+
client_options="fake_options",
2665+
method_name="fake_method",
2666+
is_streaming=False,
2667+
)
2668+
kwargs = mock_wrap.call_args.kwargs
2669+
assert "kind" not in kwargs
2670+
assert "client_options" not in kwargs
2671+
assert "method_name" not in kwargs
2672+
assert "is_streaming" not in kwargs
2673+
{% endif %}
2674+
24782675
{% endblock %}

‎packages/gapic-generator/tests/integration/goldens/asset/google/cloud/asset_v1/services/asset_service/transports/base.py‎

Lines changed: 44 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -150,20 +150,34 @@ def _wrap_method(self, func, *args, **kwargs):
150150
stripped for backward compatibility with older `google-api-core`
151151
versions.
152152
"""
153-
if _WRAP_METHOD_SUPPORTS_TRACING: # pragma: NO COVER
153+
if _WRAP_METHOD_SUPPORTS_TRACING:
154154
kwargs["client_options"] = self._client_options
155155
try:
156156
kind = self.kind
157-
except NotImplementedError: # pragma: NO COVER
157+
except NotImplementedError:
158158
kind = None
159-
if kind: # pragma: NO COVER
159+
if kind:
160160
kwargs["kind"] = kind
161161
return gapic_v1.method.wrap_method(func, *args, **kwargs)
162-
# The fallback below strips tracing-specific arguments when an older version
163-
# of google-api-core is installed (which does not accept client_options, etc.).
164-
for k in ["client_options", "method_name", "is_streaming", "kind"]: # pragma: NO COVER
165-
kwargs.pop(k, None) # pragma: NO COVER
166-
return gapic_v1.method.wrap_method(func, *args, **kwargs) # pragma: NO COVER
162+
# Fallback for older runtime versions of google-api-core:
163+
# - Era 1 (< 2.29.0): Neither tracing args nor `kind` are supported.
164+
# - Era 2 (>= 2.29.0, < 2.36.0): `kind` is supported (prevents REST from
165+
# falling back to gRPC wrapping), but tracing args are not.
166+
# (These fallbacks can be removed once google-api-core < 2.36.0 is no longer supported.)
167+
for k in ["client_options", "method_name", "is_streaming"]:
168+
kwargs.pop(k, None)
169+
170+
try:
171+
kind = self.kind
172+
except NotImplementedError:
173+
kind = None
174+
175+
if _WRAP_METHOD_SUPPORTS_KIND and kind:
176+
kwargs["kind"] = kind
177+
else:
178+
kwargs.pop("kind", None)
179+
180+
return gapic_v1.method.wrap_method(func, *args, **kwargs)
167181

168182
def _wrap_async_method(self, func, *args, **kwargs):
169183
"""Wrap an async RPC method with common client-level features.
@@ -175,20 +189,34 @@ def _wrap_async_method(self, func, *args, **kwargs):
175189
stripped for backward compatibility with older `google-api-core`
176190
versions.
177191
"""
178-
if _ASYNC_WRAP_METHOD_SUPPORTS_TRACING: # pragma: NO COVER
192+
if _ASYNC_WRAP_METHOD_SUPPORTS_TRACING:
179193
kwargs["client_options"] = self._client_options
180194
try:
181195
kind = self.kind
182-
except NotImplementedError: # pragma: NO COVER
196+
except NotImplementedError:
183197
kind = None
184-
if kind: # pragma: NO COVER
198+
if kind:
185199
kwargs["kind"] = kind
186200
return gapic_v1.method_async.wrap_method(func, *args, **kwargs)
187-
# The fallback below strips tracing-specific arguments when an older version
188-
# of google-api-core is installed (which does not accept client_options, etc.).
189-
for k in ["client_options", "method_name", "is_streaming", "kind"]: # pragma: NO COVER
190-
kwargs.pop(k, None) # pragma: NO COVER
191-
return gapic_v1.method_async.wrap_method(func, *args, **kwargs) # pragma: NO COVER
201+
# Fallback for older runtime versions of google-api-core:
202+
# - Era 1 (< 2.29.0): Neither tracing args nor `kind` are supported.
203+
# - Era 2 (>= 2.29.0, < 2.36.0): `kind` is supported (prevents REST from
204+
# falling back to gRPC wrapping), but tracing args are not.
205+
# (These fallbacks can be removed once google-api-core < 2.36.0 is no longer supported.)
206+
for k in ["client_options", "method_name", "is_streaming"]:
207+
kwargs.pop(k, None)
208+
209+
try:
210+
kind = self.kind
211+
except NotImplementedError:
212+
kind = None
213+
214+
if _ASYNC_WRAP_METHOD_SUPPORTS_KIND and kind:
215+
kwargs["kind"] = kind
216+
else:
217+
kwargs.pop("kind", None)
218+
219+
return gapic_v1.method_async.wrap_method(func, *args, **kwargs)
192220

193221
def _prep_wrapped_messages(self, client_info):
194222
"""Precompute and cache wrapped methods for RPC dispatch."""

0 commit comments

Comments
 (0)