feat: Implement REST support for MixIns (#1378)

Implements MixIns support in REST. The implementation is different than what was used for GRPC version. Special wrappers for MixIns are created to make the code less repetitive.

Co-authored-by: Anthonios Partheniou <partheniou@google.com>

Author: Aza Tulepbergenov

gapic/schema/api.py

@@ -43,6 +43,7 @@
 from google.protobuf.descriptor_pb2 import MethodDescriptorProto
 from google.api import annotations_pb2  # type: ignore
 from gapic.schema import metadata
+from gapic.schema import mixins
 from gapic.schema import wrappers
 from gapic.schema import naming as api_naming
 from gapic.utils import cached_property
@@ -515,6 +516,16 @@ def get_custom_operation_service(self, method: "wrappers.Method") -> "wrappers.S
 
         return op_serv
 
+    @cached_property
+    def mixin_api_signatures(self):
+        """Compile useful info about MixIn API signatures.
+
+        Returns:
+            Mapping[str, wrappers.MixinMethod]: Useful info
+                about MixIn methods present for the main API.
+        """
+        return {name: mixins.MIXINS_MAP[name] for name in self.mixin_api_methods}
+
     @cached_property
     def mixin_api_methods(self) -> Dict[str, MethodDescriptorProto]:
         methods: Dict[str, MethodDescriptorProto] = {}
@@ -529,6 +540,20 @@ def mixin_api_methods(self) -> Dict[str, MethodDescriptorProto]:
                 self._get_methods_from_service(operations_pb2)}
         return methods
 
+    @cached_property
+    def mixin_http_options(self):
+        """Gather HTTP options for the MixIn methods."""
+        api_methods = self.mixin_api_methods
+        res = {}
+        for s in api_methods:
+            m = api_methods[s]
+            http = m.options.Extensions[annotations_pb2.http]
+            http_options = [http] + list(http.additional_bindings)
+            opt_gen = (wrappers.MixinHttpRule.try_parse_http_rule(http_rule)
+                   for http_rule in http_options)
+            res[s] = [rule for rule in opt_gen if rule]
+        return res
+
     @cached_property
     def has_location_mixin(self) -> bool:
         return len(list(filter(lambda api: api.name == "google.cloud.location.Locations", self.service_yaml_config.apis))) > 0

gapic/schema/mixins.py

@@ -0,0 +1,68 @@
+# Copyright 2022 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from gapic.schema import wrappers
+
+MIXINS_MAP = {
+    'DeleteOperation': wrappers.MixinMethod(
+        'DeleteOperation',
+        request_type='operations_pb2.DeleteOperationRequest',
+        response_type='None'
+    ),
+    'WaitOperation': wrappers.MixinMethod(
+        'WaitOperation',
+        request_type='operations_pb2.WaitOperationRequest',
+        response_type='operations_pb2.Operation'
+    ),
+    'ListOperations': wrappers.MixinMethod(
+        'ListOperations',
+        request_type='operations_pb2.ListOperationsRequest',
+        response_type='operations_pb2.ListOperationsResponse'
+    ),
+    'CancelOperation': wrappers.MixinMethod(
+        'CancelOperation',
+        request_type='operations_pb2.CancelOperationRequest',
+        response_type='None'
+    ),
+    'GetOperation': wrappers.MixinMethod(
+        'GetOperation',
+        request_type='operations_pb2.GetOperationRequest',
+        response_type='operations_pb2.Operation'
+    ),
+    'TestIamPermissions': wrappers.MixinMethod(
+        'TestIamPermissions',
+        request_type='iam_policy_pb2.TestIamPermissionsRequest',
+        response_type='iam_policy_pb2.TestIamPermissionsResponse'
+    ),
+    'GetIamPolicy': wrappers.MixinMethod(
+        'GetIamPolicy',
+        request_type='iam_policy_pb2.GetIamPolicyRequest',
+        response_type='policy_pb2.Policy'
+    ),
+    'SetIamPolicy': wrappers.MixinMethod(
+        'SetIamPolicy',
+        request_type='iam_policy_pb2.SetIamPolicyRequest',
+        response_type='policy_pb2.Policy'
+    ),
+    'ListLocations': wrappers.MixinMethod(
+        'ListLocations',
+        request_type='locations_pb2.ListLocationsRequest',
+        response_type='locations_pb2.ListLocationsResponse'
+    ),
+    'GetLocation': wrappers.MixinMethod(
+        'GetLocation',
+        request_type='locations_pb2.GetLocationRequest',
+        response_type='locations_pb2.Location'
+    )
+}

gapic/schema/wrappers.py

@@ -1094,6 +1094,32 @@ def try_parse_http_rule(cls, http_rule) -> Optional['HttpRule']:
         return cls(method, uri, body)
 
 
+@dataclasses.dataclass(frozen=True)
+class MixinMethod:
+    name: str
+    request_type: str
+    response_type: str
+
+
+@dataclasses.dataclass(frozen=True)
+class MixinHttpRule(HttpRule):
+    def path_fields(self, uri):
+        """return list of (name, template) tuples extracted from uri."""
+        return [
+            (match.group("name"), match.group("template"))
+            for match in path_template._VARIABLE_RE.finditer(uri)
+            if match.group("name")
+        ]
+
+    @property
+    def sample_request(self):
+        req = uri_sample.sample_from_path_fields(self.path_fields(self.uri))
+        if not self.body or self.body == "" or self.body == "*":
+            return req
+        req[self.body] = {}  # just an empty json.
+        return req
+
+
 @dataclasses.dataclass(frozen=True)
 class Method:
     """Description of a method (defined with the ``rpc`` keyword)."""

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

@@ -0,0 +1,93 @@
+    {% if "rest" in opts.transport %}
+
+    {% for name, sig in api.mixin_api_signatures.items() %} 
+    @property
+    def {{ name|snake_case }}(self):
+        return self._{{ name }}(self._session, self._host, self._interceptor) # type: ignore
+    
+    class _{{ name }}({{service.name}}RestStub):
+        def __call__(self,
+            request: {{ sig.request_type }}, *,
+            retry: OptionalRetry=gapic_v1.method.DEFAULT,
+            timeout: float=None,
+            metadata: Sequence[Tuple[str, str]]=(),
+            ) -> {{ sig.response_type }}:
+
+            r"""Call the {{- ' ' -}}
+            {{ (name|snake_case).replace('_',' ')|wrap(
+                    width=70, offset=45, indent=8) }}
+            {{- ' ' -}} method over HTTP.
+
+            Args:
+                request ({{ sig.request_type }}):
+                    The request object for {{ name }} method.
+                retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                    should be retried.
+                timeout (float): The timeout for this request.
+                metadata (Sequence[Tuple[str, str]]): Strings which should be
+                    sent along with the request as metadata.
+            {% if sig.response_type != 'None' %}
+
+            Returns:
+                {{ sig.response_type }}: Response from {{ name }} method.
+            {% endif %}
+            """
+
+            http_options: List[Dict[str, str]] = [
+            {%- for rule in api.mixin_http_options["{}".format(name)] %}{
+                'method': '{{ rule.method }}',
+                'uri': '{{ rule.uri }}',
+                {% if rule.body %}
+                'body': '{{ rule.body }}',
+                {% endif %}{# rule.body #}
+            },
+            {% endfor %}
+            ]
+
+            request, metadata = self._interceptor.pre_{{ name|snake_case }}(request, metadata)
+            request_kwargs = json_format.MessageToDict(request)
+            transcoded_request = path_template.transcode(
+                http_options, **request_kwargs)
+
+            {% set body_spec = api.mixin_http_options["{}".format(name)][0].body %}
+            {%- if body_spec %}
+            body = json.loads(json.dumps(transcoded_request['body']))
+            {%- endif %}
+
+            uri = transcoded_request['uri']
+            method = transcoded_request['method']
+
+            # Jsonify the query params
+            query_params = json.loads(json.dumps(transcoded_request['query_params']))
+
+            # Send the request
+            headers = dict(metadata)
+            headers['Content-Type'] = 'application/json'
+
+            response = getattr(self._session, method)(
+                "{host}{uri}".format(host=self._host, uri=uri),
+                timeout=timeout,
+                headers=headers,
+                params=rest_helpers.flatten_query_params(query_params),
+                {% if body_spec %}
+                data=body,
+                {% endif %}
+            )
+
+            # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception
+            # subclass.
+            if response.status_code >= 400:
+                raise core_exceptions.from_http_response(response)
+            
+            {% if sig.response_type == "None" %}
+            return self._interceptor.post_{{ name|snake_case }}(None)
+            {% else %}
+            
+            resp = {{ sig.response_type }}()
+            resp = json_format.Parse(response.content.decode("utf-8"), resp)
+            resp = self._interceptor.post_{{ name|snake_case }}(resp)
+            return resp
+            {% endif %}
+
+    {% endfor %}
+    {% endif %} {# rest in opts.transport #}

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

@@ -19,6 +19,16 @@ from google.protobuf import json_format
 {% if service.has_lro %}
 from google.api_core import operations_v1
 {% endif %}
+{% if opts.add_iam_methods or api.has_iam_mixin %}
+from google.iam.v1 import iam_policy_pb2  # type: ignore
+from google.iam.v1 import policy_pb2  # type: ignore
+{% endif %}
+{% if api.has_location_mixin %}
+from google.cloud.location import locations_pb2 # type: ignore
+{% endif %}
+{% if api.has_operations_mixin %}
+from google.longrunning import operations_pb2
+{% endif %}
 from requests import __version__ as requests_version
 import dataclasses
 import re
@@ -113,7 +123,25 @@ class {{ service.name }}RestInterceptor:
         """
         return response
     {% endif %}
+    {% endfor %}
+
+    {% for name, signature in api.mixin_api_signatures.items() %}
+    def pre_{{ name|snake_case }}(self, request: {{signature.request_type}}, metadata: Sequence[Tuple[str, str]]) -> {{signature.response_type}}:
+        """Pre-rpc interceptor for {{ name|snake_case }}
 
+        Override in a subclass to manipulate the request or metadata
+        before they are sent to the {{ service.name }} server.
+        """
+        return request, metadata
+
+    def post_{{ name|snake_case }}(self, response: {{signature.request_type}}) -> {{signature.response_type}}:
+        """Post-rpc interceptor for {{ name|snake_case }}
+
+        Override in a subclass to manipulate the response
+        after it is returned by the {{ service.name }} server but before
+        it is returned to user code.
+        """
+        return response
     {% endfor %}
 
 
@@ -344,7 +372,7 @@ class {{service.name}}RestTransport({{service.name}}Transport):
                 including_default_value_fields=False,
                 use_integers_for_enums={{ opts.rest_numeric_enums }}
             )
-            {%- endif %}{# body_spec #}
+            {%- endif %}
 
             uri = transcoded_request['uri']
             method = transcoded_request['method']
@@ -430,6 +458,8 @@ class {{service.name}}RestTransport({{service.name}}Transport):
 
     {% endfor %}
 
+    {% include '%namespace/%name_%version/%sub/services/%service/transports/_rest_mixins.py.j2' %}
+
     @property
     def kind(self) -> str:
         return "rest"

gapic/templates/tests/unit/gapic/%name_%version/%sub/_test_mixins.py.j2

@@ -1,3 +1,67 @@
+{% if 'rest' in opts.transport %}
+{% for name, sig in api.mixin_api_signatures.items() %}
+
+def test_{{ name|snake_case }}_rest_bad_request(transport: str = 'rest', request_type={{ sig.request_type }}):
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport=transport,
+    )
+
+    request = request_type()
+    request = json_format.ParseDict({{ api.mixin_http_options["{}".format(name)][0].sample_request }}, request)
+
+    # Mock the http request call within the method and fake a BadRequest error.
+    with mock.patch.object(Session, 'request') as req, pytest.raises(core_exceptions.BadRequest):
+        # Wrap the value into a proper Response obj
+        response_value = Response()
+        response_value.status_code = 400
+        response_value.request = Request()
+        req.return_value = response_value
+        client.{{ name|snake_case }}(request)
+
+@pytest.mark.parametrize("request_type", [
+    {{ sig.request_type }},
+    dict,
+])
+def test_{{ name|snake_case }}_rest(request_type):
+    client = {{ service.client_name }}(
+        credentials=ga_credentials.AnonymousCredentials(),
+        transport="rest",
+    )
+    request_init = {{ api.mixin_http_options["{}".format(name)][0].sample_request }}
+    request = request_type(**request_init)
+    # Mock the http request call within the method and fake a response.
+    with mock.patch.object(type(client.transport._session), 'request') as req:
+        # Designate an appropriate value for the returned response.
+        {% if sig.response_type == "None" %}
+        return_value = None
+        {% else %}
+        return_value = {{ sig.response_type }}()
+        {% endif %}
+
+        # Wrap the value into a proper Response obj
+        response_value = Response()
+        response_value.status_code = 200
+        {% if sig.response_type == "None" %}
+        json_return_value = '{}'
+        {% else %}
+        json_return_value = json_format.MessageToJson(return_value)
+        {% endif %}
+
+        response_value._content = json_return_value.encode('UTF-8')
+        req.return_value = response_value
+
+        response = client.{{ name|snake_case }}(request)
+
+    # Establish that the response is the type that we expect.
+    {% if sig.response_type == "None" %}
+    assert response is None
+    {% else %}
+    assert isinstance(response, {{ sig.response_type }})
+    {% endif %}
+{% endfor %}
+{% endif %}
+
 {% if api.has_operations_mixin and ('grpc' in opts.transport or 'grpc_asyncio' in opts.transport) %}
 
 {% if "DeleteOperation" in api.mixin_api_methods %}

gapic/templates/tests/unit/gapic/%name_%version/%sub/test_%service.py.j2

@@ -26,6 +26,7 @@ from proto.marshal.rules.dates import DurationRule, TimestampRule
 from requests import Response
 from requests import Request, PreparedRequest
 from requests.sessions import Session
+from google.protobuf import json_format
 {% endif %}
 
 {# Import the service itself as well as every proto module that it imports. #}