feat(GraphQL support): PoC for GraphQL support in datahub-frontend (Queries + Mutations)

build.gradle

@@ -30,6 +30,8 @@ project.ext.spec = [
 ]
 
 project.ext.externalDependency = [
+    'antlr4Runtime': 'org.antlr:antlr4-runtime:4.7.2',
+    'antlr4': 'org.antlr:antlr4:4.7.2',
     'assertJ': 'org.assertj:assertj-core:3.11.1',
     'avro_1_7': 'org.apache.avro:avro:1.7.7',
     'avroCompiler_1_7': 'org.apache.avro:avro-compiler:1.7.7',
@@ -53,6 +55,7 @@ project.ext.externalDependency = [
     'gmaTestModels': "com.linkedin.datahub-gma:test-models:$gmaVersion",
     'gmaTestModelsDataTemplate': "com.linkedin.datahub-gma:test-models-data-template:$gmaVersion",
     'gmaValidators': "com.linkedin.datahub-gma:validators:$gmaVersion",
+    'graphqlJava': 'com.graphql-java:graphql-java:13.0',
     'gson': 'com.google.code.gson:gson:2.8.6',
     'guice': 'com.google.inject:guice:4.2.2',
     'guava': 'com.google.guava:guava:27.0.1-jre',

datahub-dao/build.gradle

@@ -3,6 +3,8 @@ apply plugin: 'java'
 dependencies {
   compile project(':gms:client')
   compile project(':metadata-utils')
+  compile project(':li-utils')
+
   compile externalDependency.elasticSearchRest
   compile externalDependency.hibernateCore
   compile externalDependency.guava

datahub-dao/src/main/java/com/linkedin/datahub/dao/DaoFactory.java

@@ -1,10 +1,11 @@
 package com.linkedin.datahub.dao;
 
-import com.linkedin.datahub.dao.table.DataPlatformsDao;
-import com.linkedin.datahub.dao.table.DatasetOwnerDao;
+import com.linkedin.datahub.dao.table.CorpUsersDao;
 import com.linkedin.datahub.dao.table.DatasetsDao;
-import com.linkedin.datahub.dao.table.GmsDao;
+import com.linkedin.datahub.dao.table.DatasetOwnerDao;
+import com.linkedin.datahub.dao.table.DataPlatformsDao;
 import com.linkedin.datahub.dao.table.LineageDao;
+import com.linkedin.datahub.dao.table.GmsDao;
 import com.linkedin.datahub.dao.view.BrowseDAO;
 import com.linkedin.datahub.dao.view.CorpUserViewDao;
 import com.linkedin.datahub.dao.view.DatasetViewDao;
@@ -21,6 +22,7 @@ public class DaoFactory {
   private static DocumentSearchDao datasetDocumentSearchDao;
   private static DocumentSearchDao corpUserDocumentSearchDao;
   private static CorpUserViewDao corpUserViewDao;
+  private static CorpUsersDao corpUsersDao;
   private static BrowseDAO datasetBrowseDao;
   private static OwnerViewDao ownerViewDao;
   private static DatasetViewDao datasetViewDao;
@@ -54,6 +56,13 @@ public static DocumentSearchDao getCorpUserDocumentSearchDao() {
     return corpUserDocumentSearchDao;
   }
 
+  public static CorpUsersDao getCorpUsersDao() {
+    if (corpUsersDao == null) {
+      corpUsersDao = new CorpUsersDao(getGmsDao().get_corpUsers());
+    }
+    return corpUsersDao;
+  }
+
   public static BrowseDAO getDatasetBrowseDAO() {
     if (datasetBrowseDao == null) {
       datasetBrowseDao = new BrowseDAO<>(getGmsDao().get_datasets());
@@ -92,7 +101,7 @@ public static DatasetOwnerDao getDatasetOwnerDao() {
 
   public static DatasetsDao getDatasetsDao() {
     if (datasetsDao == null) {
-      datasetsDao = new DatasetsDao(getGmsDao().get_ownerships());
+      datasetsDao = new DatasetsDao(getGmsDao().get_datasets(), getGmsDao().get_ownerships());
     }
     return datasetsDao;
   }

datahub-dao/src/main/java/com/linkedin/datahub/dao/table/CorpUsersDao.java

@@ -0,0 +1,65 @@
+package com.linkedin.datahub.dao.table;
+
+import com.linkedin.common.urn.CorpuserUrn;
+import com.linkedin.datahub.util.CorpUserUtil;
+import com.linkedin.identity.CorpUser;
+import com.linkedin.identity.client.CorpUsers;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+
+public class CorpUsersDao {
+
+    private final CorpUsers _corpUsers;
+
+    public CorpUsersDao(@Nonnull CorpUsers corpUsers) {
+        _corpUsers = corpUsers;
+    }
+
+    @Nonnull
+    public List<CorpUser> getCorpUsers(@Nonnull List<String> corpUserUrnStrs) throws Exception {
+        List<CorpuserUrn> corpUserUrns = corpUserUrnStrs
+                .stream()
+                .map(this::getOwnerUrn)
+                .collect(Collectors.toList());
+
+        Map<CorpuserUrn, CorpUser> corpUserMap = _corpUsers.batchGet(new HashSet<>(corpUserUrns));
+
+        List<CorpUser> results = new ArrayList<>();
+        for (CorpuserUrn urn : corpUserUrns) {
+            results.add(corpUserMap.getOrDefault(urn, null));
+        }
+        return results;
+    }
+
+    @Nonnull
+    public CorpUser getCorpUser(@Nonnull String corpUserUrn) throws Exception {
+        CorpuserUrn corpUser = getOwnerUrn(corpUserUrn);
+        return corpUser == null ? null : _corpUsers.get(corpUser);
+    }
+
+
+    @Nonnull
+    private Set<CorpuserUrn> getOwnerUrns(@Nonnull Set<String> corpUserUrns) {
+        return corpUserUrns.stream()
+                .map(this::getOwnerUrn)
+                .collect(Collectors.toSet());
+    }
+
+    @Nullable
+    private CorpuserUrn getOwnerUrn(@Nonnull String corpUserUrn) {
+        try {
+            return CorpUserUtil.toCorpUserUrn(corpUserUrn);
+        } catch (URISyntaxException e) {
+            return null;
+        }
+    }
+}

datahub-dao/src/main/java/com/linkedin/datahub/dao/table/DatasetsDao.java

@@ -1,21 +1,76 @@
 package com.linkedin.datahub.dao.table;
 
+import com.linkedin.common.Ownership;
+import com.linkedin.common.urn.DatasetUrn;
+import com.linkedin.dataset.Dataset;
+import com.linkedin.dataset.client.Datasets;
 import com.linkedin.dataset.client.Ownerships;
 
 import javax.annotation.Nonnull;
+import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.Map;
 import java.util.Arrays;
 import java.util.List;
+import java.util.Objects;
+import java.util.stream.Collectors;
 
+import static com.linkedin.datahub.util.DatasetUtil.toDatasetUrn;
+import static com.linkedin.datahub.util.DatasetUtil.toSnapshot;
 
+/**
+ * Data access object for Datasets and related aspects.
+ */
 public class DatasetsDao {
 
+    private final Datasets _datasets;
     private final Ownerships _ownership;
 
-    public DatasetsDao(@Nonnull Ownerships ownerships) {
-        this._ownership = ownerships;
+    public DatasetsDao(@Nonnull Datasets datasets, @Nonnull Ownerships ownership) {
+        _datasets = datasets;
+        _ownership = ownership;
     }
 
     public List<String> getDatasetOwnerTypes() {
         return Arrays.asList("DataOwner", "Producer", "Delegate", "Stakeholder", "Consumer", "Developer");
     }
+
+    @Nonnull
+    public Dataset getDataset(@Nonnull String datasetUrn) throws Exception {
+        return _datasets.get(toDatasetUrn(datasetUrn));
+    }
+
+
+    @Nonnull
+    public List<Dataset> getDatasets(@Nonnull List<String> datasetUrnStrs) throws Exception {
+        List<DatasetUrn> datasetUrns = datasetUrnStrs.stream().map(urnStr -> {
+                try {
+                    return toDatasetUrn(urnStr);
+                } catch (URISyntaxException e) {
+                    return null;
+                }
+            }).collect(Collectors.toList());
+
+        Map<DatasetUrn, Dataset> datasetMap = _datasets.batchGet(datasetUrns
+                .stream()
+                .filter(Objects::nonNull)
+                .collect(Collectors.toSet()));
+
+        List<Dataset> results = new ArrayList<>();
+        for (DatasetUrn urn : datasetUrns) {
+            results.add(datasetMap.getOrDefault(urn, null));
+        }
+        return results;
+    }
+
+    @Nonnull
+    public Ownership getOwnership(@Nonnull String datasetUrn) throws Exception {
+        return _ownership.getLatestOwnership(toDatasetUrn(datasetUrn));
+    }
+
+    @Nonnull
+    public void updateDataset(@Nonnull String datasetUrnStr, @Nonnull Dataset partialDataset) throws Exception {
+        DatasetUrn datasetUrn = toDatasetUrn(datasetUrnStr);
+        _datasets.createSnapshot(datasetUrn, toSnapshot(datasetUrn, partialDataset));
+    }
 }

datahub-dao/src/main/java/com/linkedin/datahub/dao/view/BrowseDAO.java

@@ -1,20 +1,21 @@
 package com.linkedin.datahub.dao.view;
 
 import com.fasterxml.jackson.databind.JsonNode;
-import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.linkedin.data.template.StringArray;
 import com.linkedin.datahub.util.DatasetUtil;
 import com.linkedin.dataset.client.Datasets;
 import com.linkedin.metadata.query.BrowseResult;
-import com.linkedin.metadata.query.BrowseResultEntityArray;
 import com.linkedin.r2.RemoteInvocationException;
 import java.io.IOException;
 import java.net.URISyntaxException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import javax.annotation.Nonnull;
 
+import static com.linkedin.datahub.util.BrowseUtil.getJsonFromBrowseResult;
+import static com.linkedin.datahub.util.BrowseUtil.toGraphQLBrowseResponse;
 import static com.linkedin.datahub.util.RestliUtil.*;
-import static java.util.stream.Collectors.toList;
 
 
 public class BrowseDAO<CLIENT> {
@@ -27,35 +28,47 @@ public BrowseDAO(@Nonnull CLIENT client) {
   @Nonnull
   public JsonNode browse(@Nonnull String path, @Nonnull Map<String, String> requestFilters,
       int start, int count) throws RemoteInvocationException, IOException {
-    final BrowseResult resp;
-    if (_client instanceof Datasets) {
-      resp = ((Datasets) _client).browse(path, requestFilters, start, count);
-    } else {
-      throw new IllegalArgumentException("Unexpected client type: " + _client.getClass().getName());
-    }
-    return getJsonFromBrowseResult(resp);
+    return getJsonFromBrowseResult(executeBrowse(path, requestFilters, start, count));
+  }
+
+  @Nonnull
+  public Map<String, Object> graphQLBrowse(@Nonnull String path,
+                                           @Nonnull Map<String, String> requestFilters,
+                                           int start,
+                                           int count) throws RemoteInvocationException {
+    return toGraphQLBrowseResponse(executeBrowse(path, requestFilters, start, count));
   }
 
   @Nonnull
   public JsonNode getBrowsePaths(@Nonnull String urn) throws RemoteInvocationException, URISyntaxException {
-    final StringArray response;
+    return stringCollectionToArrayNode(new ArrayList<>(executeBrowsePaths(urn)));
+  }
+
+  @Nonnull
+  public List<String> graphQLGetBrowsePaths(@Nonnull String urn) throws RemoteInvocationException, URISyntaxException {
+    return executeBrowsePaths(urn);
+  }
+
+  private BrowseResult executeBrowse(@Nonnull String path,
+                                     @Nonnull Map<String, String> requestFilters,
+                                     int start,
+                                     int count) throws RemoteInvocationException {
+    final BrowseResult res;
     if (_client instanceof Datasets) {
-      response = ((Datasets) _client).getBrowsePaths(DatasetUtil.toDatasetUrn(urn));
+      res = ((Datasets) _client).browse(path, requestFilters, start, count);
     } else {
       throw new IllegalArgumentException("Unexpected client type: " + _client.getClass().getName());
     }
-    return stringCollectionToArrayNode(response.stream().collect(toList()));
+    return res;
   }
 
-  @Nonnull
-  public JsonNode getJsonFromBrowseResult(@Nonnull BrowseResult browseResult) throws IOException {
-    final ObjectNode node = OM.createObjectNode();
-    final BrowseResultEntityArray browseResultEntityArray = browseResult.getEntities();
-    node.set("elements", collectionToArrayNode(browseResultEntityArray.subList(0, browseResultEntityArray.size())));
-    node.put("start", browseResult.getFrom());
-    node.put("count", browseResult.getPageSize());
-    node.put("total", browseResult.getNumEntities());
-    node.set("metadata", toJsonNode(browseResult.getMetadata()));
-    return node;
+  private StringArray executeBrowsePaths(@Nonnull String urn) throws RemoteInvocationException, URISyntaxException {
+    final StringArray res;
+    if (_client instanceof Datasets) {
+      res = ((Datasets) _client).getBrowsePaths(DatasetUtil.toDatasetUrn(urn));
+    } else {
+      throw new IllegalArgumentException("Unexpected client type: " + _client.getClass().getName());
+    }
+    return res;
   }
 }