Skip to content

Commit

Permalink
Decouple DynamicProto by extracting logic for descriptor pool and mes…
Browse files Browse the repository at this point in the history 
…sage factory.

PiperOrigin-RevId: 570526845
  • Loading branch information
@l46kok @copybara-github
l46kok authored and copybara-github committed Oct 12, 2023
1 parent 6cca9e4 commit 38d4045
Show file tree
Hide file tree
Showing 22 changed files with 245 additions and 391 deletions.
6 changes: 2 additions & 4 deletions common/src/main/java/dev/cel/common/internal/BUILD.bazel
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,7 +91,6 @@ java_library(
],
)

# keep sorted
DYNAMIC_PROTO_SOURCES = [
"DynamicProto.java",
"ProtoAdapter.java",
Expand Down Expand Up @@ -119,13 +118,12 @@ java_library(
],
deps = [
":converter",
":default_instance_message_factory",
":proto_message_factory",
":well_known_proto",
"//:auto_value",
"//common",
"//common:error_codes",
"//common:runtime_exception",
"//common/annotations",
"//common/types:cel_types",
"@cel_spec//proto/cel/expr:expr_java_proto",
"@maven//:com_google_code_findbugs_annotations",
"@maven//:com_google_errorprone_error_prone_annotations",
Expand Down
3 changes: 2 additions & 1 deletion common/src/main/java/dev/cel/common/internal/CombinedDescriptorPool.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,8 @@ public ExtensionRegistry getExtensionRegistry() {
private CombinedDescriptorPool(ImmutableList<CelDescriptorPool> descriptorPools) {
this.descriptorPools = descriptorPools;
// TODO: Combine the extension registry. This will become necessary once we accept
// ExtensionRegistry through runtime builder.
// ExtensionRegistry through runtime builder. Ideally, proto team should open source this
// implementation but we may have to create our own.
this.extensionRegistry =
descriptorPools.stream()
.map(CelDescriptorPool::getExtensionRegistry)
Expand Down
128 changes: 9 additions & 119 deletions common/src/main/java/dev/cel/common/internal/DynamicProto.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,29 +15,16 @@
package dev.cel.common.internal;

import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.ImmutableMap.toImmutableMap;
import static java.util.Arrays.stream;

import com.google.auto.value.AutoBuilder;
import com.google.common.collect.ImmutableCollection;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMultimap;
import com.google.errorprone.annotations.CheckReturnValue;
import com.google.errorprone.annotations.Immutable;
import com.google.protobuf.Any;
import com.google.protobuf.ByteString;
import com.google.protobuf.Descriptors.Descriptor;
import com.google.protobuf.Descriptors.FieldDescriptor;
import com.google.protobuf.DynamicMessage;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.Message;
import dev.cel.common.CelDescriptorUtil;
import dev.cel.common.CelDescriptors;
import dev.cel.common.annotations.Internal;
import dev.cel.common.types.CelTypes;
import java.util.Map.Entry;
import java.util.Optional;
import org.jspecify.nullness.Nullable;

/**
* The {@code DynamicProto} class supports the conversion of {@link Any} values to concrete {@code
Expand All @@ -49,58 +36,13 @@
@CheckReturnValue
@Internal
public final class DynamicProto {

private static final ImmutableMap<String, Descriptor> WELL_KNOWN_DESCRIPTORS =
stream(ProtoAdapter.WellKnownProto.values())
.collect(toImmutableMap(d -> d.typeName(), d -> d.descriptor()));

private final ImmutableMap<String, Descriptor> dynamicDescriptors;
private final ImmutableMultimap<String, FieldDescriptor> dynamicExtensionDescriptors;
private final ProtoMessageFactory protoMessageFactory;

/** {@code ProtoMessageFactory} provides a method to create a protobuf builder objects by name. */
@Immutable
@FunctionalInterface
public interface ProtoMessageFactory {
Message.@Nullable Builder newBuilder(String messageName);
}

/** Builder for configuring the {@link DynamicProto}. */
@AutoBuilder(ofClass = DynamicProto.class)
public abstract static class Builder {

/** Sets {@link CelDescriptors} to unpack any message types. */
public abstract Builder setDynamicDescriptors(CelDescriptors celDescriptors);

/** Sets a custom type factory to unpack any message types. */
public abstract Builder setProtoMessageFactory(ProtoMessageFactory factory);

/** Builds a new instance of {@link DynamicProto} */
@CheckReturnValue
public abstract DynamicProto build();
public static DynamicProto create(ProtoMessageFactory protoMessageFactory) {
return new DynamicProto(protoMessageFactory);
}

public static Builder newBuilder() {
return new AutoBuilder_DynamicProto_Builder()
.setDynamicDescriptors(CelDescriptors.builder().build())
.setProtoMessageFactory((typeName) -> null);
}

DynamicProto(
CelDescriptors dynamicDescriptors,
ProtoMessageFactory protoMessageFactory) {
ImmutableMap<String, Descriptor> messageTypeDescriptorMap =
CelDescriptorUtil.descriptorCollectionToMap(dynamicDescriptors.messageTypeDescriptors());
ImmutableMap<String, Descriptor> filteredDescriptors =
messageTypeDescriptorMap.entrySet().stream()
.filter(e -> !WELL_KNOWN_DESCRIPTORS.containsKey(e.getKey()))
.collect(toImmutableMap(Entry::getKey, Entry::getValue));
this.dynamicDescriptors =
ImmutableMap.<String, Descriptor>builder()
.putAll(WELL_KNOWN_DESCRIPTORS)
.putAll(filteredDescriptors)
.buildOrThrow();
this.dynamicExtensionDescriptors = checkNotNull(dynamicDescriptors.extensionDescriptors());
DynamicProto(ProtoMessageFactory protoMessageFactory) {
this.protoMessageFactory = checkNotNull(protoMessageFactory);
}

Expand All @@ -120,7 +62,8 @@ public Message unpack(Any any) throws InvalidProtocolBufferException {
String.format("malformed type URL: %s", any.getTypeUrl())));

Message.Builder builder =
newMessageBuilder(messageTypeName)
protoMessageFactory
.newBuilder(messageTypeName)
.orElseThrow(
() ->
new InvalidProtocolBufferException(
Expand All @@ -136,7 +79,7 @@ public Message unpack(Any any) throws InvalidProtocolBufferException {
*/
public Message maybeAdaptDynamicMessage(DynamicMessage input) {
Optional<Message.Builder> maybeBuilder =
newMessageBuilder(input.getDescriptorForType().getFullName());
protoMessageFactory.newBuilder(input.getDescriptorForType().getFullName());
if (!maybeBuilder.isPresent() || maybeBuilder.get() instanceof DynamicMessage.Builder) {
// Just return the same input if:
// 1. We didn't get a builder back because there's no descriptor (nothing we can do)
Expand All @@ -148,61 +91,6 @@ public Message maybeAdaptDynamicMessage(DynamicMessage input) {
return merge(maybeBuilder.get(), input.toByteString());
}

/**
* This method instantiates a builder for the given {@code typeName} assuming one is configured
* within the descriptor set provided to the {@code DynamicProto} constructor.
*
* <p>When the {@code useLinkedTypes} flag is set, the {@code Message.Builder} returned will be
* the concrete builder instance linked into the binary if it is present; otherwise, the result
* will be a {@code DynamicMessageBuilder}.
*/
public Optional<Message.Builder> newMessageBuilder(String typeName) {
if (!CelTypes.isWellKnownType(typeName)) {
// Check if the message factory can produce a concrete message via custom type factory
// first.
Message.Builder builder = protoMessageFactory.newBuilder(typeName);
if (builder != null) {
return Optional.of(builder);
}
}

Optional<Descriptor> descriptor = maybeGetDescriptor(typeName);
if (!descriptor.isPresent()) {
return Optional.empty();
}
// If the descriptor that's resolved does not match the descriptor instance in the message
// factory, the call to fetch the prototype will return null, and a dynamic proto message
// should be used as a fallback.
Optional<Message> message =
DefaultInstanceMessageFactory.getInstance().getPrototype(descriptor.get());
if (message.isPresent()) {
return Optional.of(message.get().toBuilder());
}

// Fallback to a dynamic proto instance.
return Optional.of(DynamicMessage.newBuilder(descriptor.get()));
}

private Optional<Descriptor> maybeGetDescriptor(String typeName) {

Descriptor descriptor = ProtoRegistryProvider.getTypeRegistry().find(typeName);
return Optional.ofNullable(descriptor != null ? descriptor : dynamicDescriptors.get(typeName));
}

/** Gets the corresponding field descriptor for an extension field on a message. */
public Optional<FieldDescriptor> maybeGetExtensionDescriptor(
Descriptor containingDescriptor, String fieldName) {

String typeName = containingDescriptor.getFullName();
ImmutableCollection<FieldDescriptor> fieldDescriptors =
dynamicExtensionDescriptors.get(typeName);
if (fieldDescriptors.isEmpty()) {
return Optional.empty();
}

return fieldDescriptors.stream().filter(d -> d.getFullName().equals(fieldName)).findFirst();
}

/**
* Merge takes in a Message builder and merges another message bytes into the builder. Some
* example usages are:
Expand All @@ -214,7 +102,9 @@ public Optional<FieldDescriptor> maybeGetExtensionDescriptor(
*/
private Message merge(Message.Builder builder, ByteString inputBytes) {
try {
return builder.mergeFrom(inputBytes, ProtoRegistryProvider.getExtensionRegistry()).build();
return builder
.mergeFrom(inputBytes, protoMessageFactory.getDescriptorPool().getExtensionRegistry())
.build();
} catch (InvalidProtocolBufferException e) {
throw new AssertionError("Failed to merge input message into the message builder", e);
}
Expand Down
47 changes: 0 additions & 47 deletions common/src/main/java/dev/cel/common/internal/ProtoAdapter.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,53 +143,6 @@ public final class ProtoAdapter {
stream(WellKnownProto.values())
.collect(toImmutableMap(WellKnownProto::typeName, Function.identity()));

/**
* WellKnownProto types used throughout CEL. These types are specially handled to ensure that
* bidirectional conversion between CEL native values and these well-known types is performed
* consistently across runtimes.
*/
enum WellKnownProto {
JSON_VALUE(Value.getDescriptor()),
JSON_STRUCT_VALUE(Struct.getDescriptor()),
JSON_LIST_VALUE(ListValue.getDescriptor()),
ANY_VALUE(Any.getDescriptor()),
BOOL_VALUE(BoolValue.getDescriptor(), true),
BYTES_VALUE(BytesValue.getDescriptor(), true),
DOUBLE_VALUE(DoubleValue.getDescriptor(), true),
FLOAT_VALUE(FloatValue.getDescriptor(), true),
INT32_VALUE(Int32Value.getDescriptor(), true),
INT64_VALUE(Int64Value.getDescriptor(), true),
STRING_VALUE(StringValue.getDescriptor(), true),
UINT32_VALUE(UInt32Value.getDescriptor(), true),
UINT64_VALUE(UInt64Value.getDescriptor(), true),
DURATION_VALUE(Duration.getDescriptor()),
TIMESTAMP_VALUE(Timestamp.getDescriptor());

private final Descriptor descriptor;
private final boolean isWrapperType;

WellKnownProto(Descriptor descriptor) {
this(descriptor, /* isWrapperType= */ false);
}

WellKnownProto(Descriptor descriptor, boolean isWrapperType) {
this.descriptor = descriptor;
this.isWrapperType = isWrapperType;
}

Descriptor descriptor() {
return descriptor;
}

String typeName() {
return descriptor.getFullName();
}

boolean isWrapperType() {
return isWrapperType;
}
}

private final DynamicProto dynamicProto;
private final boolean enableUnsignedLongs;

Expand Down
Loading

0 comments on commit 38d4045

Please sign in to comment.