Skip to content

Commit e7bee70

Browse files
feat: expose extended operations annotations within generator (googleapis#1145)
* feat: full diregapic LROs WIP add test, test fails * Style check * Integrate reviews * Failures * Mypy Co-authored-by: Anthonios Partheniou <[email protected]>
1 parent d84e541 commit e7bee70

File tree

5 files changed

+172
-8
lines changed

5 files changed

+172
-8
lines changed

gapic/schema/api.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
from google.api import http_pb2 # type: ignore
3131
from google.api import resource_pb2 # type: ignore
3232
from google.api import service_pb2 # type: ignore
33+
from google.cloud import extended_operations_pb2 as ex_ops_pb2 # type: ignore
3334
from google.gapic.metadata import gapic_metadata_pb2 # type: ignore
3435
from google.longrunning import operations_pb2 # type: ignore
3536
from google.protobuf import descriptor_pb2
@@ -474,6 +475,20 @@ def requires_package(self, pkg: Tuple[str, ...]) -> bool:
474475
for message in proto.all_messages.values()
475476
)
476477

478+
def get_custom_operation_service(self, method: "wrappers.Method") -> "wrappers.Service":
479+
if not method.output.is_extended_operation:
480+
raise ValueError(
481+
f"Method is not an extended operation LRO: {method.name}")
482+
483+
op_serv_name = self.naming.proto_package + "." + \
484+
method.options.Extensions[ex_ops_pb2.operation_service]
485+
op_serv = self.services[op_serv_name]
486+
if not op_serv.custom_polling_method:
487+
raise ValueError(
488+
f"Service is not an extended operation operation service: {op_serv.name}")
489+
490+
return op_serv
491+
477492

478493
class _ProtoBuilder:
479494
"""A "builder class" for Proto objects.

gapic/schema/wrappers.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -360,7 +360,7 @@ def oneof_fields(self, include_optional=False):
360360
return oneof_fields
361361

362362
@utils.cached_property
363-
def is_diregapic_operation(self) -> bool:
363+
def is_extended_operation(self) -> bool:
364364
if not self.name == "Operation":
365365
return False
366366

@@ -877,7 +877,7 @@ def __getattr__(self, name):
877877

878878
@property
879879
def is_operation_polling_method(self):
880-
return self.output.is_diregapic_operation and self.options.Extensions[ex_ops_pb2.operation_polling_method]
880+
return self.output.is_extended_operation and self.options.Extensions[ex_ops_pb2.operation_polling_method]
881881

882882
@utils.cached_property
883883
def client_output(self):

tests/unit/schema/test_api.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
from google.api import client_pb2
2323
from google.api import resource_pb2
2424
from google.api_core import exceptions
25+
from google.cloud import extended_operations_pb2 as ex_ops_pb2
2526
from google.gapic.metadata import gapic_metadata_pb2
2627
from google.longrunning import operations_pb2
2728
from google.protobuf import descriptor_pb2
@@ -1595,3 +1596,151 @@ def test_http_options(fs):
15951596
method='get', uri='/v3/{name=projects/*/locations/*/operations/*}', body=None),
15961597
wrappers.HttpRule(method='get', uri='/v3/{name=/locations/*/operations/*}', body=None)]
15971598
}
1599+
1600+
1601+
def generate_basic_extended_operations_setup():
1602+
T = descriptor_pb2.FieldDescriptorProto.Type
1603+
1604+
operation = make_message_pb2(
1605+
name="Operation",
1606+
fields=(
1607+
make_field_pb2(name=name, type=T.Value("TYPE_STRING"), number=i)
1608+
for i, name in enumerate(("name", "status", "error_code", "error_message"), start=1)
1609+
),
1610+
)
1611+
1612+
for f in operation.field:
1613+
options = descriptor_pb2.FieldOptions()
1614+
# Note: The field numbers were carefully chosen to be the corresponding enum values.
1615+
options.Extensions[ex_ops_pb2.operation_field] = f.number
1616+
f.options.MergeFrom(options)
1617+
1618+
options = descriptor_pb2.MethodOptions()
1619+
options.Extensions[ex_ops_pb2.operation_polling_method] = True
1620+
1621+
polling_method = descriptor_pb2.MethodDescriptorProto(
1622+
name="Get",
1623+
input_type="google.extended_operations.v1.stuff.GetOperation",
1624+
output_type="google.extended_operations.v1.stuff.Operation",
1625+
options=options,
1626+
)
1627+
1628+
delete_input_message = make_message_pb2(name="Input")
1629+
delete_output_message = make_message_pb2(name="Output")
1630+
ops_service = descriptor_pb2.ServiceDescriptorProto(
1631+
name="CustomOperations",
1632+
method=[
1633+
polling_method,
1634+
descriptor_pb2.MethodDescriptorProto(
1635+
name="Delete",
1636+
input_type="google.extended_operations.v1.stuff.Input",
1637+
output_type="google.extended_operations.v1.stuff.Output",
1638+
),
1639+
],
1640+
)
1641+
1642+
request = make_message_pb2(
1643+
name="GetOperation",
1644+
fields=[
1645+
make_field_pb2(name="name", type=T.Value("TYPE_STRING"), number=1)
1646+
],
1647+
)
1648+
1649+
initial_opts = descriptor_pb2.MethodOptions()
1650+
initial_opts.Extensions[ex_ops_pb2.operation_service] = ops_service.name
1651+
initial_input_message = make_message_pb2(name="Initial")
1652+
initial_method = descriptor_pb2.MethodDescriptorProto(
1653+
name="CreateTask",
1654+
input_type="google.extended_operations.v1.stuff.GetOperation",
1655+
output_type="google.extended_operations.v1.stuff.Operation",
1656+
options=initial_opts,
1657+
)
1658+
1659+
regular_service = descriptor_pb2.ServiceDescriptorProto(
1660+
name="RegularService",
1661+
method=[
1662+
initial_method,
1663+
],
1664+
)
1665+
1666+
file_protos = [
1667+
make_file_pb2(
1668+
name="extended_operations.proto",
1669+
package="google.extended_operations.v1.stuff",
1670+
messages=[
1671+
operation,
1672+
request,
1673+
delete_output_message,
1674+
delete_input_message,
1675+
initial_input_message,
1676+
],
1677+
services=[
1678+
regular_service,
1679+
ops_service,
1680+
],
1681+
),
1682+
]
1683+
1684+
return file_protos
1685+
1686+
1687+
def test_extended_operations_lro_operation_service():
1688+
file_protos = generate_basic_extended_operations_setup()
1689+
api_schema = api.API.build(file_protos)
1690+
initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"]
1691+
1692+
expected = api_schema.services['google.extended_operations.v1.stuff.CustomOperations']
1693+
actual = api_schema.get_custom_operation_service(initial_method)
1694+
1695+
assert expected is actual
1696+
1697+
assert actual.custom_polling_method is actual.methods["Get"]
1698+
1699+
1700+
def test_extended_operations_lro_operation_service_no_annotation():
1701+
file_protos = generate_basic_extended_operations_setup()
1702+
1703+
api_schema = api.API.build(file_protos)
1704+
initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"]
1705+
# It's easier to manipulate data structures after building the API.
1706+
del initial_method.options.Extensions[ex_ops_pb2.operation_service]
1707+
1708+
with pytest.raises(KeyError):
1709+
api_schema.get_custom_operation_service(initial_method)
1710+
1711+
1712+
def test_extended_operations_lro_operation_service_no_such_service():
1713+
file_protos = generate_basic_extended_operations_setup()
1714+
1715+
api_schema = api.API.build(file_protos)
1716+
initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"]
1717+
initial_method.options.Extensions[ex_ops_pb2.operation_service] = "UnrealService"
1718+
1719+
with pytest.raises(KeyError):
1720+
api_schema.get_custom_operation_service(initial_method)
1721+
1722+
1723+
def test_extended_operations_lro_operation_service_not_an_lro():
1724+
file_protos = generate_basic_extended_operations_setup()
1725+
1726+
api_schema = api.API.build(file_protos)
1727+
initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"]
1728+
# Hack to pretend that the initial_method is not an LRO
1729+
super(type(initial_method), initial_method).__setattr__(
1730+
"output", initial_method.input)
1731+
1732+
with pytest.raises(ValueError):
1733+
api_schema.get_custom_operation_service(initial_method)
1734+
1735+
1736+
def test_extended_operations_lro_operation_service_no_polling_method():
1737+
file_protos = generate_basic_extended_operations_setup()
1738+
1739+
api_schema = api.API.build(file_protos)
1740+
initial_method = api_schema.services["google.extended_operations.v1.stuff.RegularService"].methods["CreateTask"]
1741+
1742+
operation_service = api_schema.services["google.extended_operations.v1.stuff.CustomOperations"]
1743+
del operation_service.methods["Get"].options.Extensions[ex_ops_pb2.operation_polling_method]
1744+
1745+
with pytest.raises(ValueError):
1746+
api_schema.get_custom_operation_service(initial_method)

tests/unit/schema/wrappers/test_message.py

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -331,7 +331,7 @@ def test_required_fields():
331331
assert set(request.required_fields) == {mass_kg, length_m, color}
332332

333333

334-
def test_is_diregapic_operation():
334+
def test_is_extended_operation():
335335
T = descriptor_pb2.FieldDescriptorProto.Type
336336

337337
# Canonical Operation
@@ -349,7 +349,7 @@ def test_is_diregapic_operation():
349349
options.Extensions[ex_ops_pb2.operation_field] = f.number
350350
f.options.MergeFrom(options)
351351

352-
assert operation.is_diregapic_operation
352+
assert operation.is_extended_operation
353353

354354
# Missing a required field
355355

@@ -367,7 +367,7 @@ def test_is_diregapic_operation():
367367
options.Extensions[ex_ops_pb2.operation_field] = f.number
368368
f.options.MergeFrom(options)
369369

370-
assert not missing.is_diregapic_operation
370+
assert not missing.is_extended_operation
371371

372372
# Named incorrectly
373373

@@ -383,7 +383,7 @@ def test_is_diregapic_operation():
383383
options.Extensions[ex_ops_pb2.operation_field] = f.number
384384
f.options.MergeFrom(options)
385385

386-
assert not my_message.is_diregapic_operation
386+
assert not my_message.is_extended_operation
387387

388388
# Duplicated annotation
389389
for mapping in range(1, 5):
@@ -401,4 +401,4 @@ def test_is_diregapic_operation():
401401
f.options.MergeFrom(options)
402402

403403
with pytest.raises(TypeError):
404-
duplicate.is_diregapic_operation
404+
duplicate.is_extended_operation

tests/unit/schema/wrappers/test_service.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -589,7 +589,7 @@ def test_operation_polling_method():
589589
assert not user_service.custom_polling_method
590590

591591

592-
def test_diregapic_lro_detection():
592+
def test_extended_operations_lro_detection():
593593
T = descriptor_pb2.FieldDescriptorProto.Type
594594

595595
operation = make_message(

0 commit comments

Comments
 (0)