From 8841581d323f93bb7bc4fe9b24c3dee55c2c34eb Mon Sep 17 00:00:00 2001
From: Mayuri N <mayuri.nehate@gslab.com>
Date: Wed, 7 Feb 2024 11:01:31 +0530
Subject: [PATCH] feat(graphql): get raw aspects for assertions, allow
 aspectNames filter

---
 .../linkedin/datahub/graphql/GmsGraphQLEngine.java   |  4 +++-
 .../datahub/graphql/WeaklyTypedAspectsResolver.java  |  7 ++++++-
 .../src/main/resources/entity.graphql                | 12 ++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
index e3eef0688c2692..28b3a982c7b283 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -2436,7 +2436,9 @@ private void configureAssertionResolvers(final RuntimeWiring.Builder builder) {
                               ? assertion.getDataPlatformInstance().getUrn()
                               : null;
                         }))
-                .dataFetcher("runEvents", new AssertionRunEventResolver(entityClient)));
+                .dataFetcher("runEvents", new AssertionRunEventResolver(entityClient))
+                .dataFetcher(
+                    "aspects", new WeaklyTypedAspectsResolver(entityClient, entityRegistry)));
   }
 
   private void configurePolicyResolvers(final RuntimeWiring.Builder builder) {
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/WeaklyTypedAspectsResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/WeaklyTypedAspectsResolver.java
index d8665ae784bd1e..fd23cd5fdda450 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/WeaklyTypedAspectsResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/WeaklyTypedAspectsResolver.java
@@ -37,7 +37,12 @@ public class WeaklyTypedAspectsResolver implements DataFetcher<CompletableFuture
   private static final JacksonDataCodec CODEC = new JacksonDataCodec();
 
   private boolean shouldReturnAspect(AspectSpec aspectSpec, AspectParams params) {
-    return !params.getAutoRenderOnly() || aspectSpec.isAutoRender();
+    return (params.getAutoRenderOnly() == null
+            || !params.getAutoRenderOnly()
+            || aspectSpec.isAutoRender())
+        && (params.getAspectNames() == null
+            || params.getAspectNames().isEmpty()
+            || params.getAspectNames().contains(aspectSpec.getName()));
   }
 
   @Override
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
index 0074dc3fcb44c7..762514b480fcae 100644
--- a/datahub-graphql-core/src/main/resources/entity.graphql
+++ b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -1763,6 +1763,12 @@ input AspectParams {
     Only fetch auto render aspects
     """
     autoRenderOnly: Boolean
+
+    """
+    Fetch using aspect names
+    If absent, returns all aspects matching other inputs
+    """
+    aspectNames: [String!]
 }
 
 
@@ -6788,6 +6794,12 @@ type Assertion implements EntityWithRelationships & Entity {
     Edges extending from this entity grouped by direction in the lineage graph
     """
     lineage(input: LineageInput!): EntityLineageResult
+
+    """
+    Experimental API.
+    For fetching extra aspects that do not have custom UI code yet
+    """
+    aspects(input: AspectParams): [RawAspect!]
 }
 
 """