feat(assertion): enforce strict urns for assertions (#15512)

david-leifker · web-flow · commit 47f2653f0efa · 2025-12-09T21:41:24.000-06:00
diff --git a/docs/how/updating-datahub.md b/docs/how/updating-datahub.md
@@ -9,6 +9,8 @@
 - #13726: Default search results per page (new: 5000, old: 10000) can be configured with environment variable `ELASTICSEARCH_LIMIT_RESULTS_API_DEFAULT`
 - #13726: Maximum lineage visualization hops (new: 20, old: 1000) can be configured with environment variable `ELASTICSEARCH_SEARCH_GRAPH_LINEAGE_MAX_HOPS`
 - #15044: If using the new system-update sqlSetup, the system-update job must be granted privileges to create users, databases, and tables.
+- #15512: Assertion related data models will enforce strict urn validation including selective existence, format, and supported entity types. This may impact
+          existing assertions which would need to be removed from the system.
 
 ### Known Issues
 
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/aspect/validation/AssertionUrnValidationAnnotationTest.java b/metadata-io/src/test/java/com/linkedin/metadata/aspect/validation/AssertionUrnValidationAnnotationTest.java
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionRunEvent.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/AssertionRunEvent.pdl
@@ -33,6 +33,11 @@ record AssertionRunEvent {
   * Urn of entity on which the assertion is applicable
   */
   @TimeseriesField = {}
+  @UrnValidation = {
+    "exist": false,
+    "strict": true,
+    "entityTypes": [ "dataset" ]
+  }
   asserteeUrn: Urn
 
   /**
@@ -60,6 +65,11 @@ record AssertionRunEvent {
   * Urn of assertion which is evaluated
   */
   @TimeseriesField = {}
+  @UrnValidation = {
+    "exist": false,
+    "strict": true,
+    "entityTypes": [ "assertion" ]
+  }
   assertionUrn: Urn
 
   /**
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/CustomAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/CustomAssertionInfo.pdl
@@ -23,6 +23,11 @@ record CustomAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset" ]
+    }
     entity: Urn
 
     /**
@@ -34,6 +39,11 @@ record CustomAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "schemaField" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "schemaField" ]
+    }
     field: optional Urn
 
     /*
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/DatasetAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/DatasetAssertionInfo.pdl
@@ -13,6 +13,11 @@ record DatasetAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset" ]
+    }
     dataset: Urn
 
     /**
@@ -62,6 +67,11 @@ record DatasetAssertionInfo {
         "fieldType": "URN"
       }
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "schemaField" ]
+    }
     fields: optional array[Urn]
 
     /**
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FieldAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FieldAssertionInfo.pdl
@@ -33,6 +33,11 @@ record FieldAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset" ]
+    }
     entity: Urn
 
     /**
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/FreshnessAssertionInfo.pdl
@@ -33,6 +33,11 @@ record FreshnessAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset", "dataJob" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset", "dataJob" ]
+    }
     entity: Urn
 
     /**
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/SchemaAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/SchemaAssertionInfo.pdl
@@ -17,6 +17,11 @@ record SchemaAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset", "dataJob" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset", "dataJob" ]
+    }
     entity: Urn
 
     /**
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/SqlAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/SqlAssertionInfo.pdl
@@ -33,6 +33,11 @@ record SqlAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset" ]
+    }
     entity: Urn
 
     /**
diff --git a/metadata-models/src/main/pegasus/com/linkedin/assertion/VolumeAssertionInfo.pdl b/metadata-models/src/main/pegasus/com/linkedin/assertion/VolumeAssertionInfo.pdl
@@ -47,6 +47,11 @@ record VolumeAssertionInfo {
       "name": "Asserts",
       "entityTypes": [ "dataset" ]
     }
+    @UrnValidation = {
+      "exist": false,
+      "strict": true,
+      "entityTypes": [ "dataset" ]
+    }
     entity: Urn
 
     /**
diff --git a/metadata-utils/src/main/java/com/linkedin/metadata/utils/UrnValidationUtil.java b/metadata-utils/src/main/java/com/linkedin/metadata/utils/UrnValidationUtil.java
@@ -221,6 +221,14 @@ private static void traverseDataMap(
       // Recursively traverse nested DataMaps
       if (value instanceof DataMap) {
         traverseDataMap((DataMap) value, fieldPath, urnValidationSpecs, result);
+      } else if (value instanceof DataList) {
+        // Recursively traverse DataMaps within arrays (e.g., array[RecordType])
+        DataList list = (DataList) value;
+        for (Object item : list) {
+          if (item instanceof DataMap) {
+            traverseDataMap((DataMap) item, fieldPath + "/*", urnValidationSpecs, result);
+          }
+        }
       }
     }
   }
diff --git a/metadata-utils/src/test/java/com/linkedin/metadata/utils/UrnValidationUtilTest.java b/metadata-utils/src/test/java/com/linkedin/metadata/utils/UrnValidationUtilTest.java
@@ -1,7 +1,18 @@
 package com.linkedin.metadata.utils;
 
+import static org.mockito.Mockito.*;
+
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.data.DataList;
+import com.linkedin.data.DataMap;
+import com.linkedin.data.schema.DataSchema;
+import com.linkedin.data.schema.PathSpec;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.metadata.aspect.batch.BatchItem;
+import com.linkedin.metadata.models.AspectSpec;
+import com.linkedin.metadata.models.UrnValidationFieldSpec;
+import com.linkedin.metadata.models.annotation.UrnValidationAnnotation;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import io.datahubproject.test.metadata.context.TestOperationContexts;
 import java.io.BufferedReader;
@@ -11,6 +22,7 @@
 import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import javax.annotation.Nonnull;
 import org.testng.Assert;
@@ -46,6 +58,196 @@ public void testUrnNull() {
     UrnValidationUtil.validateUrn(entityRegistry, null, true);
   }
 
+  // ==================== findUrnValidationFields Tests ====================
+
+  @Test
+  public void testFindUrnValidationFields_DirectField() {
+    // Test finding a URN field directly on the aspect
+    String testUrn = "urn:li:dataset:(urn:li:dataPlatform:hive,test,PROD)";
+    DataMap dataMap = new DataMap();
+    dataMap.put("entity", testUrn);
+
+    UrnValidationAnnotation annotation =
+        new UrnValidationAnnotation(true, true, List.of("dataset"));
+    PathSpec pathSpec = new PathSpec("entity");
+    DataSchema dataSchema = mock(DataSchema.class);
+    UrnValidationFieldSpec fieldSpec = new UrnValidationFieldSpec(pathSpec, annotation, dataSchema);
+
+    AspectSpec aspectSpec = mock(AspectSpec.class);
+    when(aspectSpec.getUrnValidationFieldSpecMap()).thenReturn(Map.of("/entity", fieldSpec));
+
+    RecordTemplate recordTemplate = mock(RecordTemplate.class);
+    when(recordTemplate.data()).thenReturn(dataMap);
+
+    BatchItem batchItem = mock(BatchItem.class);
+    when(batchItem.getRecordTemplate()).thenReturn(recordTemplate);
+    when(batchItem.getAspectSpec()).thenReturn(aspectSpec);
+
+    Set<UrnValidationUtil.UrnValidationEntry> entries =
+        UrnValidationUtil.findUrnValidationFields(batchItem, aspectSpec);
+
+    Assert.assertEquals(entries.size(), 1);
+    UrnValidationUtil.UrnValidationEntry entry = entries.iterator().next();
+    Assert.assertEquals(entry.getUrn(), testUrn);
+    Assert.assertEquals(entry.getFieldPath(), "/entity");
+  }
+
+  @Test
+  public void testFindUrnValidationFields_NestedRecord() {
+    // Test finding a URN field in a nested record (e.g., source.sourceUrn)
+    String testUrn = "urn:li:assertion:test-assertion";
+    DataMap sourceMap = new DataMap();
+    sourceMap.put("sourceUrn", testUrn);
+    sourceMap.put("type", "INFERRED_ASSERTION_FAILURE");
+
+    DataMap dataMap = new DataMap();
+    dataMap.put("source", sourceMap);
+
+    UrnValidationAnnotation annotation =
+        new UrnValidationAnnotation(true, true, List.of("assertion"));
+    PathSpec pathSpec = new PathSpec("source", "sourceUrn");
+    DataSchema dataSchema = mock(DataSchema.class);
+    UrnValidationFieldSpec fieldSpec = new UrnValidationFieldSpec(pathSpec, annotation, dataSchema);
+
+    AspectSpec aspectSpec = mock(AspectSpec.class);
+    when(aspectSpec.getUrnValidationFieldSpecMap())
+        .thenReturn(Map.of("/source/sourceUrn", fieldSpec));
+
+    RecordTemplate recordTemplate = mock(RecordTemplate.class);
+    when(recordTemplate.data()).thenReturn(dataMap);
+
+    BatchItem batchItem = mock(BatchItem.class);
+    when(batchItem.getRecordTemplate()).thenReturn(recordTemplate);
+    when(batchItem.getAspectSpec()).thenReturn(aspectSpec);
+
+    Set<UrnValidationUtil.UrnValidationEntry> entries =
+        UrnValidationUtil.findUrnValidationFields(batchItem, aspectSpec);
+
+    Assert.assertEquals(entries.size(), 1);
+    UrnValidationUtil.UrnValidationEntry entry = entries.iterator().next();
+    Assert.assertEquals(entry.getUrn(), testUrn);
+    Assert.assertEquals(entry.getFieldPath(), "/source/sourceUrn");
+  }
+
+  @Test
+  public void testFindUrnValidationFields_ArrayOfRecords() {
+    // Test finding URN fields inside an array of records (e.g., assertions[*].assertion)
+    String testUrn1 = "urn:li:assertion:assertion-1";
+    String testUrn2 = "urn:li:assertion:assertion-2";
+
+    DataMap record1 = new DataMap();
+    record1.put("assertion", testUrn1);
+
+    DataMap record2 = new DataMap();
+    record2.put("assertion", testUrn2);
+
+    DataList arrayList = new DataList();
+    arrayList.add(record1);
+    arrayList.add(record2);
+
+    DataMap dataMap = new DataMap();
+    dataMap.put("assertions", arrayList);
+
+    UrnValidationAnnotation annotation =
+        new UrnValidationAnnotation(true, true, List.of("assertion"));
+    PathSpec pathSpec = new PathSpec("assertions", PathSpec.WILDCARD, "assertion");
+    DataSchema dataSchema = mock(DataSchema.class);
+    UrnValidationFieldSpec fieldSpec = new UrnValidationFieldSpec(pathSpec, annotation, dataSchema);
+
+    AspectSpec aspectSpec = mock(AspectSpec.class);
+    when(aspectSpec.getUrnValidationFieldSpecMap())
+        .thenReturn(Map.of("/assertions/*/assertion", fieldSpec));
+
+    RecordTemplate recordTemplate = mock(RecordTemplate.class);
+    when(recordTemplate.data()).thenReturn(dataMap);
+
+    BatchItem batchItem = mock(BatchItem.class);
+    when(batchItem.getRecordTemplate()).thenReturn(recordTemplate);
+    when(batchItem.getAspectSpec()).thenReturn(aspectSpec);
+
+    Set<UrnValidationUtil.UrnValidationEntry> entries =
+        UrnValidationUtil.findUrnValidationFields(batchItem, aspectSpec);
+
+    Assert.assertEquals(entries.size(), 2, "Should find URNs in both array elements");
+    Set<String> urns =
+        entries.stream()
+            .map(UrnValidationUtil.UrnValidationEntry::getUrn)
+            .collect(java.util.stream.Collectors.toSet());
+    Assert.assertTrue(urns.contains(testUrn1));
+    Assert.assertTrue(urns.contains(testUrn2));
+  }
+
+  @Test
+  public void testFindUrnValidationFields_ArrayOfUrns() {
+    // Test finding URN fields that are directly an array of URNs (e.g., fields: array[Urn])
+    String testUrn1 = "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,test,PROD),f1)";
+    String testUrn2 = "urn:li:schemaField:(urn:li:dataset:(urn:li:dataPlatform:hive,test,PROD),f2)";
+
+    DataList urnList = new DataList();
+    urnList.add(testUrn1);
+    urnList.add(testUrn2);
+
+    DataMap dataMap = new DataMap();
+    dataMap.put("fields", urnList);
+
+    UrnValidationAnnotation annotation =
+        new UrnValidationAnnotation(false, true, List.of("schemaField"));
+    PathSpec pathSpec = new PathSpec("fields");
+    DataSchema dataSchema = mock(DataSchema.class);
+    UrnValidationFieldSpec fieldSpec = new UrnValidationFieldSpec(pathSpec, annotation, dataSchema);
+
+    AspectSpec aspectSpec = mock(AspectSpec.class);
+    when(aspectSpec.getUrnValidationFieldSpecMap()).thenReturn(Map.of("/fields", fieldSpec));
+
+    RecordTemplate recordTemplate = mock(RecordTemplate.class);
+    when(recordTemplate.data()).thenReturn(dataMap);
+
+    BatchItem batchItem = mock(BatchItem.class);
+    when(batchItem.getRecordTemplate()).thenReturn(recordTemplate);
+    when(batchItem.getAspectSpec()).thenReturn(aspectSpec);
+
+    Set<UrnValidationUtil.UrnValidationEntry> entries =
+        UrnValidationUtil.findUrnValidationFields(batchItem, aspectSpec);
+
+    Assert.assertEquals(entries.size(), 2, "Should find both URNs in the array");
+    Set<String> urns =
+        entries.stream()
+            .map(UrnValidationUtil.UrnValidationEntry::getUrn)
+            .collect(java.util.stream.Collectors.toSet());
+    Assert.assertTrue(urns.contains(testUrn1));
+    Assert.assertTrue(urns.contains(testUrn2));
+  }
+
+  @Test
+  public void testFindUrnValidationFields_EmptyArray() {
+    // Test with an empty array - should return no entries
+    DataList emptyList = new DataList();
+    DataMap dataMap = new DataMap();
+    dataMap.put("assertions", emptyList);
+
+    UrnValidationAnnotation annotation =
+        new UrnValidationAnnotation(true, true, List.of("assertion"));
+    PathSpec pathSpec = new PathSpec("assertions", PathSpec.WILDCARD, "assertion");
+    DataSchema dataSchema = mock(DataSchema.class);
+    UrnValidationFieldSpec fieldSpec = new UrnValidationFieldSpec(pathSpec, annotation, dataSchema);
+
+    AspectSpec aspectSpec = mock(AspectSpec.class);
+    when(aspectSpec.getUrnValidationFieldSpecMap())
+        .thenReturn(Map.of("/assertions/*/assertion", fieldSpec));
+
+    RecordTemplate recordTemplate = mock(RecordTemplate.class);
+    when(recordTemplate.data()).thenReturn(dataMap);
+
+    BatchItem batchItem = mock(BatchItem.class);
+    when(batchItem.getRecordTemplate()).thenReturn(recordTemplate);
+    when(batchItem.getAspectSpec()).thenReturn(aspectSpec);
+
+    Set<UrnValidationUtil.UrnValidationEntry> entries =
+        UrnValidationUtil.findUrnValidationFields(batchItem, aspectSpec);
+
+    Assert.assertEquals(entries.size(), 0, "Empty array should return no entries");
+  }
+
   /**
    * Common method to validate URNs from a file and return the validation results.
    *
