fix(proto_optional_fields): loosen schema validators to only check fo…

…r existence of previously required fields
jsdonn · May 23, 2024 · 2b34a36 · 2b34a36
1 parent c79500a
commit 2b34a36
Show file tree

Hide file tree

Showing 8 changed files with 67 additions and 44 deletions.
diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/ModelUtils.java
@@ -38,6 +38,8 @@ public class ModelUtils {
 
   private static final ClassLoader CLASS_LOADER = DummySnapshot.class.getClassLoader();
   private static final String METADATA_AUDIT_EVENT_PREFIX = "METADATA_AUDIT_EVENT";
+  private static final String URN_FIELD = "urn";
+  private static final String ASPECTS_FIELD = "aspects";
 
   private ModelUtils() {
     // Util class
@@ -135,7 +137,11 @@ public static Class<? extends RecordTemplate> getMetadataSnapshotClassFromName(@
   @Nonnull
   public static <SNAPSHOT extends RecordTemplate> Urn getUrnFromSnapshot(@Nonnull SNAPSHOT snapshot) {
     SnapshotValidator.validateSnapshotSchema(snapshot.getClass());
-    return RecordUtils.getRecordTemplateField(snapshot, "urn", urnClassForSnapshot(snapshot.getClass()));
+    final Urn urn = RecordUtils.getRecordTemplateField(snapshot, URN_FIELD, urnClassForSnapshot(snapshot.getClass()));
+    if (urn == null) {
+      ValidationUtils.throwNullFieldException(URN_FIELD);
+    }
+    return urn;
   }
 
   /**
@@ -191,7 +197,7 @@ public static Urn getUrnFromSnapshotUnion(@Nonnull UnionTemplate snapshotUnion)
   @Nonnull
   public static <DELTA extends RecordTemplate> Urn getUrnFromDelta(@Nonnull DELTA delta) {
     DeltaValidator.validateDeltaSchema(delta.getClass());
-    return RecordUtils.getRecordTemplateField(delta, "urn", urnClassForDelta(delta.getClass()));
+    return RecordUtils.getRecordTemplateField(delta, URN_FIELD, urnClassForDelta(delta.getClass()));
   }
 
   /**
@@ -212,7 +218,11 @@ public static Urn getUrnFromDeltaUnion(@Nonnull UnionTemplate deltaUnion) {
   @Nonnull
   public static <DOCUMENT extends RecordTemplate> Urn getUrnFromDocument(@Nonnull DOCUMENT document) {
     DocumentValidator.validateDocumentSchema(document.getClass());
-    return RecordUtils.getRecordTemplateField(document, "urn", urnClassForDocument(document.getClass()));
+    final Urn urn = RecordUtils.getRecordTemplateField(document, URN_FIELD, urnClassForDocument(document.getClass()));
+    if (urn == null) {
+      ValidationUtils.throwNullFieldException(URN_FIELD);
+    }
+    return urn;
   }
 
   /**
@@ -225,7 +235,11 @@ public static <DOCUMENT extends RecordTemplate> Urn getUrnFromDocument(@Nonnull
   @Nonnull
   public static <ENTITY extends RecordTemplate> Urn getUrnFromEntity(@Nonnull ENTITY entity) {
     EntityValidator.validateEntitySchema(entity.getClass());
-    return RecordUtils.getRecordTemplateField(entity, "urn", urnClassForDocument(entity.getClass()));
+    final Urn urn = RecordUtils.getRecordTemplateField(entity, URN_FIELD, urnClassForDocument(entity.getClass()));
+    if (urn == null) {
+      ValidationUtils.throwNullFieldException(URN_FIELD);
+    }
+    return urn;
   }
 
   /**
@@ -240,8 +254,11 @@ public static <ENTITY extends RecordTemplate> Urn getUrnFromEntity(@Nonnull ENTI
   private static <RELATIONSHIP extends RecordTemplate> Urn getUrnFromRelationship(@Nonnull RELATIONSHIP relationship,
       @Nonnull String fieldName) {
     RelationshipValidator.validateRelationshipSchema(relationship.getClass());
-    return RecordUtils.getRecordTemplateField(relationship, fieldName,
-        urnClassForRelationship(relationship.getClass(), fieldName));
+    final Urn urn = RecordUtils.getRecordTemplateField(relationship, fieldName, urnClassForRelationship(relationship.getClass(), fieldName));
+    if (urn == null) {
+      ValidationUtils.throwNullFieldException(URN_FIELD);
+    }
+    return urn;
   }
 
   /**
@@ -272,7 +289,6 @@ public static <RELATIONSHIP extends RecordTemplate> Urn getDestinationUrnFromRel
   @Nonnull
   public static <SNAPSHOT extends RecordTemplate> List<RecordTemplate> getAspectsFromSnapshot(
       @Nonnull SNAPSHOT snapshot) {
-
     SnapshotValidator.validateSnapshotSchema(snapshot.getClass());
     return getAspects(snapshot);
   }
@@ -307,7 +323,10 @@ public static List<RecordTemplate> getAspectsFromSnapshotUnion(@Nonnull UnionTem
   private static List<RecordTemplate> getAspects(@Nonnull RecordTemplate snapshot) {
     final Class<? extends WrappingArrayTemplate> clazz = getAspectsArrayClass(snapshot.getClass());
 
-    WrappingArrayTemplate aspectArray = RecordUtils.getRecordTemplateWrappedField(snapshot, "aspects", clazz);
+    final WrappingArrayTemplate aspectArray = RecordUtils.getRecordTemplateWrappedField(snapshot, ASPECTS_FIELD, clazz);
+    if (aspectArray == null) {
+      ValidationUtils.throwNullFieldException(ASPECTS_FIELD);
+    }
 
     final List<RecordTemplate> aspects = new ArrayList<>();
     aspectArray.forEach(item -> aspects.add(RecordUtils.getSelectedRecordTemplateFromUnion((UnionTemplate) item)));
@@ -352,10 +371,16 @@ public static <SNAPSHOT extends RecordTemplate, ASPECT_UNION extends UnionTempla
 
     try {
       final SNAPSHOT snapshot = snapshotClass.newInstance();
-      RecordUtils.setRecordTemplatePrimitiveField(snapshot, "urn", urn);
+      if (urn == null) {
+        ValidationUtils.throwNullFieldException(URN_FIELD);
+      }
+      if (aspects == null) {
+        ValidationUtils.throwNullFieldException(ASPECTS_FIELD);
+      }
+      RecordUtils.setRecordTemplatePrimitiveField(snapshot, URN_FIELD, urn);
       WrappingArrayTemplate aspectArray = aspectArrayClass.newInstance();
       aspectArray.addAll(aspects);
-      RecordUtils.setRecordTemplateComplexField(snapshot, "aspects", aspectArray);
+      RecordUtils.setRecordTemplateComplexField(snapshot, ASPECTS_FIELD, aspectArray);
       return snapshot;
     } catch (InstantiationException | IllegalAccessException e) {
       throw new RuntimeException(e);
@@ -418,7 +443,7 @@ public static Class<? extends UnionTemplate> aspectClassForSnapshot(
     SnapshotValidator.validateSnapshotSchema(snapshotClass);
 
     String aspectClassName = ((TyperefDataSchema) ((ArrayDataSchema) ValidationUtils.getRecordSchema(snapshotClass)
-        .getField("aspects")
+        .getField(ASPECTS_FIELD)
         .getType()).getItems()).getBindingName();
 
     return getClassFromName(aspectClassName, UnionTemplate.class);
@@ -430,7 +455,7 @@ public static Class<? extends UnionTemplate> aspectClassForSnapshot(
   @Nonnull
   public static Class<? extends Urn> urnClassForEntity(@Nonnull Class<? extends RecordTemplate> entityClass) {
     EntityValidator.validateEntitySchema(entityClass);
-    return urnClassForField(entityClass, "urn");
+    return urnClassForField(entityClass, URN_FIELD);
   }
 
   /**
@@ -439,7 +464,7 @@ public static Class<? extends Urn> urnClassForEntity(@Nonnull Class<? extends Re
   @Nonnull
   public static Class<? extends Urn> urnClassForSnapshot(@Nonnull Class<? extends RecordTemplate> snapshotClass) {
     SnapshotValidator.validateSnapshotSchema(snapshotClass);
-    return urnClassForField(snapshotClass, "urn");
+    return urnClassForField(snapshotClass, URN_FIELD);
   }
 
   /**
@@ -448,7 +473,7 @@ public static Class<? extends Urn> urnClassForSnapshot(@Nonnull Class<? extends
   @Nonnull
   public static Class<? extends Urn> urnClassForDelta(@Nonnull Class<? extends RecordTemplate> deltaClass) {
     DeltaValidator.validateDeltaSchema(deltaClass);
-    return urnClassForField(deltaClass, "urn");
+    return urnClassForField(deltaClass, URN_FIELD);
   }
 
   /**
@@ -457,7 +482,7 @@ public static Class<? extends Urn> urnClassForDelta(@Nonnull Class<? extends Rec
   @Nonnull
   public static Class<? extends Urn> urnClassForDocument(@Nonnull Class<? extends RecordTemplate> documentClass) {
     DocumentValidator.validateDocumentSchema(documentClass);
-    return urnClassForField(documentClass, "urn");
+    return urnClassForField(documentClass, URN_FIELD);
   }
 
   /**

diff --git a/dao-api/src/main/java/com/linkedin/metadata/dao/utils/RecordUtils.java b/dao-api/src/main/java/com/linkedin/metadata/dao/utils/RecordUtils.java
@@ -281,7 +281,7 @@ public static <T extends RecordTemplate, V extends DataTemplate> void setRecordT
    * @param valueClass the expected type for the value
    * @return the value for the field
    */
-  @Nonnull
+  @Nullable
   public static <T extends RecordTemplate, V> V getRecordTemplateField(@Nonnull T recordTemplate,
       @Nonnull String fieldName, @Nonnull Class<V> valueClass) {
 
@@ -300,7 +300,7 @@ public static <T extends RecordTemplate, V> V getRecordTemplateField(@Nonnull T
    * @param valueClass the expected type for the value
    * @return the value for the field
    */
-  @Nonnull
+  @Nullable
   public static <T extends RecordTemplate, V extends DataTemplate> V getRecordTemplateWrappedField(
       @Nonnull T recordTemplate, @Nonnull String fieldName, @Nonnull Class<V> valueClass) {
 
@@ -389,7 +389,7 @@ private static Method getProtectedMethod(@Nonnull Class clazz, @Nonnull String m
     }
   }
 
-  @Nonnull
+  @Nullable
   private static <T> T invokeProtectedMethod(Object object, Method method, Object... args) {
     try {
       method.setAccessible(true);

diff --git a/validators/src/main/java/com/linkedin/metadata/validator/DocumentValidator.java b/validators/src/main/java/com/linkedin/metadata/validator/DocumentValidator.java
@@ -38,17 +38,13 @@ public static void validateDocumentSchema(@Nonnull RecordDataSchema schema) {
     final String className = schema.getBindingName();
 
     if (!ValidationUtils.schemaHasExactlyOneSuchField(schema, ValidationUtils::isValidUrnField)) {
-      ValidationUtils.invalidSchema("Document '%s' must contain an non-optional 'urn' field of URN type", className);
+      ValidationUtils.invalidSchema("Document '%s' must contain an 'urn' field of URN type", className);
     }
 
     ValidationUtils.fieldsUsingInvalidType(schema, ValidationUtils.PRIMITIVE_TYPES).forEach(field -> {
       ValidationUtils.invalidSchema("Document '%s' contains a field '%s' that makes use of a disallowed type '%s'.",
           className, field.getName(), field.getType().getType());
     });
-
-    ValidationUtils.nonOptionalFields(schema, NON_OPTIONAL_FIELDS).forEach(field -> {
-      ValidationUtils.invalidSchema("Document '%s' must contain an optional '%s' field", className, field.getName());
-    });
   }
 
   /**

diff --git a/validators/src/main/java/com/linkedin/metadata/validator/EntityValidator.java b/validators/src/main/java/com/linkedin/metadata/validator/EntityValidator.java
@@ -4,8 +4,6 @@
 import com.linkedin.data.schema.UnionDataSchema;
 import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.data.template.UnionTemplate;
-import java.util.Collections;
-import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import javax.annotation.Nonnull;
@@ -16,13 +14,6 @@
  */
 public final class EntityValidator {
 
-  // Allowed non-optional fields. All other fields must be optional.
-  private static final Set<String> NON_OPTIONAL_FIELDS = Collections.unmodifiableSet(new HashSet<String>() {
-    {
-      add("urn");
-    }
-  });
-
   // A cache of validated classes
   private static final Set<Class<? extends RecordTemplate>> VALIDATED = ConcurrentHashMap.newKeySet();
 
@@ -44,17 +35,13 @@ public static void validateEntitySchema(@Nonnull RecordDataSchema schema) {
     final String className = schema.getBindingName();
 
     if (!ValidationUtils.schemaHasExactlyOneSuchField(schema, ValidationUtils::isValidUrnField)) {
-      ValidationUtils.invalidSchema("Entity '%s' must contain a non-optional 'urn' field of URN type", className);
+      ValidationUtils.invalidSchema("Entity '%s' must contain an 'urn' field of URN type", className);
     }
 
     ValidationUtils.fieldsUsingInvalidType(schema, ValidationUtils.PRIMITIVE_TYPES).forEach(field -> {
       ValidationUtils.invalidSchema("Entity '%s' contains a field '%s' that makes use of a disallowed type '%s'.",
           className, field.getName(), field.getType().getType());
     });
-
-    ValidationUtils.nonOptionalFields(schema, NON_OPTIONAL_FIELDS).forEach(field -> {
-      ValidationUtils.invalidSchema("Entity '%s' must contain an optional '%s' field", className, field.getName());
-    });
   }
 
   /**

diff --git a/validators/src/main/java/com/linkedin/metadata/validator/NullFieldException.java b/validators/src/main/java/com/linkedin/metadata/validator/NullFieldException.java
@@ -0,0 +1,11 @@
+package com.linkedin.metadata.validator;
+
+/**
+ * Thrown when a field which is not supposed to be null is null.
+ */
+public class NullFieldException extends RuntimeException {
+
+  public NullFieldException(String message) {
+    super(message);
+  }
+}
diff --git a/validators/src/main/java/com/linkedin/metadata/validator/RelationshipValidator.java b/validators/src/main/java/com/linkedin/metadata/validator/RelationshipValidator.java
@@ -44,13 +44,13 @@ public static void validateRelationshipSchema(@Nonnull RecordDataSchema schema)
 
     if (!ValidationUtils.schemaHasExactlyOneSuchField(schema,
         field -> ValidationUtils.isValidUrnField(field, "source"))) {
-      ValidationUtils.invalidSchema("Relationship '%s' must contain an non-optional 'source' field of URN type",
+      ValidationUtils.invalidSchema("Relationship '%s' must contain a 'source' field of URN type",
           className);
     }
 
     if (!ValidationUtils.schemaHasExactlyOneSuchField(schema,
         field -> ValidationUtils.isValidUrnField(field, "destination"))) {
-      ValidationUtils.invalidSchema("Relationship '%s' must contain an non-optional 'destination' field of URN type",
+      ValidationUtils.invalidSchema("Relationship '%s' must contain a 'destination' field of URN type",
           className);
     }
 

diff --git a/validators/src/main/java/com/linkedin/metadata/validator/SnapshotValidator.java b/validators/src/main/java/com/linkedin/metadata/validator/SnapshotValidator.java
@@ -33,11 +33,11 @@ public static void validateSnapshotSchema(@Nonnull RecordDataSchema schema) {
     final String className = schema.getBindingName();
 
     if (!ValidationUtils.schemaHasExactlyOneSuchField(schema, ValidationUtils::isValidUrnField)) {
-      ValidationUtils.invalidSchema("Snapshot '%s' must contain an non-optional 'urn' field of URN type", className);
+      ValidationUtils.invalidSchema("Snapshot '%s' must contain an 'urn' field of URN type", className);
     }
 
     if (!ValidationUtils.schemaHasExactlyOneSuchField(schema, SnapshotValidator::isValidAspectsField)) {
-      ValidationUtils.invalidSchema("Snapshot '%s' must contain an non-optional 'aspects' field of ARRAY type",
+      ValidationUtils.invalidSchema("Snapshot '%s' must contain an 'aspects' field of ARRAY type",
           className);
     }
 
@@ -75,7 +75,7 @@ public static void validateUniqueUrn(@Nonnull Collection<Class<? extends RecordT
   }
 
   private static boolean isValidAspectsField(@Nonnull RecordDataSchema.Field field) {
-    return field.getName().equals("aspects") && !field.getOptional()
+    return field.getName().equals("aspects")
         && field.getType().getType() == DataSchema.Type.ARRAY;
   }
 

diff --git a/validators/src/main/java/com/linkedin/metadata/validator/ValidationUtils.java b/validators/src/main/java/com/linkedin/metadata/validator/ValidationUtils.java
@@ -89,10 +89,10 @@ public static boolean schemaHasExactlyOneSuchField(@Nonnull RecordDataSchema sch
   }
 
   /**
-   * Returns true if the non-optional field matches the field name and has a URN type.
+   * Returns true if the passed-in field matches the field name and is of URN type.
    */
   public static boolean isValidUrnField(@Nonnull RecordDataSchema.Field field, @Nonnull String fieldName) {
-    return field.getName().equals(fieldName) && !field.getOptional()
+    return field.getName().equals(fieldName)
         && field.getType().getType() == DataSchema.Type.TYPEREF && Urn.class.isAssignableFrom(getUrnClass(field));
   }
 
@@ -197,4 +197,8 @@ private static DataSchema.Type getFieldOrArrayItemType(@Nonnull RecordDataSchema
     }
     return type.getType();
   }
+
+  public static void throwNullFieldException(String field) {
+    throw new NullFieldException(String.format("The %s field cannot be null.", field));
+  }
 }