feat: Add support for reading selective GAPIC generation methods from service YAML (#2272)

Co-authored-by: Anthonios Partheniou <[email protected]>
Co-authored-by: Victor Chudnovsky <[email protected]>

Author: 3 people

DEVELOPMENT.md

@@ -60,5 +60,6 @@ Execute unit tests by running one of the sessions prefixed with `unit-`.
     bazel run //tests/integration:eventarc_update
     bazel run //tests/integration:logging_update
     bazel run //tests/integration:redis_update
+    bazel run //tests/integration:redis_selective_update
     ```
 

gapic/schema/api.py

@@ -24,7 +24,7 @@
 import os
 import sys
 from types import MappingProxyType
-from typing import Callable, Container, Dict, FrozenSet, Mapping, Optional, Sequence, Set, Tuple
+from typing import Callable, Container, Dict, FrozenSet, Iterable, Mapping, Optional, Sequence, Set, Tuple
 import yaml
 
 from google.api_core import exceptions
@@ -237,6 +237,95 @@ def disambiguate(self, string: str) -> str:
             return self.disambiguate(f'_{string}')
         return string
 
+    def add_to_address_allowlist(self, *,
+                                 address_allowlist: Set['metadata.Address'],
+                                 method_allowlist: Set[str],
+                                 resource_messages: Dict[str, 'wrappers.MessageType'],
+                                 ) -> None:
+        """Adds to the set of Addresses of wrapper objects to be included in selective GAPIC generation.
+
+        This method is used to create an allowlist of addresses to be used to filter out unneeded
+        services, methods, messages, and enums at a later step.
+
+        Args:
+            address_allowlist (Set[metadata.Address]): A set of allowlisted metadata.Address
+                objects to add to. Only the addresses of the allowlisted methods, the services
+                containing these methods, and messages/enums those methods use will be part of the
+                final address_allowlist. The set may be modified during this call.
+            method_allowlist (Set[str]): An allowlist of fully-qualified method names.
+            resource_messages (Dict[str, wrappers.MessageType]): A dictionary mapping the unified
+                resource type name of a resource message to the corresponding MessageType object
+                representing that resource message. Only resources with a message representation
+                should be included in the dictionary.
+        Returns:
+            None
+        """
+        # The method.operation_service for an extended LRO is not fully qualified, so we
+        # truncate the service names accordingly so they can be found in
+        # method.add_to_address_allowlist
+        services_in_proto = {
+            service.name: service for service in self.services.values()
+        }
+        for service in self.services.values():
+            service.add_to_address_allowlist(address_allowlist=address_allowlist,
+                                             method_allowlist=method_allowlist,
+                                             resource_messages=resource_messages,
+                                             services_in_proto=services_in_proto)
+
+    def prune_messages_for_selective_generation(self, *,
+                                           address_allowlist: Set['metadata.Address']) -> Optional['Proto']:
+        """Returns a truncated version of this Proto.
+
+        Only the services, messages, and enums contained in the allowlist
+        of visited addresses are included in the returned object. If there
+        are no services, messages, or enums left, and no file level resources,
+        return None.
+
+        Args:
+            address_allowlist (Set[metadata.Address]): A set of allowlisted metadata.Address
+                objects to filter against. Objects with addresses not the allowlist will be
+                removed from the returned Proto.
+        Returns:
+            Optional[Proto]: A truncated version of this proto. If there are no services, messages,
+                or enums left after the truncation process and there are no file level resources,
+                returns None.
+        """
+        # Once the address allowlist has been created, it suffices to only
+        # prune items at 2 different levels to truncate the Proto object:
+        #
+        #   1. At the Proto level, we remove unnecessary services, messages,
+        #      and enums.
+        #   2. For allowlisted services, at the Service level, we remove
+        #      non-allowlisted methods.
+        services = {
+            k: v.prune_messages_for_selective_generation(
+                address_allowlist=address_allowlist)
+            for k, v in self.services.items()
+            if v.meta.address in address_allowlist
+        }
+
+        all_messages = {
+            k: v
+            for k, v in self.all_messages.items()
+            if v.ident in address_allowlist
+        }
+
+        all_enums = {
+            k: v
+            for k, v in self.all_enums.items()
+            if v.ident in address_allowlist
+        }
+
+        if not services and not all_messages and not all_enums:
+            return None
+
+        return dataclasses.replace(
+            self,
+            services=services,
+            all_messages=all_messages,
+            all_enums=all_enums
+        )
+
 
 @dataclasses.dataclass(frozen=True)
 class API:
@@ -365,10 +454,52 @@ def disambiguate_keyword_sanitize_fname(
             ignore_unknown_fields=True
         )
 
-        # Done; return the API.
-        return cls(naming=naming,
-                   all_protos=protos,
-                   service_yaml_config=service_yaml_config)
+        # Third pass for various selective GAPIC settings; these require
+        # settings in the service.yaml and so we build the API object
+        # before doing another pass.
+        api = cls(naming=naming,
+                  all_protos=protos,
+                  service_yaml_config=service_yaml_config)
+
+        if package in api.all_library_settings:
+            selective_gapic_methods = set(
+                api.all_library_settings[package].python_settings.common.selective_gapic_generation.methods
+            )
+            if selective_gapic_methods:
+
+                all_resource_messages = collections.ChainMap(
+                    *(proto.resource_messages for proto in protos.values())
+                )
+
+                # Prepare a list of addresses to include in selective generation,
+                # then prune each Proto object. We look at metadata.Addresses, not objects, because
+                # objects that refer to the same thing in the proto are different Python objects
+                # in memory.
+                address_allowlist: Set['metadata.Address'] = set([])
+                for proto in api.protos.values():
+                    proto.add_to_address_allowlist(address_allowlist=address_allowlist,
+                                                   method_allowlist=selective_gapic_methods,
+                                                   resource_messages=all_resource_messages)
+
+                # The list of explicitly allow-listed protos to generate, plus all
+                # the proto dependencies regardless of the allow-list.
+                new_all_protos = {}
+
+                # We only prune services/messages/enums from protos that are not dependencies.
+                for name, proto in api.all_protos.items():
+                    if name not in api.protos:
+                        new_all_protos[name] = proto
+                    else:
+                        proto_to_generate = proto.prune_messages_for_selective_generation(
+                            address_allowlist=address_allowlist)
+                        if proto_to_generate:
+                            new_all_protos[name] = proto_to_generate
+
+                api = cls(naming=naming,
+                        all_protos=new_all_protos,
+                        service_yaml_config=service_yaml_config)
+
+        return api
 
     @cached_property
     def enums(self) -> Mapping[str, wrappers.EnumType]:
@@ -743,6 +874,21 @@ def enforce_valid_library_settings(
                 continue
             versions_seen.add(library_settings.version)
 
+            # Check to see if selective gapic generation methods are valid.
+            selective_gapic_errors = {}
+            for method_name in library_settings.python_settings.common.selective_gapic_generation.methods:
+                if method_name not in self.all_methods:
+                    selective_gapic_errors[method_name] = "Method does not exist."
+                elif not method_name.startswith(library_settings.version):
+                    selective_gapic_errors[method_name] = "Mismatched version for method."
+
+            if selective_gapic_errors:
+                all_errors[library_settings.version] = [
+                    {
+                        "selective_gapic_generation": selective_gapic_errors,
+                    }
+                ]
+
         if all_errors:
             raise ClientLibrarySettingsError(yaml.dump(all_errors))