From c1eff5982e3bab93fae1ac1d03e6f68f0c25c87e Mon Sep 17 00:00:00 2001
From: RyanHolstien <RyanHolstien@users.noreply.github.com>
Date: Fri, 29 Mar 2024 13:43:33 -0500
Subject: [PATCH] fix(openapi): fix index out of bounds for sort order (#10168)

---
 .../v2/controller/RelationshipController.java | 20 +++-------------
 .../com/linkedin/metadata/graph/Edge.java     | 24 ++++++++++++++++---
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java
index 2ac16e8c27..b71d2b19fc 100644
--- a/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java
+++ b/metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/RelationshipController.java
@@ -16,10 +16,7 @@ import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService;
 import com.linkedin.metadata.models.registry.EntityRegistry;
 import com.linkedin.metadata.query.filter.RelationshipDirection;
 import com.linkedin.metadata.query.filter.RelationshipFilter;
-import com.linkedin.metadata.query.filter.SortCriterion;
-import com.linkedin.metadata.query.filter.SortOrder;
 import com.linkedin.metadata.search.utils.QueryUtils;
-import com.linkedin.metadata.utils.SearchUtil;
 import io.datahubproject.openapi.exception.UnauthorizedException;
 import io.datahubproject.openapi.v2.models.GenericRelationship;
 import io.datahubproject.openapi.v2.models.GenericScrollResult;
@@ -28,7 +25,6 @@ import io.swagger.v3.oas.annotations.tags.Tag;
 import java.util.Arrays;
 import java.util.List;
 import java.util.stream.Collectors;
-import java.util.stream.IntStream;
 import java.util.stream.Stream;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -49,16 +45,6 @@ import org.springframework.web.bind.annotation.RestController;
     name = "Generic Relationships",
     description = "APIs for ingesting and accessing entity relationships.")
 public class RelationshipController {
-  private static final String[] SORT_ORDERS = {"ASCENDING", "ASCENDING", "ASCENDING", "ASCENDING"};
-  private static final List<SortCriterion> EDGE_SORT_CRITERION;
-
-  static {
-    EDGE_SORT_CRITERION =
-        IntStream.range(0, Edge.KEY_FIELDS.length)
-            .mapToObj(
-                idx -> SearchUtil.sortBy(Edge.KEY_FIELDS[idx], SortOrder.valueOf(SORT_ORDERS[idx])))
-            .collect(Collectors.toList());
-  }
 
   @Autowired private EntityRegistry entityRegistry;
   @Autowired private ElasticSearchGraphService graphService;
@@ -97,7 +83,7 @@ public class RelationshipController {
             null,
             List.of(relationshipType),
             new RelationshipFilter().setDirection(RelationshipDirection.UNDIRECTED),
-            EDGE_SORT_CRITERION,
+            Edge.EDGE_SORT_CRITERION,
             scrollId,
             count,
             null,
@@ -180,7 +166,7 @@ public class RelationshipController {
               new RelationshipFilter()
                   .setDirection(RelationshipDirection.UNDIRECTED)
                   .setOr(QueryUtils.newFilter("destination.urn", entityUrn).getOr()),
-              EDGE_SORT_CRITERION,
+              Edge.EDGE_SORT_CRITERION,
               scrollId,
               count,
               null,
@@ -197,7 +183,7 @@ public class RelationshipController {
               new RelationshipFilter()
                   .setDirection(RelationshipDirection.UNDIRECTED)
                   .setOr(QueryUtils.newFilter("source.urn", entityUrn).getOr()),
-              EDGE_SORT_CRITERION,
+              Edge.EDGE_SORT_CRITERION,
               scrollId,
               count,
               null,
diff --git a/metadata-service/services/src/main/java/com/linkedin/metadata/graph/Edge.java b/metadata-service/services/src/main/java/com/linkedin/metadata/graph/Edge.java
index cb74ae5acd..83965986fc 100644
--- a/metadata-service/services/src/main/java/com/linkedin/metadata/graph/Edge.java
+++ b/metadata-service/services/src/main/java/com/linkedin/metadata/graph/Edge.java
@@ -1,11 +1,17 @@
 package com.linkedin.metadata.graph;
 
 import com.linkedin.common.urn.Urn;
+import com.linkedin.metadata.query.filter.SortCriterion;
+import com.linkedin.metadata.query.filter.SortOrder;
+import com.linkedin.metadata.utils.SearchUtil;
+import com.linkedin.util.Pair;
 import java.nio.charset.StandardCharsets;
 import java.security.MessageDigest;
 import java.security.NoSuchAlgorithmException;
 import java.util.Base64;
+import java.util.List;
 import java.util.Map;
+import java.util.stream.Collectors;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
@@ -75,8 +81,20 @@ public class Edge {
     }
   }
 
-  public static final String[] KEY_FIELDS = {
-    "source.urn", "destination.urn", "relationshipType", "lifeCycleOwner"
-  };
+  public static final String SOURCE_URN_FIELD = "source.urn";
+  public static final String DESTINATION_URN_FIELD = "destination.urn";
+  public static final String RELATIONSHIP_TYPE_FIELD = "relationshipType";
+  public static final String LIFE_CYCLE_OWNER_FIELD = "lifeCycleOwner";
+
+  public static final List<Pair<String, SortOrder>> KEY_SORTS =
+      List.of(
+          new Pair<>(SOURCE_URN_FIELD, SortOrder.ASCENDING),
+          new Pair<>(DESTINATION_URN_FIELD, SortOrder.ASCENDING),
+          new Pair<>(RELATIONSHIP_TYPE_FIELD, SortOrder.ASCENDING),
+          new Pair<>(LIFE_CYCLE_OWNER_FIELD, SortOrder.ASCENDING));
+  public static List<SortCriterion> EDGE_SORT_CRITERION =
+      KEY_SORTS.stream()
+          .map(entry -> SearchUtil.sortBy(entry.getKey(), entry.getValue()))
+          .collect(Collectors.toList());
   private static final String DOC_DELIMETER = "--";
 }