feat: add support for long-running operations with rest transport.

gapic/schema/api.py

@@ -27,11 +27,14 @@
 from types import MappingProxyType
 
 from google.api_core import exceptions
+from google.api import http_pb2  # type: ignore
 from google.api import resource_pb2  # type: ignore
+from google.api import service_pb2  # type: ignore
 from google.gapic.metadata import gapic_metadata_pb2  # type: ignore
 from google.longrunning import operations_pb2  # type: ignore
 from google.protobuf import descriptor_pb2
 from google.protobuf.json_format import MessageToJson
+from google.protobuf.json_format import ParseDict
 
 import grpc  # type: ignore
 
@@ -226,6 +229,7 @@ class API:
     """
     naming: api_naming.Naming
     all_protos: Mapping[str, Proto]
+    service_yaml_config: service_pb2.Service
     subpackage_view: Tuple[str, ...] = dataclasses.field(default_factory=tuple)
 
     @classmethod
@@ -318,8 +322,14 @@ def disambiguate_keyword_fname(
             for name, proto in pre_protos.items()
         }
 
+        # Parse the google.api.Service proto from the service_yaml data.
+        service_yaml_config = service_pb2.Service()
+        ParseDict(opts.service_yaml_config, service_yaml_config)
+
         # Done; return the API.
-        return cls(naming=naming, all_protos=protos)
+        return cls(naming=naming,
+                   all_protos=protos,
+                   service_yaml_config=service_yaml_config)
 
     @cached_property
     def enums(self) -> Mapping[str, wrappers.EnumType]:
@@ -374,6 +384,23 @@ def services(self) -> Mapping[str, wrappers.Service]:
                                     *[p.services for p in self.protos.values()],
                                     )
 
+    @cached_property
+    def http_options(self) -> Mapping[str, wrappers.HttpRule]:
+        """Return a map of API-wide http rules."""
+
+        def make_http_options(rule: http_pb2.HttpRule
+                              ) -> Sequence[wrappers.HttpRule]:
+            http_options = [rule] + list(rule.additional_bindings)
+            opt_gen = (wrappers.HttpRule.try_parse_http_rule(http_rule)
+                       for http_rule in http_options)
+            return [rule for rule in opt_gen if rule]
+
+        result = {}
+        for rule in self.service_yaml_config.http.rules:
+            result[rule.selector] = make_http_options(rule)
+
+        return result
+
     @cached_property
     def subpackages(self) -> Mapping[str, 'API']:
         """Return a map of all subpackages, if any.

gapic/templates/%namespace/%name_%version/%sub/services/%service/transports/rest.py.j2

@@ -8,8 +8,10 @@ from google.api_core import retry as retries
 from google.api_core import rest_helpers
 from google.api_core import path_template
 from google.api_core import gapic_v1
+from google.api_core import operations_v1
 {% if service.has_lro %}
 from google.api_core import operations_v1
+from google.protobuf import json_format
 {% endif %}
 from requests import __version__ as requests_version
 from typing import Callable, Dict, Optional, Sequence, Tuple, Union
@@ -25,10 +27,6 @@ except AttributeError:  # pragma: NO COVER
 {% block content %}
 
 
-{% if service.has_lro %}
-{% endif %}
-
-
 {# TODO(yon-mg): re-add python_import/ python_modules from removed diff/current grpc template code #}
 {% filter sort_lines %}
 {% for method in service.methods.values() %}
@@ -134,31 +132,41 @@ class {{service.name}}RestTransport({{service.name}}Transport):
         This property caches on the instance; repeated calls return the same
         client.
         """
-        # Sanity check: Only create a new client if we do not already have one.
+        # Only create a new client if we do not already have one.
         if self._operations_client is None:
-            from google.api_core import grpc_helpers
-
-            self._operations_client = operations_v1.OperationsClient(
-                grpc_helpers.create_channel(
-                    self._host,
+            http_options = {
+            {% for selector, rules in api.http_options.items() %}
+            {% if selector.startswith('google.longrunning.Operations') %}
+                '{{ selector }}': [
+                    {% for rule in rules %}
+                    {
+                        'method': '{{ rule.method }}',
+                        'uri': '{{ rule.uri }}',
+                        {% if rule.body %}
+                        'body': '{{ rule.body }}',
+                        {% endif %}
+                    },
+                    {% endfor %}{# rules #}
+                ],
+            {% endif %}{# longrunning.Operations #}
+            {% endfor %}{# http_options #}
+            }
+
+            rest_transport = operations_v1.OperationsRestTransport(
+                    host=self._host,
                     credentials=self._credentials,
-                    default_scopes=cls.AUTH_SCOPES,
                     scopes=self._scopes,
-                    default_host=cls.DEFAULT_HOST,
-                    options=[
-                        ("grpc.max_send_message_length", -1),
-                        ("grpc.max_receive_message_length", -1),
-                    ],
-                )
-            )
+                    http_options=http_options)
+
+            self._operations_client = operations_v1.OperationsRestClient(transport=rest_transport)
 
         # Return the client from cache.
         return self._operations_client
 
 
-    {% endif %}
+    {% endif %}{# service.has_lro #}
     {% for method in service.methods.values() %}
-    {%- if method.http_options and not method.lro and not (method.server_streaming or method.client_streaming) %}
+    {%- if method.http_options and not (method.server_streaming or method.client_streaming) %}
     def _{{method.name | snake_case}}(self,
             request: {{method.input.ident}}, *,
             retry: OptionalRetry=gapic_v1.method.DEFAULT,
@@ -279,10 +287,16 @@ class {{service.name}}RestTransport({{service.name}}Transport):
         {% if not method.void %}
 
         # Return the response
+        {% if method.lro %}
+        return_op = operations_pb2.Operation()
+        json_format.Parse(response.content, return_op, ignore_unknown_fields=True)
+        return return_op
+        {% else %}
         return {{method.output.ident}}.from_json(
             response.content,
             ignore_unknown_fields=True
         )
+    {% endif %}
         {% endif %}
     {% else %}
 
@@ -296,10 +310,6 @@ class {{service.name}}RestTransport({{service.name}}Transport):
 
         raise RuntimeError(
             "Cannot define a method without a valid 'google.api.http' annotation.")
-        {%- elif method.lro %}
-
-        raise NotImplementedError(
-            "LRO over REST is not yet defined for python client.")
         {%- elif method.server_streaming or method.client_streaming %}
 
         raise NotImplementedError(

gapic/utils/options.py

@@ -20,6 +20,7 @@
 import json
 import os
 import warnings
+import yaml
 
 from gapic.samplegen_utils import utils as samplegen_utils
 
@@ -45,6 +46,8 @@ class Options:
     metadata: bool = False
     # TODO(yon-mg): should there be an enum for transport type?
     transport: List[str] = dataclasses.field(default_factory=lambda: [])
+    service_yaml_config: Dict[str, Any] = dataclasses.field(
+        default_factory=dict)
 
     # Class constants
     PYTHON_GAPIC_PREFIX: str = 'python-gapic-'
@@ -54,6 +57,7 @@ class Options:
         'metadata',             # generate GAPIC metadata JSON file
         'old-naming',           # TODO(dovs): Come up with a better comment
         'retry-config',         # takes a path
+        'service-yaml',         # takes a path
         'samples',              # output dir
         'autogen-snippets',     # produce auto-generated snippets
         # transport type(s) delineated by '+' (i.e. grpc, rest, custom.[something], etc?)
@@ -129,6 +133,16 @@ def tweak_path(p):
             with open(retry_paths[-1]) as f:
                 retry_cfg = json.load(f)
 
+        service_yaml_config = {}
+        service_yaml_paths = opts.pop('service-yaml', None)
+        if service_yaml_paths:
+            # Just use the last file specified.
+            with open(service_yaml_paths[-1]) as f:
+                service_yaml_config = yaml.load(f, Loader=yaml.Loader)
+        # The yaml service files typically have this field,
+        # but it is not a field in the gogle.api.Service proto.
+        service_yaml_config.pop('type', None)
+
         # Build the options instance.
         sample_paths = opts.pop('samples', [])
 
@@ -150,7 +164,8 @@ def tweak_path(p):
             add_iam_methods=bool(opts.pop('add-iam-methods', False)),
             metadata=bool(opts.pop('metadata', False)),
             # transport should include desired transports delimited by '+', e.g. transport='grpc+rest'
-            transport=opts.pop('transport', ['grpc'])[0].split('+')
+            transport=opts.pop('transport', ['grpc'])[0].split('+'),
+            service_yaml_config=service_yaml_config
         )
 
         # Note: if we ever need to recursively check directories for sample

rules_python_gapic/py_gapic.bzl

@@ -21,6 +21,7 @@ def py_gapic_library(
         plugin_args = None,
         opt_args = None,
         metadata = True,
+        service_yaml = None,
         **kwargs):
     #    srcjar_target_name = "%s_srcjar" % name
     srcjar_target_name = name
@@ -35,6 +36,8 @@ def py_gapic_library(
     file_args = {}
     if grpc_service_config:
         file_args[grpc_service_config] = "retry-config"
+    if service_yaml:
+        file_args[service_yaml] = "service-yaml"
 
     proto_custom_library(
         name = srcjar_target_name,

tests/unit/generator/test_generator.py

@@ -19,6 +19,7 @@
 import jinja2
 import pytest
 
+from google.api import service_pb2
 from google.protobuf import descriptor_pb2
 from google.protobuf.compiler.plugin_pb2 import CodeGeneratorResponse
 
@@ -767,9 +768,17 @@ def make_proto(
     ).proto
 
 
-def make_api(*protos, naming: naming.Naming = None, **kwargs) -> api.API:
+def make_api(
+    *protos,
+    naming: naming.Naming = None,
+    service_yaml_config: service_pb2.Service = None,
+    **kwargs
+) -> api.API:
     return api.API(
-        naming=naming or make_naming(), all_protos={i.name: i for i in protos}, **kwargs
+        naming=naming or make_naming(),
+        service_yaml_config=service_yaml_config or service_pb2.Service(),
+        all_protos={i.name: i for i in protos},
+        **kwargs
     )
 
 

tests/unit/generator/test_options.py

@@ -140,6 +140,36 @@ def test_options_service_config(fs):
     assert opts.retry == expected_cfg
 
 
+def test_options_service_yaml_config(fs):
+    opts = Options.build("")
+    assert opts.service_yaml_config == {}
+
+    service_yaml_fpath = "testapi_v1.yaml"
+    fs.create_file(service_yaml_fpath,
+                   contents=("type: google.api.Service\n"
+                             "config_version: 3\n"
+                             "name: testapi.googleapis.com\n"))
+    opt_string = f"service-yaml={service_yaml_fpath}"
+    opts = Options.build(opt_string)
+    expected_config = {
+        "config_version": 3,
+        "name": "testapi.googleapis.com"
+    }
+    assert opts.service_yaml_config == expected_config
+
+    service_yaml_fpath = "testapi_v2.yaml"
+    fs.create_file(service_yaml_fpath,
+                   contents=("config_version: 3\n"
+                             "name: testapi.googleapis.com\n"))
+    opt_string = f"service-yaml={service_yaml_fpath}"
+    opts = Options.build(opt_string)
+    expected_config = {
+        "config_version": 3,
+        "name": "testapi.googleapis.com"
+    }
+    assert opts.service_yaml_config == expected_config
+
+
 def test_options_bool_flags():
     # All these options are default False.
     # If new options violate this assumption,

tests/unit/schema/test_api.py

@@ -1560,3 +1560,38 @@ def test_gapic_metadata():
     expected = MessageToJson(expected, sort_keys=True)
     actual = api_schema.gapic_metadata_json(opts)
     assert expected == actual
+
+
+def test_http_options(fs):
+    fd = (
+        make_file_pb2(
+            name='example.proto',
+            package='google.example.v1',
+            messages=(make_message_pb2(name='ExampleRequest', fields=()),),
+        ),)
+
+    opts = Options(service_yaml_config={
+        'http': {
+            'rules': [
+                {
+                    'selector': 'Cancel',
+                    'post': '/v3/{name=projects/*/locations/*/operations/*}:cancel',
+                    'body': '*'
+                },
+                {
+                    'selector': 'Get',
+                    'get': '/v3/{name=projects/*/locations/*/operations/*}',
+                    'additional_bindings': [{'get': '/v3/{name=/locations/*/operations/*}'}],
+                }, ]
+        }
+    })
+
+    api_schema = api.API.build(fd, 'google.example.v1', opts=opts)
+    http_options = api_schema.http_options
+    assert http_options == {
+        'Cancel': [wrappers.HttpRule(method='post', uri='/v3/{name=projects/*/locations/*/operations/*}:cancel', body='*')],
+        'Get': [
+            wrappers.HttpRule(
+                method='get', uri='/v3/{name=projects/*/locations/*/operations/*}', body=None),
+            wrappers.HttpRule(method='get', uri='/v3/{name=/locations/*/operations/*}', body=None)]
+    }