fix(assertions): minor changes in custom assertion api (datahub-proje…

…ct#10794)
ajoymajumdar · Jul 1, 2024 · 3056cd8 · 3056cd8
1 parent 3e47496
commit 3056cd8
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 23 deletions.
diff --git a/.../java/com/linkedin/datahub/graphql/resolvers/assertion/ReportAssertionResultResolver.java b/.../java/com/linkedin/datahub/graphql/resolvers/assertion/ReportAssertionResultResolver.java
@@ -72,9 +72,10 @@ public CompletableFuture<Boolean> get(DataFetchingEnvironment environment) throw
                 context.getOperationContext(),
                 assertionUrn,
                 asserteeUrn,
-                input.getTimestampMillis(),
-                assertionResult,
-                mapContextParameters(input.getProperties()));
+                input.getTimestampMillis() != null
+                    ? input.getTimestampMillis()
+                    : System.currentTimeMillis(),
+                assertionResult);
             return true;
           }
           throw new AuthorizationException(
@@ -99,6 +100,9 @@ private AssertionResult mapAssertionResult(AssertionResultInput input) {
     if (assertionResult.getType() == AssertionResultType.ERROR && input.getError() != null) {
       assertionResult.setError(mapAssertionResultError(input));
     }
+    if (input.getProperties() != null) {
+      assertionResult.setNativeResults(mapContextParameters(input.getProperties()));
+    }
     return assertionResult;
   }
 

diff --git a/datahub-graphql-core/src/main/resources/assertions.graphql b/datahub-graphql-core/src/main/resources/assertions.graphql
@@ -87,7 +87,9 @@ input AssertionResultInput {
   type: AssertionResultType!
 
   """
-  Additional key-value pairs representing runtime context
+  Additional metadata representing about the native results of the assertion.
+  These will be displayed alongside the result.
+  It should be used to capture additional context that is useful for the user.
   """
   properties: [StringMapEntryInput!]
 

diff --git a/...a/com/linkedin/datahub/graphql/resolvers/assertion/ReportAssertionResultResolverTest.java b/...a/com/linkedin/datahub/graphql/resolvers/assertion/ReportAssertionResultResolverTest.java
@@ -61,8 +61,8 @@ public class ReportAssertionResultResolverTest {
                           .setType(com.linkedin.assertion.AssertionResultErrorType.UNKNOWN_ERROR)
                           .setProperties(
                               new StringMap(Map.of("message", "an unknown error occurred"))))
-                  .setExternalUrl(customAssertionUrl))
-          .setRuntimeContext(new StringMap(Map.of("prop1", "value1")));
+                  .setExternalUrl(customAssertionUrl)
+                  .setNativeResults(new StringMap(Map.of("prop1", "value1"))));
 
   @Test
   public void testGetSuccessReportAssertionRunEvent() throws Exception {
@@ -91,8 +91,7 @@ public void testGetSuccessReportAssertionRunEvent() throws Exception {
             Mockito.eq(TEST_ASSERTION_URN),
             Mockito.eq(TEST_DATASET_URN),
             Mockito.eq(TEST_ASSERTION_RUN_EVENT.getTimestampMillis()),
-            Mockito.eq(TEST_ASSERTION_RUN_EVENT.getResult()),
-            Mockito.eq(TEST_ASSERTION_RUN_EVENT.getRuntimeContext()));
+            Mockito.eq(TEST_ASSERTION_RUN_EVENT.getResult()));
   }
 
   @Test
@@ -126,7 +125,6 @@ public void testGetUpdateAssertionUnauthorized() throws Exception {
             Mockito.any(),
             Mockito.any(),
             Mockito.any(),
-            Mockito.any(),
             Mockito.any());
   }
 
@@ -146,7 +144,6 @@ public void testGetAssertionServiceException() {
             Mockito.any(),
             Mockito.any(),
             Mockito.any(),
-            Mockito.any(),
             Mockito.any());
 
     ReportAssertionResultResolver resolver = new ReportAssertionResultResolver(mockService);

diff --git a/metadata-io/src/test/java/com/linkedin/metadata/service/AssertionServiceTest.java b/metadata-io/src/test/java/com/linkedin/metadata/service/AssertionServiceTest.java
@@ -1,7 +1,6 @@
 package com.linkedin.metadata.service;
 
 import static com.linkedin.metadata.Constants.*;
-import static com.linkedin.metadata.service.AssertionService.*;
 import static org.mockito.ArgumentMatchers.*;
 import static org.mockito.Mockito.*;
 
@@ -221,8 +220,7 @@ public void testAddAssertionRunEventRequiredFields() throws Exception, RemoteInv
         TEST_ASSERTION_URN,
         TEST_DATASET_URN,
         eventtime,
-        new AssertionResult().setType(AssertionResultType.SUCCESS),
-        null);
+        new AssertionResult().setType(AssertionResultType.SUCCESS));
 
     Mockito.verify(mockedEntityClient, Mockito.times(1))
         .ingestProposal(any(OperationContext.class), Mockito.any(), Mockito.eq(false));
@@ -234,7 +232,7 @@ public void testAddAssertionRunEventAllFields() throws Exception, RemoteInvocati
     AssertionService assertionService =
         new AssertionService(mockedEntityClient, mock(GraphClient.class));
     Long eventtime = 1718619000000L;
-    StringMap customProps = new StringMap(Map.of("prop-1", "value-1"));
+    StringMap nativeResults = new StringMap(Map.of("prop-1", "value-1"));
     StringMap errorProps = new StringMap(Map.of("message", "errorMessage"));
     String externalUrlOfAssertion = "https://abc/xyz";
 
@@ -254,7 +252,7 @@ public void testAddAssertionRunEventAllFields() throws Exception, RemoteInvocati
               Assert.assertEquals(runEvent.getAsserteeUrn(), TEST_DATASET_URN);
               Assert.assertEquals(runEvent.getTimestampMillis(), eventtime);
               Assert.assertEquals(runEvent.getStatus(), AssertionRunStatus.COMPLETE);
-              Assert.assertEquals(runEvent.getRuntimeContext(), customProps);
+              Assert.assertEquals(runEvent.getResult().getNativeResults(), nativeResults);
 
               AssertionResult result = runEvent.getResult();
               Assert.assertEquals(result.getType(), AssertionResultType.ERROR);
@@ -276,11 +274,11 @@ public void testAddAssertionRunEventAllFields() throws Exception, RemoteInvocati
         new AssertionResult()
             .setType(AssertionResultType.ERROR)
             .setExternalUrl(externalUrlOfAssertion)
+            .setNativeResults(nativeResults)
             .setError(
                 new AssertionResultError()
                     .setType(AssertionResultErrorType.UNKNOWN_ERROR)
-                    .setProperties(errorProps)),
-        customProps);
+                    .setProperties(errorProps)));
 
     Mockito.verify(mockedEntityClient, Mockito.times(1))
         .ingestProposal(any(OperationContext.class), Mockito.any(), Mockito.eq(false));

diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/service/AssertionService.java b/metadata-service/services/src/main/java/com/linkedin/metadata/service/AssertionService.java
@@ -14,7 +14,6 @@
 import com.linkedin.common.EntityRelationships;
 import com.linkedin.common.url.Url;
 import com.linkedin.common.urn.Urn;
-import com.linkedin.data.template.StringMap;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.SystemEntityClient;
 import com.linkedin.metadata.Constants;
@@ -157,18 +156,14 @@ public void addAssertionRunEvent(
       @Nonnull Urn assertionUrn,
       @Nonnull Urn asserteeUrn,
       @Nonnull Long timestampMillis,
-      @Nonnull AssertionResult assertionResult,
-      @Nullable StringMap contextParameters) {
+      @Nonnull AssertionResult assertionResult) {
     AssertionRunEvent assertionRunEvent = new AssertionRunEvent();
     assertionRunEvent.setTimestampMillis(timestampMillis);
     assertionRunEvent.setRunId(timestampMillis.toString());
     assertionRunEvent.setAssertionUrn(assertionUrn);
     assertionRunEvent.setAsserteeUrn(asserteeUrn);
     assertionRunEvent.setStatus(AssertionRunStatus.COMPLETE);
     assertionRunEvent.setResult(assertionResult);
-    if (contextParameters != null) {
-      assertionRunEvent.setRuntimeContext(contextParameters);
-    }
 
     try {
       this.entityClient.ingestProposal(