From 4f6bd96e10132ebf1888ea3fc5e553fe578a1e60 Mon Sep 17 00:00:00 2001
From: si-chakraborty <si.chakraborty@adevinta.com>
Date: Tue, 3 Oct 2023 16:40:00 +0200
Subject: [PATCH] feat(grahql): Added datafetcher for DataPlatformInstance
 entity

---
 .../datahub/graphql/GmsGraphQLEngine.java     |  1 +
 .../DataPlatformInstanceType.java             | 34 ++++++++++++++++++-
 .../src/main/resources/entity.graphql         |  5 +++
 3 files changed, 39 insertions(+), 1 deletion(-)

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 682710ad5d539..86a37855f89e8 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
@@ -815,6 +815,7 @@ private void configureQueryResolvers(final RuntimeWiring.Builder builder) {
             .dataFetcher("glossaryNode", getResolver(glossaryNodeType))
             .dataFetcher("domain", getResolver((domainType)))
             .dataFetcher("dataPlatform", getResolver(dataPlatformType))
+            .dataFetcher("dataPlatformInstance", getResolver(dataPlatformInstanceType))
             .dataFetcher("mlFeatureTable", getResolver(mlFeatureTableType))
             .dataFetcher("mlFeature", getResolver(mlFeatureType))
             .dataFetcher("mlPrimaryKey", getResolver(mlPrimaryKeyType))
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java
index 2423fc31ea52e..87614e1332528 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/dataplatforminstance/DataPlatformInstanceType.java
@@ -4,16 +4,25 @@
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
 import com.linkedin.datahub.graphql.generated.DataPlatformInstance;
 import com.linkedin.datahub.graphql.generated.Entity;
 import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.FacetFilterInput;
+import com.linkedin.datahub.graphql.generated.SearchResults;
 import com.linkedin.datahub.graphql.types.dataplatforminstance.mappers.DataPlatformInstanceMapper;
+import com.linkedin.datahub.graphql.types.mappers.AutoCompleteResultsMapper;
+import com.linkedin.datahub.graphql.types.SearchableEntityType;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.AutoCompleteResult;
+import com.linkedin.metadata.query.filter.Filter;
 import graphql.execution.DataFetcherResult;
+import org.apache.commons.lang3.NotImplementedException;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
@@ -22,7 +31,10 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 
-public class DataPlatformInstanceType implements com.linkedin.datahub.graphql.types.EntityType<DataPlatformInstance, String> {
+import static com.linkedin.metadata.Constants.DATA_PLATFORM_INSTANCE_ENTITY_NAME;
+
+public class DataPlatformInstanceType implements SearchableEntityType<DataPlatformInstance, String>,
+        com.linkedin.datahub.graphql.types.EntityType<DataPlatformInstance, String> {
 
     static final Set<String> ASPECTS_TO_FETCH = ImmutableSet.of(
         Constants.DATA_PLATFORM_INSTANCE_KEY_ASPECT_NAME,
@@ -84,4 +96,24 @@ public List<DataFetcherResult<DataPlatformInstance>> batchLoad(@Nonnull List<Str
         }
     }
 
+    @Override
+    public SearchResults search(@Nonnull String query,
+                                @Nullable List<FacetFilterInput> filters,
+                                int start,
+                                int count,
+                                @Nonnull final QueryContext context) throws Exception {
+        throw new NotImplementedException("Searchable type (deprecated) not implemented on DataPlatformInstance entity type");
+    }
+
+    @Override
+    public AutoCompleteResults autoComplete(@Nonnull String query,
+                                            @Nullable String field,
+                                            @Nullable Filter filters,
+                                            int limit,
+                                            @Nonnull final QueryContext context) throws Exception {
+        final AutoCompleteResult result = _entityClient.autoComplete(DATA_PLATFORM_INSTANCE_ENTITY_NAME, query,
+                filters, limit, context.getAuthentication());
+        return AutoCompleteResultsMapper.map(result);
+    }
+
 }
diff --git a/datahub-graphql-core/src/main/resources/entity.graphql b/datahub-graphql-core/src/main/resources/entity.graphql
index 044c405942a3c..c5dc652b9fcbe 100644
--- a/datahub-graphql-core/src/main/resources/entity.graphql
+++ b/datahub-graphql-core/src/main/resources/entity.graphql
@@ -226,6 +226,11 @@ type Query {
     listOwnershipTypes(
         "Input required for listing custom ownership types"
         input: ListOwnershipTypesInput!): ListOwnershipTypesResult!
+
+    """
+    Fetch a Data Platform Instance by primary key (urn)
+    """
+    dataPlatformInstance(urn: String!): DataPlatformInstance
 }
 
 """