Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 3ef95d6

Browse filesBrowse files
partheaparthea
authored
fix: restore messages not attached to rpc in selective_gapic_generation (#16951)
Towards b/507893758, b/507889482, b/501132869, b/503326310
1 parent 0f7c68b commit 3ef95d6
Copy full SHA for 3ef95d6

7 files changed

+538-144Lines changed: 538 additions & 144 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

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

Copy file name to clipboardExpand all lines: packages/gapic-generator/gapic/schema/api.py
+106-76Lines changed: 106 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,7 @@ def add_to_address_allowlist(
265265
address_allowlist: Set["metadata.Address"],
266266
method_allowlist: Set[str],
267267
resource_messages: Dict[str, "wrappers.MessageType"],
268+
services_in_proto: Dict[str, "wrappers.Service"],
268269
) -> None:
269270
"""Adds to the set of Addresses of wrapper objects to be included in selective GAPIC generation.
270271
@@ -281,15 +282,12 @@ def add_to_address_allowlist(
281282
resource type name of a resource message to the corresponding MessageType object
282283
representing that resource message. Only resources with a message representation
283284
should be included in the dictionary.
285+
services_in_proto (Dict[str, wrappers.Service]): A dictionary mapping the names of Service
286+
objects in the proto containing this method to the Service objects. This is necessary
287+
for traversing the operation service in the case of extended LROs.
284288
Returns:
285289
None
286290
"""
287-
# The method.operation_service for an extended LRO is not fully qualified, so we
288-
# truncate the service names accordingly so they can be found in
289-
# method.add_to_address_allowlist
290-
services_in_proto = {
291-
service.name: service for service in self.services.values()
292-
}
293291
for service in self.services.values():
294292
service.add_to_address_allowlist(
295293
address_allowlist=address_allowlist,
@@ -298,75 +296,60 @@ def add_to_address_allowlist(
298296
services_in_proto=services_in_proto,
299297
)
300298

301-
def prune_messages_for_selective_generation(
302-
self, *, address_allowlist: Set["metadata.Address"]
299+
def with_selective_generation(
300+
self,
301+
*,
302+
generate_omitted_as_internal: bool,
303+
public_methods: Set[str],
304+
excluded_addresses: Set["metadata.Address"],
303305
) -> Optional["Proto"]:
304-
"""Returns a truncated version of this Proto.
305-
306-
Only the services, messages, and enums contained in the allowlist
307-
of visited addresses are included in the returned object. If there
308-
are no services, messages, or enums left, and no file level resources,
309-
return None.
306+
"""Returns a version of this Proto for selective generation.
310307
311308
Args:
312-
address_allowlist (Set[metadata.Address]): A set of allowlisted metadata.Address
313-
objects to filter against. Objects with addresses not the allowlist will be
314-
removed from the returned Proto.
309+
generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
310+
public_methods (Set[str]): The set of fully-qualified method names to keep as public.
311+
excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
312+
315313
Returns:
316-
Optional[Proto]: A truncated version of this proto. If there are no services, messages,
317-
or enums left after the truncation process and there are no file level resources,
318-
returns None.
314+
Optional[Proto]: A version of this Proto with services/methods filtered.
315+
Returns None if the Proto becomes empty and generate_omitted_as_internal is False.
319316
"""
320-
# Once the address allowlist has been created, it suffices to only
321-
# prune items at 2 different levels to truncate the Proto object:
317+
services = {}
318+
for k, v in self.services.items():
319+
new_v = v.with_selective_generation(
320+
generate_omitted_as_internal=generate_omitted_as_internal,
321+
public_methods=public_methods,
322+
excluded_addresses=excluded_addresses)
323+
if new_v:
324+
services[k] = new_v
325+
326+
# We only prune messages/enums from protos that are not dependencies.
327+
# A message or enum is excluded IF AND ONLY IF:
328+
# 1. It is a top-level request or response message for an omitted RPC.
329+
# 2. It is NOT reachable from any publicly allowed RPC.
322330
#
323-
# 1. At the Proto level, we remove unnecessary services, messages,
324-
# and enums.
325-
# 2. For allowlisted services, at the Service level, we remove
326-
# non-allowlisted methods.
327-
services = {
328-
k: v.prune_messages_for_selective_generation(
329-
address_allowlist=address_allowlist
330-
)
331-
for k, v in self.services.items()
332-
if v.meta.address in address_allowlist
333-
}
334-
331+
# This ensures that shared messages, messages not attached to any RPC,
332+
# and messages reachable via other paths (like LRO response types) are KEPT.
335333
all_messages = {
336-
k: v for k, v in self.all_messages.items() if v.ident in address_allowlist
334+
k: v for k, v in self.all_messages.items() if v.ident not in excluded_addresses
337335
}
338336

339337
all_enums = {
340-
k: v for k, v in self.all_enums.items() if v.ident in address_allowlist
338+
k: v for k, v in self.all_enums.items() if v.ident not in excluded_addresses
341339
}
342340

343-
if not services and not all_messages and not all_enums:
341+
# If the proto becomes empty after pruning, we return None to signal
342+
# that it should be excluded from generation.
343+
if not generate_omitted_as_internal and not services and not all_messages and not all_enums:
344344
return None
345345

346346
return dataclasses.replace(
347-
self, services=services, all_messages=all_messages, all_enums=all_enums
347+
self,
348+
services=services,
349+
all_messages=all_messages,
350+
all_enums=all_enums,
348351
)
349352

350-
def with_internal_methods(self, *, public_methods: Set[str]) -> "Proto":
351-
"""Returns a version of this Proto with some Methods marked as internal.
352-
353-
The methods not in the public_methods set will be marked as internal and
354-
services containing these methods will also be marked as internal by extension.
355-
(See :meth:`Service.is_internal` for more details).
356-
357-
Args:
358-
public_methods (Set[str]): An allowlist of fully-qualified method names.
359-
Methods not in this allowlist will be marked as internal.
360-
Returns:
361-
Proto: A version of this Proto with Method objects corresponding to methods
362-
not in `public_methods` marked as internal.
363-
"""
364-
services = {
365-
k: v.with_internal_methods(public_methods=public_methods)
366-
for k, v in self.services.items()
367-
}
368-
return dataclasses.replace(self, services=services)
369-
370353

371354
@dataclasses.dataclass(frozen=True)
372355
class API:
@@ -532,32 +515,79 @@ def disambiguate_keyword_sanitize_fname(
532515

533516
if selective_gapic_settings.generate_omitted_as_internal:
534517
for name, proto in api.protos.items():
535-
new_all_protos[name] = proto.with_internal_methods(
536-
public_methods=selective_gapic_methods
518+
new_all_protos[name] = proto.with_selective_generation(
519+
generate_omitted_as_internal=True,
520+
public_methods=selective_gapic_methods,
521+
excluded_addresses=set([]),
537522
)
538523
else:
539-
all_resource_messages = collections.ChainMap(
540-
*(proto.resource_messages for proto in protos.values())
541-
)
542-
543-
# Prepare a list of addresses to include in selective generation,
544-
# then prune each Proto object. We look at metadata.Addresses, not objects, because
545-
# objects that refer to the same thing in the proto are different Python objects
546-
# in memory.
547-
address_allowlist: Set["metadata.Address"] = set([])
548-
for proto in api.protos.values():
524+
all_resource_messages = dict(collections.ChainMap(
525+
*(proto.resource_messages for proto in api.all_protos.values())
526+
))
527+
528+
# Create a global map of services to support cross-proto lookup
529+
# for extended LROs.
530+
#
531+
# Note: This is keyed by the fully qualified proto name (as a string)
532+
# to ensure compatibility with Address.resolve() lookups in wrappers.py.
533+
all_services: Dict[str, wrappers.Service] = {}
534+
for p in api.all_protos.values():
535+
for s in p.services.values():
536+
all_services[s.meta.address.proto] = s
537+
538+
# Calculate addresses of omitted RPCs and their top-level request/response messages.
539+
# These are "candidates" for exclusion.
540+
#
541+
# We only consider top-level request/response messages of omitted RPCs as
542+
# candidates for exclusion. This is conservative: it ensures that:
543+
# - Messages NOT used by any RPC are KEPT (e.g. for user convenience).
544+
# - Messages shared between an omitted and a kept RPC are KEPT.
545+
# - Messages reachable from a kept RPC but NOT as a top-level request/response
546+
# (e.g. nested messages) are KEPT.
547+
candidate_excluded_addresses: Set["metadata.Address"] = set([])
548+
for proto in api.all_protos.values():
549+
for service in proto.services.values():
550+
for method in service.methods.values():
551+
if method.ident.proto not in selective_gapic_methods:
552+
# Candidate for exclusion: the method itself and its direct request/response types.
553+
candidate_excluded_addresses.add(method.meta.address)
554+
candidate_excluded_addresses.add(method.input.ident)
555+
candidate_excluded_addresses.add(method.output.ident)
556+
557+
# If this is an LRO, add its response and metadata types to candidates.
558+
if method.lro:
559+
candidate_excluded_addresses.add(method.lro.response_type.ident)
560+
candidate_excluded_addresses.add(method.lro.metadata_type.ident)
561+
562+
# If this is an extended LRO, add its request and operation types to candidates.
563+
if method.extended_lro:
564+
candidate_excluded_addresses.add(method.extended_lro.request_type.ident)
565+
candidate_excluded_addresses.add(method.extended_lro.operation_type.ident)
566+
567+
# Calculate publicly reachable addresses (API-wide).
568+
# This includes all types reachable from the allowlisted (public) methods.
569+
public_rpc_addresses: Set["metadata.Address"] = set([])
570+
for proto in api.all_protos.values():
549571
proto.add_to_address_allowlist(
550-
address_allowlist=address_allowlist,
572+
address_allowlist=public_rpc_addresses,
551573
method_allowlist=selective_gapic_methods,
552574
resource_messages=all_resource_messages,
575+
services_in_proto=all_services,
553576
)
554577

555-
# We only prune services/messages/enums from protos that are not dependencies.
578+
# Addresses to exclude: those that are candidates for exclusion but NOT
579+
# reachable from any PUBLIC RPC.
580+
#
581+
# This set difference effectively "vets" the candidates. If a candidate
582+
# message is actually reachable from a public RPC, it's removed from
583+
# the exclusion list.
584+
excluded_addresses = candidate_excluded_addresses - public_rpc_addresses
585+
556586
for name, proto in api.protos.items():
557-
proto_to_generate = (
558-
proto.prune_messages_for_selective_generation(
559-
address_allowlist=address_allowlist
560-
)
587+
proto_to_generate = proto.with_selective_generation(
588+
generate_omitted_as_internal=False,
589+
public_methods=selective_gapic_methods,
590+
excluded_addresses=excluded_addresses,
561591
)
562592
if proto_to_generate:
563593
new_all_protos[name] = proto_to_generate
Collapse file

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

Copy file name to clipboardExpand all lines: packages/gapic-generator/gapic/schema/wrappers.py
+54-51Lines changed: 54 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -2032,7 +2032,9 @@ def add_to_address_allowlist(
20322032
# the allowlist, as it might not have been specified by
20332033
# the methods under selective_gapic_generation.
20342034
# We assume that the operation service lives in the same proto file as this one.
2035-
operation_service = services_in_proto[self.operation_service]
2035+
operation_service = services_in_proto[
2036+
self.meta.address.resolve(self.operation_service)
2037+
]
20362038
address_allowlist.add(operation_service.meta.address)
20372039
operation_service.operation_polling_method.add_to_address_allowlist(
20382040
address_allowlist=address_allowlist,
@@ -2055,26 +2057,39 @@ def add_to_address_allowlist(
20552057
resource_messages=resource_messages,
20562058
)
20572059

2058-
def with_internal_methods(self, *, public_methods: Set[str]) -> "Method":
2059-
"""Returns a version of this ``Method`` marked as internal
2060-
2061-
The methods not in the public_methods set will be marked as internal and
2062-
this ``Service`` will as well by extension (see :meth:`Service.is_internal`).
2060+
def with_selective_generation(
2061+
self,
2062+
*,
2063+
generate_omitted_as_internal: bool,
2064+
public_methods: Set[str],
2065+
excluded_addresses: Set["metadata.Address"],
2066+
) -> Optional["Method"]:
2067+
"""Returns a version of this Method for selective generation.
20632068
20642069
Args:
2065-
public_methods (Set[str]): An allowlist of fully-qualified method names.
2066-
Methods not in this allowlist will be marked as internal.
2070+
generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
2071+
public_methods (Set[str]): The set of fully-qualified method names to keep as public.
2072+
excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
2073+
20672074
Returns:
2068-
Service: A version of this `Service` with `Method` objects corresponding to methods
2069-
not in `public_methods` marked as internal.
2075+
Optional[Method]: The method, possibly marked as internal, or None if it should be removed.
20702076
"""
20712077
if self.ident.proto in public_methods:
20722078
return self
20732079

2074-
return dataclasses.replace(
2075-
self,
2076-
is_internal=True,
2077-
)
2080+
# Not public.
2081+
# We mark it as internal if either:
2082+
# 1. generate_omitted_as_internal is set in selective_gapic_generation.
2083+
# 2. The method is NOT in excluded_addresses, which means it is reachable
2084+
# from some public method (e.g. as a polling method for an extended LRO).
2085+
if generate_omitted_as_internal or self.meta.address not in excluded_addresses:
2086+
return dataclasses.replace(
2087+
self,
2088+
is_internal=True,
2089+
)
2090+
else:
2091+
return None
2092+
20782093

20792094

20802095
@dataclasses.dataclass(frozen=True)
@@ -2463,45 +2478,33 @@ def add_to_address_allowlist(
24632478
services_in_proto=services_in_proto,
24642479
)
24652480

2466-
def prune_messages_for_selective_generation(
2467-
self, *, address_allowlist: Set["metadata.Address"]
2468-
) -> "Service":
2469-
"""Returns a truncated version of this Service.
2470-
2471-
Only the methods, messages, and enums contained in the address allowlist
2472-
are included in the returned object.
2481+
def with_selective_generation(
2482+
self,
2483+
*,
2484+
generate_omitted_as_internal: bool,
2485+
public_methods: Set[str],
2486+
excluded_addresses: Set["metadata.Address"],
2487+
) -> Optional["Service"]:
2488+
"""Returns a version of this Service for selective generation.
24732489
24742490
Args:
2475-
address_allowlist (Set[metadata.Address]): A set of allowlisted metadata.Address
2476-
objects to filter against. Objects with addresses not the allowlist will be
2477-
removed from the returned Proto.
2478-
Returns:
2479-
Service: A truncated version of this proto.
2480-
"""
2481-
return dataclasses.replace(
2482-
self,
2483-
methods={
2484-
k: v for k, v in self.methods.items() if v.ident in address_allowlist
2485-
},
2486-
)
2487-
2488-
def with_internal_methods(self, *, public_methods: Set[str]) -> "Service":
2489-
"""Returns a version of this ``Service`` with some Methods marked as internal.
2490-
2491-
The methods not in the public_methods set will be marked as internal and
2492-
this ``Service`` will as well by extension (see :meth:`Service.is_internal`).
2491+
generate_omitted_as_internal (bool): Whether to mark omitted methods as internal.
2492+
public_methods (Set[str]): The set of fully-qualified method names to keep as public.
2493+
excluded_addresses (Set[metadata.Address]): The set of addresses to exclude from generation.
24932494
2494-
Args:
2495-
public_methods (Set[str]): An allowlist of fully-qualified method names.
2496-
Methods not in this allowlist will be marked as internal.
24972495
Returns:
2498-
Service: A version of this `Service` with `Method` objects corresponding to methods
2499-
not in `public_methods` marked as internal.
2496+
Optional[Service]: The service with filtered methods, or None if it should be removed.
25002497
"""
2501-
return dataclasses.replace(
2502-
self,
2503-
methods={
2504-
k: v.with_internal_methods(public_methods=public_methods)
2505-
for k, v in self.methods.items()
2506-
},
2507-
)
2498+
methods = {}
2499+
for k, v in self.methods.items():
2500+
new_v = v.with_selective_generation(
2501+
generate_omitted_as_internal=generate_omitted_as_internal,
2502+
public_methods=public_methods,
2503+
excluded_addresses=excluded_addresses)
2504+
if new_v:
2505+
methods[k] = new_v
2506+
2507+
if not generate_omitted_as_internal and not methods:
2508+
return None
2509+
2510+
return dataclasses.replace(self, methods=methods)

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.