Skip to content

Commit 73854e8

Browse files
busunkim96busunkim96
authored
feat: allow gRPC metadata to be passed to operations client (#127)
1 parent c5fee89 commit 73854e8

File tree

8 files changed

+96
-31
lines changed

8 files changed

+96
-31
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[List[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[List[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/operation_async.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,7 @@ async def cancelled(self):
189189
)
190190

191191

192-
def from_gapic(operation, operations_client, result_type, **kwargs):
192+
def from_gapic(operation, operations_client, result_type, grpc_metadata=None, **kwargs):
193193
"""Create an operation future from a gapic client.
194194
195195
This interacts with the long-running operations `service`_ (specific
@@ -204,12 +204,14 @@ def from_gapic(operation, operations_client, result_type, **kwargs):
204204
operations_client (google.api_core.operations_v1.OperationsClient):
205205
The operations client.
206206
result_type (:func:`type`): The protobuf result type.
207+
grpc_metadata (Optional[List[Tuple[str, str]]]): Additional metadata to pass
208+
to the rpc.
207209
kwargs: Keyword args passed into the :class:`Operation` constructor.
208210
209211
Returns:
210212
~.api_core.operation.Operation: The operation future to track the given
211213
operation.
212214
"""
213-
refresh = functools.partial(operations_client.get_operation, operation.name)
214-
cancel = functools.partial(operations_client.cancel_operation, operation.name)
215+
refresh = functools.partial(operations_client.get_operation, operation.name, metadata=grpc_metadata)
216+
cancel = functools.partial(operations_client.cancel_operation, operation.name, metadata=grpc_metadata)
215217
return AsyncOperation(operation, refresh, cancel, result_type, **kwargs)

‎google/api_core/operations_v1/operations_async_client.py

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

7979
async def get_operation(
80-
self, name, retry=gapic_v1.method_async.DEFAULT, timeout=gapic_v1.method_async.DEFAULT
80+
self,
81+
name,
82+
retry=gapic_v1.method_async.DEFAULT,
83+
timeout=gapic_v1.method_async.DEFAULT,
84+
metadata=None,
8185
):
8286
"""Gets the latest state of a long-running operation.
8387
@@ -103,6 +107,8 @@ async def get_operation(
103107
unspecified, the the default timeout in the client
104108
configuration is used. If ``None``, then the RPC method will
105109
not time out.
110+
metadata (Optional[List[Tuple[str, str]]]):
111+
Additional gRPC metadata.
106112
107113
Returns:
108114
google.longrunning.operations_pb2.Operation: The state of the
@@ -114,14 +120,15 @@ async def get_operation(
114120
subclass will be raised.
115121
"""
116122
request = operations_pb2.GetOperationRequest(name=name)
117-
return await self._get_operation(request, retry=retry, timeout=timeout)
123+
return await self._get_operation(request, retry=retry, timeout=timeout, metadata=metadata)
118124

119125
async def list_operations(
120126
self,
121127
name,
122128
filter_,
123129
retry=gapic_v1.method_async.DEFAULT,
124130
timeout=gapic_v1.method_async.DEFAULT,
131+
metadata=None,
125132
):
126133
"""
127134
Lists operations that match the specified filter in the request.
@@ -157,6 +164,8 @@ async def list_operations(
157164
unspecified, the the default timeout in the client
158165
configuration is used. If ``None``, then the RPC method will
159166
not time out.
167+
metadata (Optional[List[Tuple[str, str]]]): Additional gRPC
168+
metadata.
160169
161170
Returns:
162171
google.api_core.page_iterator.Iterator: An iterator that yields
@@ -174,7 +183,7 @@ async def list_operations(
174183
request = operations_pb2.ListOperationsRequest(name=name, filter=filter_)
175184

176185
# Create the method used to fetch pages
177-
method = functools.partial(self._list_operations, retry=retry, timeout=timeout)
186+
method = functools.partial(self._list_operations, retry=retry, timeout=timeout, metadata=metadata)
178187

179188
iterator = page_iterator_async.AsyncGRPCIterator(
180189
client=None,
@@ -188,7 +197,11 @@ async def list_operations(
188197
return iterator
189198

190199
async def cancel_operation(
191-
self, name, retry=gapic_v1.method_async.DEFAULT, timeout=gapic_v1.method_async.DEFAULT
200+
self,
201+
name,
202+
retry=gapic_v1.method_async.DEFAULT,
203+
timeout=gapic_v1.method_async.DEFAULT,
204+
metadata=None,
192205
):
193206
"""Starts asynchronous cancellation on a long-running operation.
194207
@@ -228,13 +241,19 @@ async def cancel_operation(
228241
google.api_core.exceptions.GoogleAPICallError: If an error occurred
229242
while invoking the RPC, the appropriate ``GoogleAPICallError``
230243
subclass will be raised.
244+
metadata (Optional[List[Tuple[str, str]]]): Additional gRPC
245+
metadata.
231246
"""
232247
# Create the request object.
233248
request = operations_pb2.CancelOperationRequest(name=name)
234-
await self._cancel_operation(request, retry=retry, timeout=timeout)
249+
await self._cancel_operation(request, retry=retry, timeout=timeout, metadata=metadata)
235250

236251
async def delete_operation(
237-
self, name, retry=gapic_v1.method_async.DEFAULT, timeout=gapic_v1.method_async.DEFAULT
252+
self,
253+
name,
254+
retry=gapic_v1.method_async.DEFAULT,
255+
timeout=gapic_v1.method_async.DEFAULT,
256+
metadata=None,
238257
):
239258
"""Deletes a long-running operation.
240259
@@ -260,6 +279,8 @@ async def delete_operation(
260279
unspecified, the the default timeout in the client
261280
configuration is used. If ``None``, then the RPC method will
262281
not time out.
282+
metadata (Optional[List[Tuple[str, str]]]): Additional gRPC
283+
metadata.
263284
264285
Raises:
265286
google.api_core.exceptions.MethodNotImplemented: If the server
@@ -271,4 +292,4 @@ async def delete_operation(
271292
"""
272293
# Create the request object.
273294
request = operations_pb2.DeleteOperationRequest(name=name)
274-
await self._delete_operation(request, retry=retry, timeout=timeout)
295+
await self._delete_operation(request, retry=retry, timeout=timeout, metadata=metadata)

‎google/api_core/operations_v1/operations_client.py

Lines changed: 28 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,8 @@ 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[List[Tuple[str, str]]]):
125+
Additional gRPC metadata.
120126
121127
Returns:
122128
google.longrunning.operations_pb2.Operation: The state of the
@@ -128,14 +134,15 @@ def get_operation(
128134
subclass will be raised.
129135
"""
130136
request = operations_pb2.GetOperationRequest(name=name)
131-
return self._get_operation(request, retry=retry, timeout=timeout)
137+
return self._get_operation(request, retry=retry, timeout=timeout, metadata=metadata)
132138

133139
def list_operations(
134140
self,
135141
name,
136142
filter_,
137143
retry=gapic_v1.method.DEFAULT,
138144
timeout=gapic_v1.method.DEFAULT,
145+
metadata=None,
139146
):
140147
"""
141148
Lists operations that match the specified filter in the request.
@@ -171,6 +178,8 @@ def list_operations(
171178
unspecified, the the default timeout in the client
172179
configuration is used. If ``None``, then the RPC method will
173180
not time out.
181+
metadata (Optional[List[Tuple[str, str]]]): Additional gRPC
182+
metadata.
174183
175184
Returns:
176185
google.api_core.page_iterator.Iterator: An iterator that yields
@@ -188,7 +197,7 @@ def list_operations(
188197
request = operations_pb2.ListOperationsRequest(name=name, filter=filter_)
189198

190199
# Create the method used to fetch pages
191-
method = functools.partial(self._list_operations, retry=retry, timeout=timeout)
200+
method = functools.partial(self._list_operations, retry=retry, timeout=timeout, metadata=metadata)
192201

193202
iterator = page_iterator.GRPCIterator(
194203
client=None,
@@ -202,7 +211,11 @@ def list_operations(
202211
return iterator
203212

204213
def cancel_operation(
205-
self, name, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT
214+
self,
215+
name,
216+
retry=gapic_v1.method.DEFAULT,
217+
timeout=gapic_v1.method.DEFAULT,
218+
metadata=None,
206219
):
207220
"""Starts asynchronous cancellation on a long-running operation.
208221
@@ -234,6 +247,8 @@ def cancel_operation(
234247
unspecified, the the default timeout in the client
235248
configuration is used. If ``None``, then the RPC method will
236249
not time out.
250+
metadata (Optional[List[Tuple[str, str]]]): Additional gRPC
251+
metadata.
237252
238253
Raises:
239254
google.api_core.exceptions.MethodNotImplemented: If the server
@@ -245,10 +260,14 @@ def cancel_operation(
245260
"""
246261
# Create the request object.
247262
request = operations_pb2.CancelOperationRequest(name=name)
248-
self._cancel_operation(request, retry=retry, timeout=timeout)
263+
self._cancel_operation(request, retry=retry, timeout=timeout, metadata=metadata)
249264

250265
def delete_operation(
251-
self, name, retry=gapic_v1.method.DEFAULT, timeout=gapic_v1.method.DEFAULT
266+
self,
267+
name,
268+
retry=gapic_v1.method.DEFAULT,
269+
timeout=gapic_v1.method.DEFAULT,
270+
metadata=None,
252271
):
253272
"""Deletes a long-running operation.
254273
@@ -274,6 +293,8 @@ def delete_operation(
274293
unspecified, the the default timeout in the client
275294
configuration is used. If ``None``, then the RPC method will
276295
not time out.
296+
metadata (Optional[List[Tuple[str, str]]]): Additional gRPC
297+
metadata.
277298
278299
Raises:
279300
google.api_core.exceptions.MethodNotImplemented: If the server
@@ -285,4 +306,4 @@ def delete_operation(
285306
"""
286307
# Create the request object.
287308
request = operations_pb2.DeleteOperationRequest(name=name)
288-
self._delete_operation(request, retry=retry, timeout=timeout)
309+
self._delete_operation(request, retry=retry, timeout=timeout, metadata=metadata)

‎tests/asyncio/operations_v1/test_operations_async_client.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,10 @@ async def test_get_operation():
3636
operations_pb2.Operation(name="meep"))
3737
client = operations_v1.OperationsAsyncClient(mocked_channel)
3838

39-
response = await client.get_operation("name")
39+
response = await client.get_operation("name", metadata=[("x-goog-request-params", "foo")])
4040
assert method.call_count == 1
4141
assert tuple(method.call_args_list[0])[0][0].name == "name"
42+
assert ("x-goog-request-params", "foo") in tuple(method.call_args_list[0])[1]["metadata"]
4243
assert response == fake_call.response
4344

4445

@@ -53,7 +54,7 @@ async def test_list_operations():
5354
mocked_channel, method, fake_call = _mock_grpc_objects(list_response)
5455
client = operations_v1.OperationsAsyncClient(mocked_channel)
5556

56-
pager = await client.list_operations("name", "filter")
57+
pager = await client.list_operations("name", "filter", metadata=[("x-goog-request-params", "foo")])
5758

5859
assert isinstance(pager, page_iterator_async.AsyncIterator)
5960
responses = []
@@ -63,6 +64,7 @@ async def test_list_operations():
6364
assert responses == operations
6465

6566
assert method.call_count == 1
67+
assert ("x-goog-request-params", "foo") in tuple(method.call_args_list[0])[1]["metadata"]
6668
request = tuple(method.call_args_list[0])[0][0]
6769
assert isinstance(request, operations_pb2.ListOperationsRequest)
6870
assert request.name == "name"
@@ -75,10 +77,11 @@ async def test_delete_operation():
7577
empty_pb2.Empty())
7678
client = operations_v1.OperationsAsyncClient(mocked_channel)
7779

78-
await client.delete_operation("name")
80+
await client.delete_operation("name", metadata=[("x-goog-request-params", "foo")])
7981

8082
assert method.call_count == 1
8183
assert tuple(method.call_args_list[0])[0][0].name == "name"
84+
assert ("x-goog-request-params", "foo") in tuple(method.call_args_list[0])[1]["metadata"]
8285

8386

8487
@pytest.mark.asyncio
@@ -87,7 +90,8 @@ async def test_cancel_operation():
8790
empty_pb2.Empty())
8891
client = operations_v1.OperationsAsyncClient(mocked_channel)
8992

90-
await client.cancel_operation("name")
93+
await client.cancel_operation("name", metadata=[("x-goog-request-params", "foo")])
9194

9295
assert method.call_count == 1
9396
assert tuple(method.call_args_list[0])[0][0].name == "name"
97+
assert ("x-goog-request-params", "foo") in tuple(method.call_args_list[0])[1]["metadata"]

‎tests/asyncio/test_operation_async.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -177,12 +177,15 @@ def test_from_gapic():
177177
operations_client,
178178
struct_pb2.Struct,
179179
metadata_type=struct_pb2.Struct,
180+
grpc_metadata=[('x-goog-request-params', 'foo')]
180181
)
181182

182183
assert future._result_type == struct_pb2.Struct
183184
assert future._metadata_type == struct_pb2.Struct
184185
assert future.operation.name == TEST_OPERATION_NAME
185186
assert future.done
187+
assert future._refresh.keywords["metadata"] == [('x-goog-request-params', 'foo')]
188+
assert future._cancel.keywords["metadata"] == [('x-goog-request-params', 'foo')]
186189

187190

188191
def test_deserialize():

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy