diff --git a/packages/gapic-generator/gapic/schema/api.py b/packages/gapic-generator/gapic/schema/api.py index 7d982eef3aaa..797eb5718070 100644 --- a/packages/gapic-generator/gapic/schema/api.py +++ b/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 diff --git a/packages/gapic-generator/gapic/schema/wrappers.py b/packages/gapic-generator/gapic/schema/wrappers.py index 13fd99729320..2374b0def6bb 100644 --- a/packages/gapic-generator/gapic/schema/wrappers.py +++ b/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) diff --git a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis/__init__.py b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis/__init__.py index a053f78a5382..fade9bb95ef5 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis/__init__.py +++ b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis/__init__.py @@ -23,33 +23,45 @@ from google.cloud.redis_v1.types.cloud_redis import CreateInstanceRequest from google.cloud.redis_v1.types.cloud_redis import DeleteInstanceRequest +from google.cloud.redis_v1.types.cloud_redis import GcsDestination +from google.cloud.redis_v1.types.cloud_redis import GcsSource from google.cloud.redis_v1.types.cloud_redis import GetInstanceRequest +from google.cloud.redis_v1.types.cloud_redis import InputConfig from google.cloud.redis_v1.types.cloud_redis import Instance from google.cloud.redis_v1.types.cloud_redis import ListInstancesRequest from google.cloud.redis_v1.types.cloud_redis import ListInstancesResponse +from google.cloud.redis_v1.types.cloud_redis import LocationMetadata from google.cloud.redis_v1.types.cloud_redis import MaintenancePolicy from google.cloud.redis_v1.types.cloud_redis import MaintenanceSchedule from google.cloud.redis_v1.types.cloud_redis import NodeInfo from google.cloud.redis_v1.types.cloud_redis import OperationMetadata +from google.cloud.redis_v1.types.cloud_redis import OutputConfig from google.cloud.redis_v1.types.cloud_redis import PersistenceConfig from google.cloud.redis_v1.types.cloud_redis import TlsCertificate from google.cloud.redis_v1.types.cloud_redis import UpdateInstanceRequest from google.cloud.redis_v1.types.cloud_redis import WeeklyMaintenanceWindow +from google.cloud.redis_v1.types.cloud_redis import ZoneMetadata __all__ = ('CloudRedisClient', 'CloudRedisAsyncClient', 'CreateInstanceRequest', 'DeleteInstanceRequest', + 'GcsDestination', + 'GcsSource', 'GetInstanceRequest', + 'InputConfig', 'Instance', 'ListInstancesRequest', 'ListInstancesResponse', + 'LocationMetadata', 'MaintenancePolicy', 'MaintenanceSchedule', 'NodeInfo', 'OperationMetadata', + 'OutputConfig', 'PersistenceConfig', 'TlsCertificate', 'UpdateInstanceRequest', 'WeeklyMaintenanceWindow', + 'ZoneMetadata', ) diff --git a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/__init__.py b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/__init__.py index 6b792cc004dc..f9ca3ee685f8 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/__init__.py +++ b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/__init__.py @@ -28,18 +28,24 @@ from .types.cloud_redis import CreateInstanceRequest from .types.cloud_redis import DeleteInstanceRequest +from .types.cloud_redis import GcsDestination +from .types.cloud_redis import GcsSource from .types.cloud_redis import GetInstanceRequest +from .types.cloud_redis import InputConfig from .types.cloud_redis import Instance from .types.cloud_redis import ListInstancesRequest from .types.cloud_redis import ListInstancesResponse +from .types.cloud_redis import LocationMetadata from .types.cloud_redis import MaintenancePolicy from .types.cloud_redis import MaintenanceSchedule from .types.cloud_redis import NodeInfo from .types.cloud_redis import OperationMetadata +from .types.cloud_redis import OutputConfig from .types.cloud_redis import PersistenceConfig from .types.cloud_redis import TlsCertificate from .types.cloud_redis import UpdateInstanceRequest from .types.cloud_redis import WeeklyMaintenanceWindow +from .types.cloud_redis import ZoneMetadata if hasattr(api_core, "check_python_version") and hasattr(api_core, "check_dependency_versions"): # pragma: NO COVER api_core.check_python_version("google.cloud.redis_v1") # type: ignore @@ -121,16 +127,22 @@ def _get_version(dependency_name): 'CloudRedisClient', 'CreateInstanceRequest', 'DeleteInstanceRequest', +'GcsDestination', +'GcsSource', 'GetInstanceRequest', +'InputConfig', 'Instance', 'ListInstancesRequest', 'ListInstancesResponse', +'LocationMetadata', 'MaintenancePolicy', 'MaintenanceSchedule', 'NodeInfo', 'OperationMetadata', +'OutputConfig', 'PersistenceConfig', 'TlsCertificate', 'UpdateInstanceRequest', 'WeeklyMaintenanceWindow', +'ZoneMetadata', ) diff --git a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/__init__.py b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/__init__.py index 5b4432955218..e1ae30d8179e 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/__init__.py +++ b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/__init__.py @@ -16,33 +16,45 @@ from .cloud_redis import ( CreateInstanceRequest, DeleteInstanceRequest, + GcsDestination, + GcsSource, GetInstanceRequest, + InputConfig, Instance, ListInstancesRequest, ListInstancesResponse, + LocationMetadata, MaintenancePolicy, MaintenanceSchedule, NodeInfo, OperationMetadata, + OutputConfig, PersistenceConfig, TlsCertificate, UpdateInstanceRequest, WeeklyMaintenanceWindow, + ZoneMetadata, ) __all__ = ( 'CreateInstanceRequest', 'DeleteInstanceRequest', + 'GcsDestination', + 'GcsSource', 'GetInstanceRequest', + 'InputConfig', 'Instance', 'ListInstancesRequest', 'ListInstancesResponse', + 'LocationMetadata', 'MaintenancePolicy', 'MaintenanceSchedule', 'NodeInfo', 'OperationMetadata', + 'OutputConfig', 'PersistenceConfig', 'TlsCertificate', 'UpdateInstanceRequest', 'WeeklyMaintenanceWindow', + 'ZoneMetadata', ) diff --git a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/cloud_redis.py b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/cloud_redis.py index a6dd39f7c72d..64d88e277c10 100755 --- a/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/cloud_redis.py +++ b/packages/gapic-generator/tests/integration/goldens/redis_selective/google/cloud/redis_v1/types/cloud_redis.py @@ -41,7 +41,13 @@ 'CreateInstanceRequest', 'UpdateInstanceRequest', 'DeleteInstanceRequest', + 'GcsSource', + 'InputConfig', + 'GcsDestination', + 'OutputConfig', 'OperationMetadata', + 'LocationMetadata', + 'ZoneMetadata', 'TlsCertificate', }, ) @@ -920,6 +926,79 @@ class DeleteInstanceRequest(proto.Message): ) +class GcsSource(proto.Message): + r"""The Cloud Storage location for the input content + + Attributes: + uri (str): + Required. Source data URI. (e.g. + 'gs://my_bucket/my_object'). + """ + + uri: str = proto.Field( + proto.STRING, + number=1, + ) + + +class InputConfig(proto.Message): + r"""The input content + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + gcs_source (google.cloud.redis_v1.types.GcsSource): + Google Cloud Storage location where input + content is located. + + This field is a member of `oneof`_ ``source``. + """ + + gcs_source: 'GcsSource' = proto.Field( + proto.MESSAGE, + number=1, + oneof='source', + message='GcsSource', + ) + + +class GcsDestination(proto.Message): + r"""The Cloud Storage location for the output content + + Attributes: + uri (str): + Required. Data destination URI (e.g. + 'gs://my_bucket/my_object'). Existing files will be + overwritten. + """ + + uri: str = proto.Field( + proto.STRING, + number=1, + ) + + +class OutputConfig(proto.Message): + r"""The output content + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + gcs_destination (google.cloud.redis_v1.types.GcsDestination): + Google Cloud Storage destination for output + content. + + This field is a member of `oneof`_ ``destination``. + """ + + gcs_destination: 'GcsDestination' = proto.Field( + proto.MESSAGE, + number=1, + oneof='destination', + message='GcsDestination', + ) + + class OperationMetadata(proto.Message): r"""Represents the v1 metadata of the long-running operation. @@ -973,6 +1052,36 @@ class OperationMetadata(proto.Message): ) +class LocationMetadata(proto.Message): + r"""This location metadata represents additional configuration options + for a given location where a Redis instance may be created. All + fields are output only. It is returned as content of the + ``google.cloud.location.Location.metadata`` field. + + Attributes: + available_zones (MutableMapping[str, google.cloud.redis_v1.types.ZoneMetadata]): + Output only. The set of available zones in the location. The + map is keyed by the lowercase ID of each zone, as defined by + GCE. These keys can be specified in ``location_id`` or + ``alternative_location_id`` fields when creating a Redis + instance. + """ + + available_zones: MutableMapping[str, 'ZoneMetadata'] = proto.MapField( + proto.STRING, + proto.MESSAGE, + number=1, + message='ZoneMetadata', + ) + + +class ZoneMetadata(proto.Message): + r"""Defines specific information for a particular zone. Currently + empty and reserved for future use only. + + """ + + class TlsCertificate(proto.Message): r"""TlsCertificate Resource diff --git a/packages/gapic-generator/tests/unit/schema/test_api.py b/packages/gapic-generator/tests/unit/schema/test_api.py index cb3661ee580a..13ca7a009c86 100644 --- a/packages/gapic-generator/tests/unit/schema/test_api.py +++ b/packages/gapic-generator/tests/unit/schema/test_api.py @@ -3346,6 +3346,8 @@ def test_selective_gapic_api_build(): assert "google.example.v1.GetFooResponse" in api_schema.messages assert "google.example.v1.DeleteFooRequest" not in api_schema.messages assert "google.example.v1.DeleteFooResponse" not in api_schema.messages + # Baz is used by DeleteFooRequest, which is omitted, but Baz itself should be kept. + assert "google.example.v1.common.Baz" in api_schema.messages assert "google.example.v1.FooService" in api_schema.services assert len(api_schema.enums) == 0 assert api_schema.protos["foo.proto"].python_modules == ( @@ -3365,8 +3367,9 @@ def test_selective_gapic_api_build(): sub = api_schema.subpackages["common"] assert len(sub.protos) == 1 assert "google.example.v1.common.Bar" in sub.messages - assert "google.example.v1.common.Baz" not in sub.messages - + # Baz is reachable from DeleteFooRequest, which is omitted, but + # Baz itself should be part of the generation. + assert "google.example.v1.common.Baz" in sub.messages # Establish that methods have been truncated assert "google.example.v1.FooService.GetFoo" in api_schema.all_methods assert "google.example.v1.FooService.DeleteFoo" not in api_schema.all_methods @@ -3476,11 +3479,17 @@ def test_selective_gapic_api_build_with_lro(): assert "google.example.v1.AsyncCreateFooResponse" in api_schema.messages assert "google.example.v1.AsyncCreateFooMetadata" in api_schema.messages - assert "google.example.v1.Bar" not in api_schema.messages + # Top-level request/response messages for omitted RPCs are still omitted. assert "google.example.v1.AsyncCreateBarRequest" not in api_schema.messages + # AsyncCreateBarResponse and AsyncCreateBarMetadata are LRO response/metadata types + # for an omitted RPC, so they should be omitted. assert "google.example.v1.AsyncCreateBarResponse" not in api_schema.messages assert "google.example.v1.AsyncCreateBarMetadata" not in api_schema.messages + # Bar is reachable from AsyncCreateBarMetadata, and since Bar itself is NOT + # a top-level request/response message of an omitted RPC, it should be KEPT. + assert "google.example.v1.Bar" in api_schema.messages + def test_selective_gapic_api_build_remove_unnecessary_services(): # Put together a couple of minimal protos. @@ -3547,7 +3556,8 @@ def test_selective_gapic_api_build_remove_unnecessary_services(): assert "google.example.v1.GetFooRequest" in api_schema.messages assert "google.example.v1.GetFooResponse" in api_schema.messages - assert "google.example.v1.Bar" not in api_schema.messages + # Bar is used by GetBarRequest (omitted), so Bar itself should be kept. + assert "google.example.v1.Bar" in api_schema.messages assert "google.example.v1.GetBarRequest" not in api_schema.messages assert "google.example.v1.GetBarResponse" not in api_schema.messages @@ -3606,8 +3616,13 @@ def test_selective_gapic_api_build_remove_unnecessary_proto_files(): name="GetBarRequest", fields=( make_field_pb2( - name="bar", + name="bar_local", number=1, + type_name=".google.example.v1.Bar", + ), + make_field_pb2( + name="bar_common", + number=2, type_name=".google.example.v1.bar_common.Bar", ), ), @@ -3641,7 +3656,9 @@ def test_selective_gapic_api_build_remove_unnecessary_proto_files(): assert "google.example.v1.GetFooRequest" in api_schema.messages assert "google.example.v1.GetFooResponse" in api_schema.messages - assert "google.example.v1.bar_common.Bar" not in api_schema.messages + # Bar (in bar_common) is used by GetBarRequest (in bar.proto, omitted), + # so Bar itself should be kept. + assert "google.example.v1.bar_common.Bar" in api_schema.messages assert "google.example.v1.GetBarRequest" not in api_schema.messages assert "google.example.v1.GetBarResponse" not in api_schema.messages @@ -3650,8 +3667,10 @@ def test_selective_gapic_api_build_remove_unnecessary_proto_files(): assert "foo.proto" in api_schema.protos assert "foo_common.proto" in api_schema.protos - assert "bar.proto" not in api_schema.protos - assert "bar_common.proto" not in api_schema.protos + # bar.proto is NOT omitted because it contains the kept Bar message + # (even if the RPC that uses it is omitted). + assert "bar.proto" in api_schema.protos + assert "bar_common.proto" in api_schema.protos # Check that the sub-packages that have been completely pruned are excluded from generation, # but the ones that have only been partially pruned will still be appropriately included. @@ -3659,7 +3678,79 @@ def test_selective_gapic_api_build_remove_unnecessary_proto_files(): sub = api_schema.subpackages["foo_common"] assert len(sub.protos) == 1 assert "google.example.v1.foo_common.Foo" in sub.messages - assert "bar_common" not in api_schema.subpackages + assert "bar_common" in api_schema.subpackages + sub = api_schema.subpackages["bar_common"] + assert len(sub.protos) == 1 + assert "google.example.v1.bar_common.Bar" in sub.messages + + +def test_selective_gapic_api_build_keep_messages_within_omitted_rpcs(): + fd = ( + make_file_pb2( + name="foobar.proto", + package="google.example.v1", + messages=( + make_message_pb2( + name="IncludedRequest", + ), + make_message_pb2( + name="IncludedResponse", + ), + make_message_pb2( + name="OmittedRequest", + fields=( + make_field_pb2( + name="nested", + number=1, + type_name=".google.example.v1.NestedMessage", + ), + ), + ), + make_message_pb2( + name="OmittedResponse", + ), + make_message_pb2( + name="NestedMessage", + ), + ), + services=( + descriptor_pb2.ServiceDescriptorProto( + name="FooService", + method=( + descriptor_pb2.MethodDescriptorProto( + name="IncludedMethod", + input_type="google.example.v1.IncludedRequest", + output_type="google.example.v1.IncludedResponse", + ), + descriptor_pb2.MethodDescriptorProto( + name="OmittedMethod", + input_type="google.example.v1.OmittedRequest", + output_type="google.example.v1.OmittedResponse", + ), + ), + ), + ), + ), + ) + + service_yaml_config = get_service_yaml_for_selective_gapic_tests( + apis=["google.example.v1.FooService"], + methods=["google.example.v1.FooService.IncludedMethod"], + ) + opts = Options(service_yaml_config=service_yaml_config) + + api_schema = api.API.build(fd, "google.example.v1", opts=opts) + + # Included RPC and its messages should be present + assert "google.example.v1.IncludedRequest" in api_schema.messages + assert "google.example.v1.IncludedResponse" in api_schema.messages + + # Omitted RPC top-level messages should be ABSENT + assert "google.example.v1.OmittedRequest" not in api_schema.messages + assert "google.example.v1.OmittedResponse" not in api_schema.messages + + # Nested message ONLY reachable from omitted RPC should be PRESENT + assert "google.example.v1.NestedMessage" in api_schema.messages def test_selective_gapic_api_build_with_enums(): @@ -3742,8 +3833,11 @@ def test_selective_gapic_api_build_with_enums(): api_schema = api.API.build(fd, "google.example.v1", opts=opts) assert "google.example.v1.FooStatus" in api_schema.enums - assert "google.example.v1.BarStatus" not in api_schema.enums + # BarStatus is reachable from Bar, which is used by GetBarRequest (omitted), + # but Bar and its enums should be kept. + assert "google.example.v1.BarStatus" in api_schema.enums assert "google.example.v1.FooStatus" in api_schema.top_level_enums + assert "google.example.v1.BarStatus" in api_schema.top_level_enums def test_selective_gapic_api_build_with_nested_fields(): @@ -3839,10 +3933,14 @@ def test_selective_gapic_api_build_with_nested_fields(): assert "google.example.v1.Foo.FooStatus" in api_schema.enums assert "google.example.v1.Foo.Bar" in api_schema.messages - # Check that we can exclude nested types as well - assert "google.example.v1.Spam" not in api_schema.messages - assert "google.example.v1.Spam.SpamStatus" not in api_schema.enums - assert "google.example.v1.Spam.Ham" not in api_schema.messages + # Spam is used by GetBarRequest (omitted), so Spam and its nested types should be kept. + assert "google.example.v1.Spam" in api_schema.messages + assert "google.example.v1.Spam.SpamStatus" in api_schema.enums + assert "google.example.v1.Spam.Ham" in api_schema.messages + + # Top-level request/response messages for omitted RPCs are still omitted. + assert "google.example.v1.GetBarRequest" not in api_schema.messages + assert "google.example.v1.GetBarResponse" not in api_schema.messages @pytest.mark.parametrize("reference_attr", ["type", "child_type"]) @@ -3936,15 +4034,18 @@ def test_selective_gapic_api_build_with_resources(reference_attr): assert "google.example.v1.GetFooRequest" in api_schema.messages assert "google.example.v1.GetFooResponse" in api_schema.messages - assert "google.example.v1.Bar" not in api_schema.messages - assert "google.example.v1.BarDep" not in api_schema.messages + # Bar and BarDep are reachable from GetBarRequest (omitted), + # but they should be kept. + assert "google.example.v1.Bar" in api_schema.messages + assert "google.example.v1.BarDep" in api_schema.messages assert "google.example.v1.GetBarRequest" not in api_schema.messages assert "google.example.v1.GetBarResponse" not in api_schema.messages # Ensure we're also pruning resource messages for the files resource_messages = api_schema.protos["foo.proto"].resource_messages assert "foo.bar/Foo" in resource_messages - assert "foo.bar/Bar" not in resource_messages + # Bar is reachable via GetBarRequest, so it should be kept in resource_messages. + assert "foo.bar/Bar" in resource_messages def test_selective_gapic_api_build_extended_lro(): @@ -4529,3 +4630,118 @@ def test_file_level_resources_with_aliases(): # Verify that the resource exists with the overriden type assert "nomenclature.linnaen.com/Species" in resource_msgs assert resource_msgs["nomenclature.linnaen.com/Species"].resource_type == "CustomSpecies" + + +from test_utils.test_utils import ( + make_file_pb2, + make_message_pb2, +) + +def test_proto_with_selective_generation_returns_none(): + # Proto A: will be omitted because its only method is not in the allow-list. + fd_a = make_file_pb2( + name="proto_a.proto", + package="google.example.v1", + services=( + descriptor_pb2.ServiceDescriptorProto( + name="ServiceA", + method=( + descriptor_pb2.MethodDescriptorProto( + name="MethodA", + input_type="google.example.v1.ReqA", + output_type="google.example.v1.RespA", + ), + ), + ), + ), + messages=( + make_message_pb2("ReqA"), + make_message_pb2("RespA"), + ), + ) + + # Proto B: will be kept because its method IS in the allow-list. + fd_b = make_file_pb2( + name="proto_b.proto", + package="google.example.v1", + services=( + descriptor_pb2.ServiceDescriptorProto( + name="ServiceB", + method=( + descriptor_pb2.MethodDescriptorProto( + name="MethodB", + input_type="google.example.v1.ReqB", + output_type="google.example.v1.RespB", + ), + ), + ), + ), + messages=( + make_message_pb2("ReqB"), + make_message_pb2("RespB"), + ), + ) + + # Selective generation only including ServiceB.MethodB. + service_yaml_config = get_service_yaml_for_selective_gapic_tests( + methods=["google.example.v1.ServiceB.MethodB"] + ) + opts = Options(service_yaml_config=service_yaml_config) + + # Building the API will call with_selective_generation for both protos. + api_schema = api.API.build([fd_a, fd_b], package="google.example.v1", opts=opts) + + # proto_a.proto should be excluded. + assert "proto_a.proto" not in api_schema.protos + assert "proto_b.proto" in api_schema.protos + +def test_api_build_selective_multiple_protos_kept(): + fd1 = make_file_pb2( + name="proto1.proto", + package="google.example.v1", + services=( + descriptor_pb2.ServiceDescriptorProto( + name="Service1", + method=( + descriptor_pb2.MethodDescriptorProto( + name="Method1", + input_type="google.example.v1.Req1", + output_type="google.example.v1.Resp1", + ), + ), + ), + ), + messages=(make_message_pb2("Req1"), make_message_pb2("Resp1")), + ) + fd2 = make_file_pb2( + name="proto2.proto", + package="google.example.v1", + services=( + descriptor_pb2.ServiceDescriptorProto( + name="Service2", + method=( + descriptor_pb2.MethodDescriptorProto( + name="Method2", + input_type="google.example.v1.Req2", + output_type="google.example.v1.Resp2", + ), + ), + ), + ), + messages=(make_message_pb2("Req2"), make_message_pb2("Resp2")), + ) + + service_yaml_config = get_service_yaml_for_selective_gapic_tests( + methods=[ + "google.example.v1.Service1.Method1", + "google.example.v1.Service2.Method2", + ] + ) + opts = Options(service_yaml_config=service_yaml_config) + api_schema = api.API.build([fd1, fd2], package="google.example.v1", opts=opts) + + # Both protos should be kept. This covers the case where multiple protos + # are processed in the selective generation loop, taking the 534->528 branch. + assert len(api_schema.protos) == 2 + assert "proto1.proto" in api_schema.protos + assert "proto2.proto" in api_schema.protos