From 06fae51ca38b5c294b3f16ed1a0c6bba24b41c47 Mon Sep 17 00:00:00 2001
From: Kerem Sahin <ksahin@linkedin.com>
Date: Thu, 26 Sep 2019 20:53:38 -0700
Subject: [PATCH] Fix NPE issue with downstream lineage resource

---
 .../resources/DownstreamLineageResource.java     | 16 +++++++++++-----
 .../main/java/com/linkedin/common/urn/Urn.java   | 10 ++++++++++
 .../com/linkedin/metadata/dao/utils/ESUtils.java |  6 +++++-
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/gms/impl/src/main/java/com/linkedin/dataset/rest/resources/DownstreamLineageResource.java b/gms/impl/src/main/java/com/linkedin/dataset/rest/resources/DownstreamLineageResource.java
index 782d107c926..4e08da174ff 100644
--- a/gms/impl/src/main/java/com/linkedin/dataset/rest/resources/DownstreamLineageResource.java
+++ b/gms/impl/src/main/java/com/linkedin/dataset/rest/resources/DownstreamLineageResource.java
@@ -21,6 +21,7 @@ import javax.annotation.Nonnull;
 import javax.inject.Inject;
 import javax.inject.Named;
 import java.util.Collections;
+import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
 
@@ -52,7 +53,7 @@ public final class DownstreamLineageResource extends SimpleResourceTemplate<Down
     final Filter filter = SearchUtils.getFilter(Collections.singletonMap("upstreams", datasetUrn.toString()));
 
     return RestliUtils.toTask(() -> {
-      final SearchResult<DatasetDocument> searchResult = _searchDAO.search("*", filter, 0, Integer.MAX_VALUE);
+      final SearchResult<DatasetDocument> searchResult = _searchDAO.search("*", filter, 0, 10000);
       final Set<DatasetUrn> downstreamDatasets = searchResult.getDocumentList()
               .stream()
               .map(d -> (DatasetUrn) ModelUtils.getUrnFromDocument(d))
@@ -60,11 +61,16 @@ public final class DownstreamLineageResource extends SimpleResourceTemplate<Down
       final DownstreamArray downstreamArray = new DownstreamArray(downstreamDatasets.stream()
               .map(ds -> {
                 final UpstreamLineage upstreamLineage = (UpstreamLineage) _localDAO.get(UpstreamLineage.class, ds).get();
-                final Upstream upstream = upstreamLineage.getUpstreams().stream()
+                final List<Upstream> upstreams = upstreamLineage.getUpstreams().stream()
                         .filter(us -> us.getDataset().equals(datasetUrn))
-                        .collect(Collectors.toList())
-                        .get(0);
-                return new Downstream().setDataset(ds).setType(upstream.getType()).setAuditStamp(upstream.getAuditStamp());
+                        .collect(Collectors.toList());
+                if (upstreams.size() != 1) {
+                  throw new RuntimeException(String.format("There is no relation or more than 1 relation between the datasets!"));
+                }
+                return new Downstream()
+                        .setDataset(ds)
+                        .setType(upstreams.get(0).getType())
+                        .setAuditStamp(upstreams.get(0).getAuditStamp());
               })
               .collect(Collectors.toList())
       );
diff --git a/li-utils/src/main/java/com/linkedin/common/urn/Urn.java b/li-utils/src/main/java/com/linkedin/common/urn/Urn.java
index 29245926706..36ebc0fa84b 100644
--- a/li-utils/src/main/java/com/linkedin/common/urn/Urn.java
+++ b/li-utils/src/main/java/com/linkedin/common/urn/Urn.java
@@ -1,5 +1,6 @@
 package com.linkedin.common.urn;
 
+import javax.annotation.Nonnull;
 import java.net.URISyntaxException;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -68,4 +69,13 @@ public class Urn {
   public String toString() {
     return _urn;
   }
+
+  public static boolean isUrn(@Nonnull String urn) {
+    try {
+      final Urn dummy = Urn.createFromString(urn);
+      return true;
+    } catch(URISyntaxException e) {
+      return false;
+    }
+  }
 }
\ No newline at end of file
diff --git a/metadata-dao-impl/elasticsearch-dao/src/main/java/com/linkedin/metadata/dao/utils/ESUtils.java b/metadata-dao-impl/elasticsearch-dao/src/main/java/com/linkedin/metadata/dao/utils/ESUtils.java
index 882a2eec936..7f82c9e46ee 100644
--- a/metadata-dao-impl/elasticsearch-dao/src/main/java/com/linkedin/metadata/dao/utils/ESUtils.java
+++ b/metadata-dao-impl/elasticsearch-dao/src/main/java/com/linkedin/metadata/dao/utils/ESUtils.java
@@ -1,8 +1,11 @@
 package com.linkedin.metadata.dao.utils;
 
+import java.net.URISyntaxException;
 import java.util.Arrays;
 import java.util.Map;
 import javax.annotation.Nonnull;
+
+import com.linkedin.common.urn.Urn;
 import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilders;
 
@@ -26,7 +29,8 @@ public class ESUtils {
     BoolQueryBuilder boolFilter = new BoolQueryBuilder();
     for (Map.Entry<String, String> entry : requestMap.entrySet()) {
       BoolQueryBuilder filters = new BoolQueryBuilder();
-      Arrays.stream(entry.getValue().split(","))
+      // TODO: Remove checking for urn after solving META-10102
+      Arrays.stream(Urn.isUrn(entry.getValue()) ? new String[]{entry.getValue()} : entry.getValue().split(","))
           .forEach(elem -> filters.should(QueryBuilders.matchQuery(entry.getKey(), elem)));
       boolFilter.must(filters);
     }