Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
parthea wants to merge 2 commits into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
92621ee
fix: restore messages not attached to rpc in selective_gapic_generation
parthea May 14, 2026
defd209
mypy
parthea May 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
182 changes: 106 additions & 76 deletions packages/gapic-generator/gapic/schema/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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,
Expand All @@ -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"]:
Comment thread
parthea marked this conversation as resolved.
"""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:
Expand Down Expand Up @@ -532,32 +515,79 @@ def disambiguate_keyword_sanitize_fname(

Comment thread
parthea marked this conversation as resolved.
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
Expand Down
105 changes: 54 additions & 51 deletions packages/gapic-generator/gapic/schema/wrappers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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"]:
Comment thread
parthea marked this conversation as resolved.
"""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
Comment thread
parthea marked this conversation as resolved.



@dataclasses.dataclass(frozen=True)
Expand Down Expand Up @@ -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"]:
Comment thread
parthea marked this conversation as resolved.
"""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)
Comment thread
parthea marked this conversation as resolved.
Loading
Loading