feat(graphql): Add sorting option to single-entity search GraphQL API

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/businessattribute/ListBusinessAttributesResolver.java

@@ -61,6 +61,7 @@ public CompletableFuture<ListBusinessAttributesResult> get(
                     Constants.BUSINESS_ATTRIBUTE_ENTITY_NAME,
                     query,
                     Collections.emptyMap(),
+                    Collections.emptyList(),
                     start,
                     count);
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java

@@ -64,7 +64,7 @@ public CompletableFuture<ListGroupsResult> get(final DataFetchingEnvironment env
                           .withSearchFlags(flags -> flags.setFulltext(true)),
                       CORP_GROUP_ENTITY_NAME,
                       query,
-                      null,
+                      (Map<String, String>) null,
                       Collections.singletonList(
                           new SortCriterion()
                               .setField(CORP_GROUP_CREATED_TIME_INDEX_FIELD_NAME)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java

@@ -73,7 +73,7 @@ public CompletableFuture<ListSecretsResult> get(final DataFetchingEnvironment en
                           .withSearchFlags(flags -> flags.setFulltext(true)),
                       Constants.SECRETS_ENTITY_NAME,
                       query,
-                      null,
+                      (Map<String, String>) null,
                       Collections.singletonList(
                           new SortCriterion()
                               .setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/post/ListPostsResolver.java

@@ -63,7 +63,7 @@ public CompletableFuture<ListPostsResult> get(final DataFetchingEnvironment envi
                     context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
                     POST_ENTITY_NAME,
                     query,
-                    null,
+                    (Map<String, String>) null,
                     sortCriteria,
                     start,
                     count);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/role/ListRolesResolver.java

@@ -58,6 +58,7 @@ public CompletableFuture<ListRolesResult> get(final DataFetchingEnvironment envi
                     DATAHUB_ROLE_ENTITY_NAME,
                     query,
                     Collections.emptyMap(),
+                    Collections.emptyList(),
                     start,
                     count);
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java

@@ -2,7 +2,7 @@
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
 import static com.linkedin.metadata.Constants.*;
-import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags;
+import static com.linkedin.metadata.search.utils.SearchUtils.*;
 
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
@@ -17,11 +17,14 @@
 import com.linkedin.metadata.query.GroupingCriterionArray;
 import com.linkedin.metadata.query.GroupingSpec;
 import com.linkedin.metadata.query.SearchFlags;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.opentelemetry.extension.annotations.WithSpan;
 import java.util.Collections;
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
+import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
@@ -66,16 +69,34 @@ public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment)
     } else {
       searchFlags = applyDefaultSearchFlags(null, sanitizedQuery, SEARCH_RESOLVER_DEFAULTS);
     }
+    List<SortCriterion> sortCriteria;
+    if (input.getSortInput() != null) {
+      if (input.getSortInput().getSortCriteria() != null) {
+        sortCriteria =
+            input.getSortInput().getSortCriteria().stream()
+                .map(SearchUtils::mapSortCriterion)
+                .collect(Collectors.toList());
+      } else {
+        sortCriteria =
+            input.getSortInput().getSortCriterion() != null
+                ? Collections.singletonList(
+                    SearchUtils.mapSortCriterion(input.getSortInput().getSortCriterion()))
+                : Collections.emptyList();
+      }
+    } else {
+      sortCriteria = Collections.emptyList();
+    }
 
     return GraphQLConcurrencyUtils.supplyAsync(
         () -> {
           try {
             log.debug(
-                "Executing search. entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}, searchFlags: {}",
+                "Executing search. entity type {}, query {}, filters: {}, orFilters: {}, sortCriteria: {}, start: {}, count: {} searchFlags: {}",
                 input.getType(),
                 input.getQuery(),
                 input.getFilters(),
                 input.getOrFilters(),
+                input.getSortInput(),
                 start,
                 count,
                 searchFlags);
@@ -87,27 +108,29 @@ public CompletableFuture<SearchResults> get(DataFetchingEnvironment environment)
                     entityName,
                     sanitizedQuery,
                     ResolverUtils.buildFilter(input.getFilters(), input.getOrFilters()),
-                    Collections.emptyList(),
+                    sortCriteria,
                     start,
                     count));
           } catch (Exception e) {
             log.error(
-                "Failed to execute search: entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}, searchFlags: {}",
+                "Failed to execute search: entity type {}, query {}, filters: {}, orFilters: {}, sortCriteria: {}, start: {}, count: {}, searchFlags: {}",
                 input.getType(),
                 input.getQuery(),
                 input.getFilters(),
                 input.getOrFilters(),
+                input.getSortInput(),
                 start,
                 count,
                 searchFlags);
             throw new RuntimeException(
                 "Failed to execute search: "
                     + String.format(
-                        "entity type %s, query %s, filters: %s, orFilters: %s, start: %s, count: %s, searchFlags: %s",
+                        "entity type %s, query %s, filters: %s, orFilters: %s, sortCriteria: %s, start: %s, count: %s, searchFlags: %s",
                         input.getType(),
                         input.getQuery(),
                         input.getFilters(),
                         input.getOrFilters(),
+                        input.getSortInput(),
                         start,
                         count,
                         searchFlags),

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/test/ListTestsResolver.java

@@ -63,6 +63,7 @@ public CompletableFuture<ListTestsResult> get(final DataFetchingEnvironment envi
                       Constants.TEST_ENTITY_NAME,
                       query,
                       Collections.emptyMap(),
+                      Collections.emptyList(),
                       start,
                       count);
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/user/ListUsersResolver.java

@@ -64,6 +64,7 @@ public CompletableFuture<ListUsersResult> get(final DataFetchingEnvironment envi
                       CORP_USER_ENTITY_NAME,
                       query,
                       Collections.emptyMap(),
+                      Collections.emptyList(),
                       start,
                       count);
 

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/SearchableEntityType.java

@@ -6,6 +6,7 @@
 import com.linkedin.datahub.graphql.generated.FacetFilterInput;
 import com.linkedin.datahub.graphql.generated.SearchResults;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import java.util.List;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -40,6 +41,7 @@ public interface SearchableEntityType<T extends Entity, K> extends EntityType<T,
   SearchResults search(
       @Nonnull String query,
       @Nullable List<FacetFilterInput> filters,
+      @Nullable SortCriterion sort,
       int start,
       int count,
       @Nonnull final QueryContext context)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/businessattribute/BusinessAttributeType.java

@@ -26,12 +26,15 @@
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchResult;
 import graphql.execution.DataFetcherResult;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -103,17 +106,28 @@ public List<DataFetcherResult<BusinessAttribute>> batchLoad(
   public SearchResults search(
       @Nonnull String query,
       @Nullable List<FacetFilterInput> filters,
+      @Nullable SortCriterion sort,
       int start,
       int count,
       @Nonnull QueryContext context)
       throws Exception {
     final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
+
+    final SortCriterion sortCriterion = new SortCriterion();
+    Optional.ofNullable(sort)
+        .ifPresent(
+            s -> {
+              sortCriterion.setField(s.getField());
+              sortCriterion.setOrder(s.getOrder());
+            });
+
     final SearchResult searchResult =
         _entityClient.search(
             context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
             "businessAttribute",
             query,
             facetFilters,
+            Collections.singletonList(sortCriterion),
             start,
             count);
     return UrnSearchResultsMapper.map(context, searchResult);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/chart/ChartType.java

@@ -40,16 +40,19 @@
 import com.linkedin.metadata.browse.BrowseResult;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.mxe.MetadataChangeProposal;
 import com.linkedin.r2.RemoteInvocationException;
 import graphql.execution.DataFetcherResult;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -147,17 +150,28 @@ public List<DataFetcherResult<Chart>> batchLoad(
   public SearchResults search(
       @Nonnull String query,
       @Nullable List<FacetFilterInput> filters,
+      @Nullable SortCriterion sort,
       int start,
       int count,
       @Nonnull QueryContext context)
       throws Exception {
     final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
+
+    final SortCriterion sortCriterion = new SortCriterion();
+    Optional.ofNullable(sort)
+        .ifPresent(
+            s -> {
+              sortCriterion.setField(s.getField());
+              sortCriterion.setOrder(s.getOrder());
+            });
+
     final SearchResult searchResult =
         _entityClient.search(
             context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
             "chart",
             query,
             facetFilters,
+            Collections.singletonList(sortCriterion),
             start,
             count);
     return UrnSearchResultsMapper.map(context, searchResult);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/container/ContainerType.java

@@ -19,13 +19,16 @@
 import com.linkedin.metadata.Constants;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchResult;
 import graphql.execution.DataFetcherResult;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -120,17 +123,28 @@ private Urn getUrn(final String urnStr) {
   public SearchResults search(
       @Nonnull String query,
       @Nullable List<FacetFilterInput> filters,
+      @Nullable SortCriterion sort,
       int start,
       int count,
       @Nonnull final QueryContext context)
       throws Exception {
     final Map<String, String> facetFilters = ResolverUtils.buildFacetFilters(filters, FACET_FIELDS);
+
+    final SortCriterion sortCriterion = new SortCriterion();
+    Optional.ofNullable(sort)
+        .ifPresent(
+            s -> {
+              sortCriterion.setField(s.getField());
+              sortCriterion.setOrder(s.getOrder());
+            });
+
     final SearchResult searchResult =
         _entityClient.search(
             context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
             ENTITY_NAME,
             query,
             facetFilters,
+            Collections.singletonList(sortCriterion),
             start,
             count);
     return UrnSearchResultsMapper.map(context, searchResult);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpgroup/CorpGroupType.java

@@ -32,6 +32,7 @@
 import com.linkedin.metadata.authorization.PoliciesConfig;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.execution.DataFetcherResult;
@@ -40,6 +41,7 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import javax.annotation.Nonnull;
@@ -110,16 +112,25 @@ public List<DataFetcherResult<CorpGroup>> batchLoad(
   public SearchResults search(
       @Nonnull String query,
       @Nullable List<FacetFilterInput> filters,
+      @Nullable SortCriterion sort,
       int start,
       int count,
       @Nonnull final QueryContext context)
       throws Exception {
+    final SortCriterion sortCriterion = new SortCriterion();
+    Optional.ofNullable(sort)
+        .ifPresent(
+            s -> {
+              sortCriterion.setField(s.getField());
+              sortCriterion.setOrder(s.getOrder());
+            });
     final SearchResult searchResult =
         _entityClient.search(
             context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
             "corpGroup",
             query,
             Collections.emptyMap(),
+            Collections.singletonList(sortCriterion),
             start,
             count);
     return UrnSearchResultsMapper.map(context, searchResult);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java

@@ -37,6 +37,7 @@
 import com.linkedin.metadata.authorization.PoliciesConfig;
 import com.linkedin.metadata.query.AutoCompleteResult;
 import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.mxe.MetadataChangeProposal;
 import graphql.execution.DataFetcherResult;
@@ -113,16 +114,26 @@ public List<DataFetcherResult<CorpUser>> batchLoad(
   public SearchResults search(
       @Nonnull String query,
       @Nullable List<FacetFilterInput> filters,
+      @Nullable SortCriterion sort,
       int start,
       int count,
       @Nonnull final QueryContext context)
       throws Exception {
+
+    final SortCriterion sortCriterion = new SortCriterion();
+    Optional.ofNullable(sort)
+        .ifPresent(
+            s -> {
+              sortCriterion.setField(s.getField());
+              sortCriterion.setOrder(s.getOrder());
+            });
     final SearchResult searchResult =
         _entityClient.search(
             context.getOperationContext().withSearchFlags(flags -> flags.setFulltext(true)),
             "corpuser",
             query,
             Collections.emptyMap(),
+            Collections.singletonList(sortCriterion),
             start,
             count);
     return UrnSearchResultsMapper.map(context, searchResult);