fix: restore messages not attached to rpc in selective_gapic_generation (#16951)

Towards b/507893758, b/507889482, b/501132869, b/503326310

Author: parthea

packages/gapic-generator/gapic/schema/api.py

@@ -265,6 +265,7 @@ def add_to_address_allowlist(
         address_allowlist: Set["metadata.Address"],
         method_allowlist: Set[str],
         resource_messages: Dict[str, "wrappers.MessageType"],
+        services_in_proto: Dict[str, "wrappers.Service"],
     ) -> None:
         """Adds to the set of Addresses of wrapper objects to be included in selective GAPIC generation.
 
@@ -281,15 +282,12 @@ def add_to_address_allowlist(
                 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
         """
-        # 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,
@@ -298,75 +296,60 @@ def add_to_address_allowlist(
                 services_in_proto=services_in_proto,
             )
 
-    def prune_messages_for_selective_generation(
-        self, *, address_allowlist: Set["metadata.Address"]
+    def with_selective_generation(
+        self,
+        *,
+        generate_omitted_as_internal: bool,
+        public_methods: Set[str],
+        excluded_addresses: 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.
+        """Returns a version of this Proto for selective generation.
 
         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.
+            generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
+            public_methods (Set[str]): The set of fully-qualified method names to keep as public.
+            excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
+
         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.
+            Optional[Proto]: A version of this Proto with services/methods filtered.
+                Returns None if the Proto becomes empty and generate_omitted_as_internal is False.
         """
-        # Once the address allowlist has been created, it suffices to only
-        # prune items at 2 different levels to truncate the Proto object:
+        services = {}
+        for k, v in self.services.items():
+            new_v = v.with_selective_generation(
+                generate_omitted_as_internal=generate_omitted_as_internal,
+                public_methods=public_methods,
+                excluded_addresses=excluded_addresses)
+            if new_v:
+                services[k] = new_v
+
+        # We only prune messages/enums from protos that are not dependencies.
+        # A message or enum is excluded IF AND ONLY IF:
+        #   1. It is a top-level request or response message for an omitted RPC.
+        #   2. It is NOT reachable from any publicly allowed RPC.
         #
-        #   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
-        }
-
+        # This ensures that shared messages, messages not attached to any RPC,
+        # and messages reachable via other paths (like LRO response types) are KEPT.
         all_messages = {
-            k: v for k, v in self.all_messages.items() if v.ident in address_allowlist
+            k: v for k, v in self.all_messages.items() if v.ident not in excluded_addresses
         }
 
         all_enums = {
-            k: v for k, v in self.all_enums.items() if v.ident in address_allowlist
+            k: v for k, v in self.all_enums.items() if v.ident not in excluded_addresses
         }
 
-        if not services and not all_messages and not all_enums:
+        # If the proto becomes empty after pruning, we return None to signal
+        # that it should be excluded from generation.
+        if not generate_omitted_as_internal and 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
+            self,
+            services=services,
+            all_messages=all_messages,
+            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:
@@ -532,32 +515,79 @@ def disambiguate_keyword_sanitize_fname(
 
                 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
+                        new_all_protos[name] = proto.with_selective_generation(
+                            generate_omitted_as_internal=True,
+                            public_methods=selective_gapic_methods,
+                            excluded_addresses=set([]),
                         )
                 else:
-                    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():
+                    all_resource_messages = dict(collections.ChainMap(
+                        *(proto.resource_messages for proto in api.all_protos.values())
+                    ))
+
+                    # Create a global map of services to support cross-proto lookup
+                    # for extended LROs.
+                    #
+                    # Note: This is keyed by the fully qualified proto name (as a string)
+                    # to ensure compatibility with Address.resolve() lookups in wrappers.py.
+                    all_services: Dict[str, wrappers.Service] = {}
+                    for p in api.all_protos.values():
+                        for s in p.services.values():
+                            all_services[s.meta.address.proto] = s
+
+                    # Calculate addresses of omitted RPCs and their top-level request/response messages.
+                    # These are "candidates" for exclusion.
+                    #
+                    # We only consider top-level request/response messages of omitted RPCs as
+                    # candidates for exclusion. This is conservative: it ensures that:
+                    #   - Messages NOT used by any RPC are KEPT (e.g. for user convenience).
+                    #   - Messages shared between an omitted and a kept RPC are KEPT.
+                    #   - Messages reachable from a kept RPC but NOT as a top-level request/response
+                    #     (e.g. nested messages) are KEPT.
+                    candidate_excluded_addresses: Set["metadata.Address"] = set([])
+                    for proto in api.all_protos.values():
+                        for service in proto.services.values():
+                            for method in service.methods.values():
+                                if method.ident.proto not in selective_gapic_methods:
+                                    # Candidate for exclusion: the method itself and its direct request/response types.
+                                    candidate_excluded_addresses.add(method.meta.address)
+                                    candidate_excluded_addresses.add(method.input.ident)
+                                    candidate_excluded_addresses.add(method.output.ident)
+
+                                    # If this is an LRO, add its response and metadata types to candidates.
+                                    if method.lro:
+                                        candidate_excluded_addresses.add(method.lro.response_type.ident)
+                                        candidate_excluded_addresses.add(method.lro.metadata_type.ident)
+
+                                    # If this is an extended LRO, add its request and operation types to candidates.
+                                    if method.extended_lro:
+                                        candidate_excluded_addresses.add(method.extended_lro.request_type.ident)
+                                        candidate_excluded_addresses.add(method.extended_lro.operation_type.ident)
+
+                    # Calculate publicly reachable addresses (API-wide).
+                    # This includes all types reachable from the allowlisted (public) methods.
+                    public_rpc_addresses: Set["metadata.Address"] = set([])
+                    for proto in api.all_protos.values():
                         proto.add_to_address_allowlist(
-                            address_allowlist=address_allowlist,
+                            address_allowlist=public_rpc_addresses,
                             method_allowlist=selective_gapic_methods,
                             resource_messages=all_resource_messages,
+                            services_in_proto=all_services,
                         )
 
-                    # We only prune services/messages/enums from protos that are not dependencies.
+                    # Addresses to exclude: those that are candidates for exclusion but NOT
+                    # reachable from any PUBLIC RPC.
+                    #
+                    # This set difference effectively "vets" the candidates. If a candidate
+                    # message is actually reachable from a public RPC, it's removed from
+                    # the exclusion list.
+                    excluded_addresses = candidate_excluded_addresses - public_rpc_addresses
+
                     for name, proto in api.protos.items():
-                        proto_to_generate = (
-                            proto.prune_messages_for_selective_generation(
-                                address_allowlist=address_allowlist
-                            )
+                        proto_to_generate = proto.with_selective_generation(
+                            generate_omitted_as_internal=False,
+                            public_methods=selective_gapic_methods,
+                            excluded_addresses=excluded_addresses,
                         )
                         if proto_to_generate:
                             new_all_protos[name] = proto_to_generate

packages/gapic-generator/gapic/schema/wrappers.py

@@ -2032,7 +2032,9 @@ def add_to_address_allowlist(
             # the allowlist, as it might not have been specified by
             # the methods under selective_gapic_generation.
             # We assume that the operation service lives in the same proto file as this one.
-            operation_service = services_in_proto[self.operation_service]
+            operation_service = services_in_proto[
+                self.meta.address.resolve(self.operation_service)
+            ]
             address_allowlist.add(operation_service.meta.address)
             operation_service.operation_polling_method.add_to_address_allowlist(
                 address_allowlist=address_allowlist,
@@ -2055,26 +2057,39 @@ def add_to_address_allowlist(
             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`).
+    def with_selective_generation(
+        self,
+        *,
+        generate_omitted_as_internal: bool,
+        public_methods: Set[str],
+        excluded_addresses: Set["metadata.Address"],
+    ) -> Optional["Method"]:
+        """Returns a version of this Method for selective generation.
 
         Args:
-            public_methods (Set[str]): An allowlist of fully-qualified method names.
-                Methods not in this allowlist will be marked as internal.
+            generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
+            public_methods (Set[str]): The set of fully-qualified method names to keep as public.
+            excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
+
         Returns:
-            Service: A version of this `Service` with `Method` objects corresponding to methods
-                not in `public_methods` marked as internal.
+            Optional[Method]: The method, possibly marked as internal, or None if it should be removed.
         """
         if self.ident.proto in public_methods:
             return self
 
-        return dataclasses.replace(
-            self,
-            is_internal=True,
-        )
+        # Not public.
+        # We mark it as internal if either:
+        #   1. generate_omitted_as_internal is set in selective_gapic_generation.
+        #   2. The method is NOT in excluded_addresses, which means it is reachable
+        #      from some public method (e.g. as a polling method for an extended LRO).
+        if generate_omitted_as_internal or self.meta.address not in excluded_addresses:
+            return dataclasses.replace(
+                self,
+                is_internal=True,
+            )
+        else:
+            return None
+
 
 
 @dataclasses.dataclass(frozen=True)
@@ -2463,45 +2478,33 @@ def add_to_address_allowlist(
                     services_in_proto=services_in_proto,
                 )
 
-    def prune_messages_for_selective_generation(
-        self, *, address_allowlist: Set["metadata.Address"]
-    ) -> "Service":
-        """Returns a truncated version of this Service.
-
-        Only the methods, messages, and enums contained in the address allowlist
-        are included in the returned object.
+    def with_selective_generation(
+        self,
+        *,
+        generate_omitted_as_internal: bool,
+        public_methods: Set[str],
+        excluded_addresses: Set["metadata.Address"],
+    ) -> Optional["Service"]:
+        """Returns a version of this Service for selective generation.
 
         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:
-            Service: A truncated version of this proto.
-        """
-        return dataclasses.replace(
-            self,
-            methods={
-                k: v 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`).
+            generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
+            public_methods (Set[str]): The set of fully-qualified method names to keep as public.
+            excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
 
-        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.
+            Optional[Service]: The service with filtered methods, or None if it should be removed.
         """
-        return dataclasses.replace(
-            self,
-            methods={
-                k: v.with_internal_methods(public_methods=public_methods)
-                for k, v in self.methods.items()
-            },
-        )
+        methods = {}
+        for k, v in self.methods.items():
+            new_v = v.with_selective_generation(
+                generate_omitted_as_internal=generate_omitted_as_internal,
+                public_methods=public_methods,
+                excluded_addresses=excluded_addresses)
+            if new_v:
+                methods[k] = new_v
+
+        if not generate_omitted_as_internal and not methods:
+            return None
+
+        return dataclasses.replace(self, methods=methods)