Re-attach proto2 extensions when loading meta-data from JSON

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/command/CommandUtil.java

@@ -21,7 +21,8 @@
 package com.apple.foundationdb.relational.yamltests.command;
 
 import com.apple.foundationdb.record.IndexState;
 import com.apple.foundationdb.record.RecordMetaData;
+import com.apple.foundationdb.record.RecordMetaDataOptionsProto;
 import com.apple.foundationdb.record.RecordMetaDataProto;
 import com.apple.foundationdb.record.util.pair.Pair;
 import com.apple.foundationdb.relational.api.metadata.SchemaTemplate;
@@ -31,11 +32,16 @@
 import com.google.gson.JsonElement;
 import com.google.gson.JsonObject;
 import com.google.gson.JsonParser;
+import com.google.protobuf.ByteString;
 import com.google.protobuf.Descriptors;
+import com.google.protobuf.ExtensionRegistry;
+import com.google.protobuf.InvalidProtocolBufferException;
+import com.google.protobuf.Message;
 import com.google.protobuf.util.JsonFormat;
 import org.junit.jupiter.api.Assertions;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.io.IOException;
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
@@ -44,11 +50,14 @@
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.StringTokenizer;
 import java.util.regex.Matcher;
@@ -102,68 +111,234 @@
     }
 
     private static RecordMetaData loadRecordMetaDataFromJson(String jsonFileName) {
         RecordMetaDataProto.MetaData.Builder builder = RecordMetaDataProto.MetaData.newBuilder();
         Set<String> neededDependencies = new LinkedHashSet<>();
         Set<String> includedDependencies = new HashSet<>();
 
         // These dependencies are automatically added, so we can treat them like they are bundled with the file dependencies
         includedDependencies.add("record_metadata.proto");
         includedDependencies.add("record_metadata_options.proto");
         includedDependencies.add("tuple_fields.proto");
 
+        final JsonObject obj;
         try {
             String jsonStr = Files.readString(Paths.get(jsonFileName), StandardCharsets.UTF_8);
 
             // Load the definition into the meta-data proto
             JsonFormat.parser().ignoringUnknownFields().merge(jsonStr, builder);
 
             // Find the list of dependencies of the top-level file
-            JsonObject obj = JsonParser.parseString(jsonStr).getAsJsonObject();
+            obj = JsonParser.parseString(jsonStr).getAsJsonObject();
             JsonArray dependencyArray = obj.getAsJsonObject("records").getAsJsonArray("dependency");
             for (JsonElement element : dependencyArray) {
                 String curDep = element.getAsString();
                 neededDependencies.add(curDep);
             }
 
             // Some dependencies may be included in the JSON descriptor itself and do not need to be
             // provided from the environment
             JsonArray includedDependencyDefinitions = obj.getAsJsonArray("dependencies");
             if (includedDependencyDefinitions != null) {
                 for (JsonElement element : includedDependencyDefinitions) {
                     JsonObject definition = element.getAsJsonObject();
                     includedDependencies.add(definition.get("name").getAsString());
                     if (definition.has("dependency")) {
                         definition.getAsJsonArray("dependency")
                                 .forEach(dep -> neededDependencies.add(dep.getAsString()));
                     }
                 }
             }
         } catch (IOException e) {
             throw new RuntimeException(e);
         }
 
         List<Descriptors.FileDescriptor> fileDescriptors = new ArrayList<>();
+        List<Class<?>> dependencyClasses = new ArrayList<>();
         for (String dep: neededDependencies) {
             if (includedDependencies.contains(dep)) {
                 continue;
             }
             try {
                 String fullClassName = getFullClassName(dep);
                 Class<?> act = Class.forName(fullClassName);
                 Method method = act.getMethod("getDescriptor");
                 fileDescriptors.add((Descriptors.FileDescriptor) method.invoke(null));
+                dependencyClasses.add(act);
             } catch (NoSuchMethodException | InvocationTargetException | IllegalAccessException |
                      IOException | ClassNotFoundException e) {
                 throw new RuntimeException(e);
             }
         }
 
+        mergeExtensions(builder, obj, extensionRegistry(dependencyClasses));
+
         return RecordMetaData.newBuilder()
                 .addDependencies(fileDescriptors.toArray(new Descriptors.FileDescriptor[0]))
                 .setRecords(builder.build())
                 .getRecordMetaData();
     }
 
+    /**
+     * Collect the extensions declared by the meta-data protos themselves and by every resolved dependency, so that the
+     * extensions the JSON carries can be looked up by their full name, see
+     * {@link #mergeExtensions(Message.Builder, JsonObject, ExtensionRegistry)}.
+     *
+     * @param dependencyClasses the generated outer classes of the resolved dependencies
+     * @return a registry holding all extensions those classes declare
+     */
+    @Nonnull
+    private static ExtensionRegistry extensionRegistry(@Nonnull List<Class<?>> dependencyClasses) {
+        final ExtensionRegistry registry = ExtensionRegistry.newInstance();
+        RecordMetaDataOptionsProto.registerAllExtensions(registry);
+        for (Class<?> dependencyClass : dependencyClasses) {
+            for (Method method : dependencyClass.getMethods()) {
+                // a generated class without any extension to register does not have that method at all
+                if (!"registerAllExtensions".equals(method.getName())
+                        || !Arrays.equals(method.getParameterTypes(), new Class<?>[] {ExtensionRegistry.class})) {
+                    continue;
+                }
+                try {
+                    method.invoke(null, registry);
+                } catch (IllegalAccessException | InvocationTargetException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+        }
+        return registry;
+    }
+
+    /**
+     * Re-attach the proto2 extensions that {@link JsonFormat} discarded while parsing the meta-data.
+     * <p>
+     * {@code JsonFormat} implements the proto3 JSON mapping, which has no notion of extensions: they are treated as if
+     * they did not exist, both when the meta-data is written and when it is read back. Extensions do, however, carry
+     * information the meta-data cannot do without. A vector field, for instance, is a {@code bytes} field whose
+     * dimensions and precision live in an extension of {@code google.protobuf.FieldOptions}, so dropping the extension
+     * silently turns the field into a plain {@code bytes} field:
+     * <pre>{@code
+     * "embedding": { "type": "TYPE_BYTES",
+     *                "options": { "com.apple.foundationdb.record.field": {
+     *                             "vectorOptions": { "precision": 64, "dimensions": 512 }}}}
+     * }</pre>
+     * The value of an extension is itself an ordinary message, though, so this walks the JSON alongside the builder it
+     * was parsed into, and for every key naming a known extension of the message at hand, parses that value and sets it
+     * on the builder.
+     * </p>
+     *
+     * @param builder the builder the {@code json} object was parsed into
+     * @param json the JSON object the builder was parsed from
+     * @param registry the extensions to look the keys of {@code json} up in
+     */
+    private static void mergeExtensions(@Nonnull Message.Builder builder,
+                                        @Nonnull JsonObject json,
+                                        @Nonnull ExtensionRegistry registry) {
+        final Descriptors.Descriptor descriptor = builder.getDescriptorForType();
+        for (Map.Entry<String, JsonElement> entry : json.entrySet()) {
+            final ExtensionRegistry.ExtensionInfo extension = registry.findImmutableExtensionByName(entry.getKey());
+            if (extension != null && descriptor.equals(extension.descriptor.getContainingType())) {
+                setExtension(builder, extension.descriptor, extension.defaultInstance, entry.getValue());
+                continue;
+            }
+            // not an extension of this message, but it may well be a message holding extensions further down
+            final Descriptors.FieldDescriptor field = findField(descriptor, entry.getKey());
+            if (field == null || field.getJavaType() != Descriptors.FieldDescriptor.JavaType.MESSAGE) {
+                continue;
+            }
+            if (field.isRepeated()) {
+                final JsonArray elements = entry.getValue().getAsJsonArray();
+                // the builder was parsed from this very array, so the two are in the same order
+                for (int i = 0; i < Math.min(elements.size(), builder.getRepeatedFieldCount(field)); i++) {
+                    mergeExtensions(builder.getRepeatedFieldBuilder(field, i), elements.get(i).getAsJsonObject(),
+                            registry);
+                }
+            } else {
+                mergeExtensions(builder.getFieldBuilder(field), entry.getValue().getAsJsonObject(), registry);
+            }
+        }
+    }
+
+    /**
+     * Parse the JSON representation of an extension value and set it on the given builder.
+     *
+     * @param builder the builder to set the extension on
+     * @param extension the field descriptor of the extension
+     * @param defaultInstance the default instance of the extension value, {@code null} unless it is a message
+     * @param value the JSON representation of the extension value
+     */
+    private static void setExtension(@Nonnull Message.Builder builder,
+                                     @Nonnull Descriptors.FieldDescriptor extension,
+                                     @Nullable Message defaultInstance,
+                                     @Nonnull JsonElement value) {
+        if (extension.isRepeated()) {
+            for (JsonElement element : value.getAsJsonArray()) {
+                builder.addRepeatedField(extension, extensionValue(extension, defaultInstance, element));
+            }
+        } else {
+            builder.setField(extension, extensionValue(extension, defaultInstance, value));
+        }
+    }
+
+    @Nonnull
+    private static Object extensionValue(@Nonnull Descriptors.FieldDescriptor extension,
+                                         @Nullable Message defaultInstance,
+                                         @Nonnull JsonElement value) {
+        switch (extension.getJavaType()) {
+            case MESSAGE:
+                // the fields of the extension value are regular fields, so JsonFormat handles them just fine
+                final Message.Builder valueBuilder = Objects.requireNonNull(defaultInstance).newBuilderForType();
+                try {
+                    JsonFormat.parser().ignoringUnknownFields().merge(value.toString(), valueBuilder);
+                } catch (InvalidProtocolBufferException e) {
+                    throw new RuntimeException("unable to parse value of extension " + extension.getFullName(), e);
+                }
+                return valueBuilder.build();
+            case BOOLEAN:
+                return value.getAsBoolean();
+            case INT:
+                return value.getAsInt();
+            case LONG:
+                return value.getAsLong();
+            case FLOAT:
+                return value.getAsFloat();
+            case DOUBLE:
+                return value.getAsDouble();
+            case STRING:
+                return value.getAsString();
+            case BYTE_STRING:
+                return ByteString.copyFrom(Base64.getDecoder().decode(value.getAsString()));
+            case ENUM:
+                final Descriptors.EnumValueDescriptor enumValue =
+                        extension.getEnumType().findValueByName(value.getAsString());
+                return Objects.requireNonNull(enumValue, () -> "unknown value of enum extension "
+                        + extension.getFullName() + ": " + value);
+            default:
+                throw new RuntimeException("unsupported type of extension " + extension.getFullName());
+        }
+    }
+
+    /**
+     * Find a field by the name the proto3 JSON mapping allows for it, which is either the name it is declared with or
+     * its lower camel case form.
+     *
+     * @param descriptor the message to look the field up in
+     * @param name the name of the field as it appears in the JSON
+     * @return the field, or {@code null} if the message has no such field
+     */
+    @Nullable
+    private static Descriptors.FieldDescriptor findField(@Nonnull Descriptors.Descriptor descriptor,
+                                                        @Nonnull String name) {
+        final Descriptors.FieldDescriptor byName = descriptor.findFieldByName(name);
+        if (byName != null) {
+            return byName;
+        }
+        for (Descriptors.FieldDescriptor field : descriptor.getFields()) {
+            if (field.getJsonName().equals(name)) {
+                return field;
+            }
+        }
+        return null;
+    }
+
     private static Pair<String, String> parseLoadTemplateString(String loadCommandString) {
         StringTokenizer lcsTokenizer = new StringTokenizer(loadCommandString, " ");
         if (lcsTokenizer.countTokens() != 3) {