8000 feat: allow gRPC metadata to be passed to operations client methods · sushicw/python-api-core@3c6bbe8 · GitHub
[go: up one dir, main page]
Skip to content

Commit 3c6bbe8

Browse files
busunkim96busunkim96
committed
feat: allow gRPC metadata to be passed to operations client methods
1 parent c5fee89 commit 3c6bbe8

File tree

2 files changed

+33
-13
lines changed

2 files changed

+33
-13
lines changed

‎google/api_core/operation.py

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ def _cancel_grpc(operations_stub, operation_name):
287287
operations_stub.CancelOperation(request_pb)
288288

289289

290-
def from_grpc(operation, operations_stub, result_type, **kwargs):
290+
def from_grpc(operation, operations_stub, result_type, grpc_metadata=None, **kwargs):
291291
"""Create an operation future using a gRPC client.
292292
293293
This interacts with the long-running operations `service`_ (specific
@@ -302,18 +302,20 @@ def from_grpc(operation, operations_stub, result_type, **kwargs):
302302
operations_stub (google.longrunning.operations_pb2.OperationsStub):
303303
The operations stub.
304304
result_type (:func:`type`): The protobuf result type.
305+
grpc_metadata (Optional[Tuple[str, str]]): Additional metadata to pass
306+
to the rpc.
305307
kwargs: Keyword args passed into the :class:`Operation` constructor.
306308
307309
Returns:
308310
~.api_core.operation.Operation: The operation future to track the given
309311
operation.
310312
"""
311-
refresh = functools.partial(_refresh_grpc, operations_stub, operation.name)
312-
cancel = functools.partial(_cancel_grpc, operations_stub, operation.name)
313+
refresh = functools.partial(_refresh_grpc, operations_stub, operation.name, metadata=grpc_metadata)
314+
cancel = functools.partial(_cancel_grpc, operations_stub, operation.name, metadata=grpc_metadata)
313315
return Operation(operation, refresh, cancel, result_type, **kwargs)
314316

315317

316-
def from_gapic(operation, operations_client, result_type, **kwargs):
318+
def from_gapic(operation, operations_client, result_type, grpc_metadata=None, **kwargs):
317319
"""Create an operation future from a gapic client.
318320
319321
This interacts with the long-running operations `service`_ (specific
@@ -328,12 +330,14 @@ def from_gapic(operation, operations_client, result_type, **kwargs):
328330
operations_client (google.api_core.operations_v1.OperationsClient):
329331
The operations client.
330332
result_type (:func:`type`): The protobuf result type.
333+
grpc_metadata (Optional[Tuple[str, str]]): Additional metadata to pass
334+
to the rpc.
331335
kwargs: Keyword args passed into the :class:`Operation` constructor.
332336
333337
Returns:
334338
~.api_core.operation.Operation: The operation future to track the given
335339
operation.
336340
"""
337-
refresh = functools.partial(operations_client.get_operation, operation.name)
338-
cancel = functools.partial(operations_client.cancel_operation, operation.name)
341+
refresh = functools.partial(operations_client.get_operation, operation.name, metadata=grpc_metadata)
342+
cancel = functools.partial(operations_client.cancel_operation, operation.name, metadata=grpc_metadata)
339343
return Operation(operation, refresh, cancel, result_type, **kwargs)

‎google/api_core/operations_v1/operations_client.py

Lines changed: 23 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,11 @@ def __init__(self, channel, client_config=operations_client_config.config):
9191

9292
# Service calls
9393
def get_operation(
94-
self, name, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT
94+
self,
95+
name,
96+
retry=gapic_v1.method.DEFAULT,
97+
timeout=gapic_v1.method.DEFAULT,
98+
metadata=None,
9599
):
96100
"""Gets the latest state of a long-running operation.
97101
@@ -117,6 +121,7 @@ def get_operation(
117121
unspecified, the the default timeout in the client
118122
configuration is used. If ``None``, then the RPC method will
119123
not time out.
124+
metadata (Optional[Tuple[str, str]]): Additional gRPC metadata.
120125
121126
Returns:
122127
google.longrunning.operations_pb2.Operation: The state of the
@@ -128,14 +133,15 @@ def get_operation(
128133
subclass will be raised.
129134
"""
130135
request = operations_pb2.GetOperationRequest(name=name)
131-
return self._get_operation(request, retry=retry, timeout=timeout)
136+
return self._get_operation(request, retry=retry, timeout=timeout, metadata=metadata)
132137

133138
def list_operations(
134139
self,
135140
name,
136141
filter_,
137142
retry=gapic_v1.method.DEFAULT,
138143
timeout=gapic_v1.method.DEFAULT,
144+
metadata=None,
139145
):
140146
"""
141147
Lists operations that match the specified filter in the request.
@@ -188,7 +194,7 @@ def list_operations(
188194
request = operations_pb2.ListOperationsRequest(name=name, filter=filter_)
189195

190196
# Create the method used to fetch pages
191-
method = functools.partial(self._list_operations, retry=retry, timeout=timeout)
197+
method = functools.partial(self._list_operations, retry=retry, timeout=timeout, metadata=metadata)
192198

193199
iterator = page_iterator.GRPCIterator(
194200
client=None,
@@ -202,7 +208,11 @@ def list_operations(
202208
return iterator
203209

204210
def cancel_operation(
205-
self, name, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT
211+
self,
212+
name,
213+
retry=gapic_v1.method.DEFAULT,
214+
timeout=gapic_v1.method.DEFAULT,
215+
metadata=None,
206216
):
207217
"""Starts asynchronous cancellation on a long-running operation.
208218
@@ -234,6 +244,7 @@ def cancel_operation(
234244
unspecified, the the default timeout in the client
235245
configuration is used. If ``None``, then the RPC method will
236246
not time out.
247+
metadata (Optional[Tuple[str, str]]): Additional gRPC metadata.
237248
238249
Raises:
239250
google.api_core.exceptions.MethodNotImplemented: If the server
@@ -245,10 +256,14 @@ def cancel_operation(
245256
"""
246257
# Create the request object.
247258
request = operations_pb2.CancelOperationRequest(name=name)
248-
self._cancel_operation(request, retry=retry, timeout=timeout)
259+
self._cancel_operation(request, retry=retry, timeout=timeout, metadata=metadata)
249260

250261
def delete_operation(
251-
self, name, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT
262+
self,
263+
name,
264+
retry=gapic_v1.method.DEFAULT,
265+
timeout=gapic_v1.method.DEFAULT,
266+
metadata=None,
252267
):
253268
"""Deletes a long-running operation.
254269
@@ -274,6 +289,7 @@ def delete_operation(
274289
unspecified, the the default timeout in the client
275290
configuration is used. If ``None``, then the RPC method will
276291
not time out.
292+
metadata (Optional[Tuple[str, str]]): Additional gRPC metadata.
277293
278294
Raises:
279295
google.api_core.exceptions.MethodNotImplemented: If the server
@@ -285,4 +301,4 @@ def delete_operation(
285301
"""
286302
# Create the request object.
287303
request = operations_pb2.DeleteOperationRequest(name=name)
288-
self._delete_operation(request, retry=retry, timeout=timeout)
304+
self._delete_operation(request, retry=retry, timeout=timeout, metadata=metadata)

0 commit comments

Comments
 (0)
0