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)