feat: Added support for internal methods in selective GAPIC generation (#2325)

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

Author: gkevinzheng

DEVELOPMENT.md

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

gapic/generator/generator.py

@@ -59,6 +59,7 @@ def __init__(self, opts: Options) -> None:
 
         # Add filters which templates require.
         self._env.filters["rst"] = utils.rst
+        self._env.filters["make_private"] = utils.make_private
         self._env.filters["snake_case"] = utils.to_snake_case
         self._env.filters["camel_case"] = utils.to_camel_case
         self._env.filters["sort_lines"] = utils.sort_lines

gapic/samplegen/samplegen.py

@@ -1036,6 +1036,12 @@ def generate_sample_specs(api_schema: api.API, *, opts) -> Generator[Dict[str, A
                 # Region Tag Format:
                 # [{START|END} ${apishortname}_${apiVersion}_generated_${serviceName}_${rpcName}_{sync|async|rest}]
                 region_tag = f"{api_short_name}_{api_version}_generated_{service_name}_{rpc_name}_{sync_or_async}"
+
+                # We assume that the only methods that start with an underscore are internal methods.
+                is_internal = rpc_name.startswith("_")
+                if is_internal:
+                    region_tag += "_internal"
+
                 spec = {
                     "rpc": rpc_name,
                     "transport": transport,

gapic/schema/api.py

@@ -326,6 +326,26 @@ def prune_messages_for_selective_generation(self, *,
             all_enums=all_enums
         )
 
+    def with_internal_methods(self, *, public_methods: Set[str]) -> 'Proto':
+        """Returns a version of this Proto with some Methods marked as internal.
+
+        The methods not in the public_methods set will be marked as internal and
+        services containing these methods will also be marked as internal by extension.
+        (See :meth:`Service.is_internal` for more details).
+
+        Args:
+            public_methods (Set[str]): An allowlist of fully-qualified method names.
+                Methods not in this allowlist will be marked as internal.
+        Returns:
+            Proto: A version of this Proto with Method objects corresponding to methods
+                not in `public_methods` marked as internal.
+        """
+        services = {
+            k: v.with_internal_methods(public_methods=public_methods)
+            for k, v in self.services.items()
+        }
+        return dataclasses.replace(self, services=services)
+
 
 @dataclasses.dataclass(frozen=True)
 class API:
@@ -462,34 +482,43 @@ def disambiguate_keyword_sanitize_fname(
                   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)
+            selective_gapic_settings = api.all_library_settings[package].python_settings.\
+                common.selective_gapic_generation
 
+            selective_gapic_methods = set(selective_gapic_settings.methods)
+            if selective_gapic_methods:
                 # The list of explicitly allow-listed protos to generate, plus all
                 # the proto dependencies regardless of the allow-list.
-                new_all_protos = {}
+                #
+                # Both selective GAPIC generation settings (omitting + internal) only alter
+                # protos that are not dependencies, so we iterate over api.all_protos and copy
+                # all dependencies as is here.
+                new_all_protos = {
+                    k: v for k, v in api.all_protos.items()
+                    if k not in api.protos
+                }
+
+                if selective_gapic_settings.generate_omitted_as_internal:
+                    for name, proto in api.protos.items():
+                        new_all_protos[name] = proto.with_internal_methods(
+                            public_methods=selective_gapic_methods)
+                else:
+                    all_resource_messages = collections.ChainMap(
+                        *(proto.resource_messages for proto in protos.values())
+                    )
 
-                # 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:
+                    # 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)
+
+                    # We only prune services/messages/enums from protos that are not dependencies.
+                    for name, proto in api.protos.items():
                         proto_to_generate = proto.prune_messages_for_selective_generation(
                             address_allowlist=address_allowlist)
                         if proto_to_generate:

gapic/schema/wrappers.py

@@ -53,6 +53,7 @@
 from gapic import utils
 from gapic.schema import metadata
 from gapic.utils import uri_sample
+from gapic.utils import make_private
 
 
 @dataclasses.dataclass(frozen=True)
@@ -900,6 +901,14 @@ def add_to_address_allowlist(self, *,
 
         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.
+        Returns:
+            None
         """
         address_allowlist.add(self.ident)
 
@@ -1014,6 +1023,18 @@ def add_to_address_allowlist(self, *,
 
         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.
+            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
         """
 
         self.request_type.add_to_address_allowlist(
@@ -1061,6 +1082,18 @@ def add_to_address_allowlist(self, *,
 
         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.
+            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
         """
         self.response_type.add_to_address_allowlist(
             address_allowlist=address_allowlist,
@@ -1362,6 +1395,7 @@ class Method:
     method_pb: descriptor_pb2.MethodDescriptorProto
     input: MessageType
     output: MessageType
+    is_internal: bool = False
     lro: Optional[OperationInfo] = dataclasses.field(default=None)
     extended_lro: Optional[ExtendedOperationInfo] = dataclasses.field(
         default=None)
@@ -1403,6 +1437,10 @@ def transport_safe_name(self) -> str:
     def is_operation_polling_method(self):
         return self.output.is_extended_operation and self.options.Extensions[ex_ops_pb2.operation_polling_method]
 
+    @utils.cached_property
+    def name(self):
+        return make_private(self.method_pb.name) if self.is_internal else self.method_pb.name
+
     @utils.cached_property
     def client_output(self):
         return self._client_output(enable_asyncio=False)
@@ -1814,6 +1852,22 @@ def add_to_address_allowlist(self, *,
 
         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.
+            services_in_proto (Dict[str, wrappers.Service]): A dictionary mapping the names of Service
+                objects in the proto containing this method to the Service objects. This is necessary
+                for traversing the operation service in the case of extended LROs.
+        Returns:
+            None
         """
 
         address_allowlist.add(self.ident)
@@ -1851,6 +1905,27 @@ def add_to_address_allowlist(self, *,
             resource_messages=resource_messages,
         )
 
+    def with_internal_methods(self, *, public_methods: Set[str]) -> 'Method':
+        """Returns a version of this ``Method`` marked as internal
+
+        The methods not in the public_methods set will be marked as internal and
+        this ``Service`` will as well by extension (see :meth:`Service.is_internal`).
+
+        Args:
+            public_methods (Set[str]): An allowlist of fully-qualified method names.
+                Methods not in this allowlist will be marked as internal.
+        Returns:
+            Service: A version of this `Service` with `Method` objects corresponding to methods
+                not in `public_methods` marked as internal.
+        """
+        if self.ident.proto in public_methods:
+            return self
+
+        return dataclasses.replace(
+            self,
+            is_internal=True,
+        )
+
 
 @dataclasses.dataclass(frozen=True)
 class CommonResource:
@@ -1928,7 +2003,7 @@ def __getattr__(self, name):
     @property
     def client_name(self) -> str:
         """Returns the name of the generated client class"""
-        return self.name + "Client"
+        return ("Base" if self.is_internal else "") + self.name + "Client"
 
     @property
     def client_package_version(self) -> str:
@@ -1937,7 +2012,7 @@ def client_package_version(self) -> str:
     @property
     def async_client_name(self) -> str:
         """Returns the name of the generated AsyncIO client class"""
-        return self.name + "AsyncClient"
+        return ("Base" if self.is_internal else "") + self.name + "AsyncClient"
 
     @property
     def transport_name(self):
@@ -2143,6 +2218,10 @@ def operation_polling_method(self) -> Optional[Method]:
             None
         )
 
+    @utils.cached_property
+    def is_internal(self) -> bool:
+        return any(m.is_internal for m in self.methods.values())
+
     def with_context(self, *,
                      collisions: Set[str],
                      visited_messages: Optional[Set["MessageType"]] = None,
@@ -2224,3 +2303,32 @@ def prune_messages_for_selective_generation(self, *,
                 for k, v in self.methods.items() if v.ident in address_allowlist
             }
         )
+
+    def with_internal_methods(self, *,
+                              public_methods: Set[str]) -> 'Service':
+        """Returns a version of this ``Service`` with some Methods marked as internal.
+
+        The methods not in the public_methods set will be marked as internal and
+        this ``Service`` will as well by extension (see :meth:`Service.is_internal`).
+
+        Args:
+            public_methods (Set[str]): An allowlist of fully-qualified method names.
+                Methods not in this allowlist will be marked as internal.
+        Returns:
+            Service: A version of this `Service` with `Method` objects corresponding to methods
+                not in `public_methods` marked as internal.
+        """
+
+        # Internal methods need to be keyed with underscore prefixed method names
+        # (e.g. google.Service.Method -> google.Service._Method) in order for
+        # samplegen to work properly.
+        return dataclasses.replace(
+            self,
+            methods={
+                meth.name: meth
+                for meth in (
+                    meth.with_internal_methods(public_methods=public_methods)
+                    for meth in self.methods.values()
+                )
+            }
+        )

gapic/templates/%namespace/%name_%version/%sub/services/%service/_shared_macros.j2

@@ -452,9 +452,9 @@ class {{ async_method_name_prefix }}{{ service.name }}RestInterceptor:
 
 @property
 def {{ name|snake_case }}(self):
-    return self._{{ name }}(self._session, self._host, self._interceptor) # type: ignore
+    return self.{{ name|make_private }}(self._session, self._host, self._interceptor) # type: ignore
 
-class _{{ name }}(_Base{{ service.name }}RestTransport._Base{{name}}, {{ async_method_name_prefix }}{{service.name}}RestStub):
+class {{ name|make_private }}(_Base{{ service.name }}RestTransport._Base{{name}}, {{ async_method_name_prefix }}{{service.name}}RestStub):
     def __hash__(self):
         return hash("{{ async_method_name_prefix }}{{ service.name }}RestTransport.{{ name }}")
 

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

@@ -366,7 +366,7 @@ class {{ service.name }}Transport(abc.ABC):
 
     {% for operations_service in api.get_extended_operations_services(service)|sort(attribute="name") %}
     @property
-    def _{{ operations_service.client_name|snake_case }}(self) -> {{ operations_service.name|snake_case }}.{{ operations_service.client_name }}:
+    def {{ operations_service.client_name|snake_case|make_private }}(self) -> {{ operations_service.name|snake_case }}.{{ operations_service.client_name }}:
         ex_op_service = self._extended_operations_services.get("{{ operations_service.name|snake_case }}")
         if not ex_op_service:
             ex_op_service = {{ operations_service.name|snake_case }}.{{ operations_service.client_name }}(

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

@@ -198,7 +198,7 @@ class {{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
 
     {% endif %}{# service.has_lro #}
     {% for method in service.methods.values()|sort(attribute="name") %}
-    class _{{method.name}}(_Base{{ service.name }}RestTransport._Base{{method.name}}, {{service.name}}RestStub):
+    class {{ method.name|make_private }}(_Base{{ service.name }}RestTransport._Base{{method.name}}, {{service.name}}RestStub):
         def __hash__(self):
             return hash("{{service.name}}RestTransport.{{method.name}}")
         {% if method.http_options and not method.client_streaming %}

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

@@ -159,7 +159,7 @@ class Async{{service.name}}RestTransport(_Base{{ service.name }}RestTransport):
     {{ shared_macros.wrap_async_method_macro()|indent(4) }}
 
     {% for method in service.methods.values()|sort(attribute="name") %}
-    class _{{method.name}}(_Base{{ service.name }}RestTransport._Base{{method.name}}, Async{{service.name}}RestStub):
+    class {{ method.name|make_private }}(_Base{{ service.name }}RestTransport._Base{{method.name}}, Async{{service.name}}RestStub):
         def __hash__(self):
             return hash("Async{{service.name}}RestTransport.{{method.name}}")
 

gapic/utils/__init__.py

@@ -20,6 +20,7 @@
 from gapic.utils.code import empty
 from gapic.utils.code import nth
 from gapic.utils.code import partition
+from gapic.utils.code import make_private
 from gapic.utils.doc import doc
 from gapic.utils.filename import to_valid_filename
 from gapic.utils.filename import to_valid_module_name
@@ -38,6 +39,7 @@
     'empty',
     'is_msg_field_pb',
     'is_str_field_pb',
+    'make_private',
     'nth',
     'Options',
     'partition',