From a27f82cae27c87381a404b1b600793b855b3e766 Mon Sep 17 00:00:00 2001
From: david-leifker <114954101+david-leifker@users.noreply.github.com>
Date: Mon, 20 Mar 2023 18:07:25 -0500
Subject: [PATCH] =?UTF-8?q?fix(search):=20tags=20with=20colons=20exercises?=
 =?UTF-8?q?=20search=20with=20urns,=20must=20follow=E2=80=A6=20(#7602)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../resolvers/search/SearchResolver.java      |  33 +--
 .../mappers/SearchFlagsInputMapper.java       |   2 +
 .../resolvers/search/SearchResolverTest.java  | 190 ++++++++++++++++++
 .../query/request/SearchRequestHandler.java   |   2 +-
 .../metadata/search/utils/SearchUtils.java    |   3 +-
 .../search/utils/SearchUtilsTest.java         |  18 ++
 6 files changed, 234 insertions(+), 14 deletions(-)
 create mode 100644 datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java

diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java
index 2f0950be73..0e66d6e601 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java
@@ -16,6 +16,7 @@ import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.metadata.search.utils.SearchUtils.applyDefaultSearchFlags;
 
 
 /**
@@ -24,7 +25,12 @@ import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
 @Slf4j
 @RequiredArgsConstructor
 public class SearchResolver implements DataFetcher<CompletableFuture<SearchResults>> {
-
+  private static final SearchFlags SEARCH_RESOLVER_DEFAULTS = new SearchFlags()
+          .setFulltext(true)
+          .setMaxAggValues(20)
+          .setSkipCache(false)
+          .setSkipAggregates(false)
+          .setSkipHighlighting(false);
   private static final int DEFAULT_START = 0;
   private static final int DEFAULT_COUNT = 10;
 
@@ -40,26 +46,29 @@ public class SearchResolver implements DataFetcher<CompletableFuture<SearchResul
 
     final int start = input.getStart() != null ? input.getStart() : DEFAULT_START;
     final int count = input.getCount() != null ? input.getCount() : DEFAULT_COUNT;
+    final SearchFlags searchFlags;
+    com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags();
+    if (inputFlags != null) {
+      searchFlags = SearchFlagsInputMapper.INSTANCE.apply(inputFlags);
+    } else {
+      searchFlags = applyDefaultSearchFlags(null, sanitizedQuery, SEARCH_RESOLVER_DEFAULTS);
+    }
 
     return CompletableFuture.supplyAsync(() -> {
       try {
-        log.debug("Executing search. entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}", input.getType(),
-            input.getQuery(), input.getFilters(), input.getOrFilters(), start, count);
-        SearchFlags searchFlags = null;
-        com.linkedin.datahub.graphql.generated.SearchFlags inputFlags = input.getSearchFlags();
-        if (inputFlags != null) {
-          searchFlags = SearchFlagsInputMapper.INSTANCE.apply(inputFlags);
-        }
+        log.debug("Executing search. entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}, searchFlags: {}",
+                input.getType(), input.getQuery(), input.getFilters(), input.getOrFilters(), start, count, searchFlags);
+
         return UrnSearchResultsMapper.map(
             _entityClient.search(entityName, sanitizedQuery, ResolverUtils.buildFilter(input.getFilters(),
                     input.getOrFilters()), null, start, count, ResolverUtils.getAuthentication(environment),
                     searchFlags));
       } catch (Exception e) {
-        log.error("Failed to execute search: entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}",
-            input.getType(), input.getQuery(), input.getFilters(), input.getOrFilters(), start, count);
+        log.error("Failed to execute search: entity type {}, query {}, filters: {}, orFilters: {}, start: {}, count: {}, searchFlags: {}",
+            input.getType(), input.getQuery(), input.getFilters(), input.getOrFilters(), 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",
-                input.getType(), input.getQuery(), input.getFilters(), input.getOrFilters(), start, count), e);
+            "Failed to execute search: " + String.format("entity type %s, query %s, filters: %s, orFilters: %s, start: %s, count: %s, searchFlags: %s",
+                input.getType(), input.getQuery(), input.getFilters(), input.getOrFilters(), start, count, searchFlags), e);
       }
     });
   }
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java
index 5cf6e20bc1..6435d6ee4c 100644
--- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java
+++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/common/mappers/SearchFlagsInputMapper.java
@@ -24,6 +24,8 @@ public class SearchFlagsInputMapper implements ModelMapper<SearchFlags, com.link
     com.linkedin.metadata.query.SearchFlags result = new com.linkedin.metadata.query.SearchFlags();
     if (searchFlags.getFulltext() != null) {
       result.setFulltext(searchFlags.getFulltext());
+    } else {
+      result.setFulltext(true);
     }
     if (searchFlags.getSkipCache() != null) {
       result.setSkipCache(searchFlags.getSkipCache());
diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java
new file mode 100644
index 0000000000..6ba8b3cefe
--- /dev/null
+++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolverTest.java
@@ -0,0 +1,190 @@
+package com.linkedin.datahub.graphql.resolvers.search;
+
+import com.datahub.authentication.Authentication;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.generated.EntityType;
+import com.linkedin.datahub.graphql.generated.SearchFlags;
+import com.linkedin.datahub.graphql.generated.SearchInput;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.filter.Filter;
+import com.linkedin.metadata.query.filter.SortCriterion;
+import com.linkedin.metadata.search.SearchEntityArray;
+import com.linkedin.metadata.search.SearchResult;
+import com.linkedin.metadata.search.SearchResultMetadata;
+import graphql.schema.DataFetchingEnvironment;
+import org.mockito.Mockito;
+import org.testng.annotations.Test;
+
+import static com.linkedin.datahub.graphql.TestUtils.getMockAllowContext;
+
+
+public class SearchResolverTest {
+    @Test
+    public void testDefaultSearchFlags() throws Exception {
+        EntityClient mockClient = initMockSearchEntityClient();
+        final SearchResolver resolver = new SearchResolver(mockClient);
+
+        final SearchInput testInput = new SearchInput(
+                EntityType.DATASET,
+                "",
+                0,
+                10,
+                null,
+                null,
+                null
+        );
+        DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+        QueryContext mockContext = getMockAllowContext();
+        Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(testInput);
+        Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+
+        resolver.get(mockEnv).get();
+
+        verifyMockSearchEntityClient(
+                mockClient,
+                Constants.DATASET_ENTITY_NAME, // Verify that merged entity types were used.
+                "",
+                null,
+                null,
+                0,
+                10,
+                new com.linkedin.metadata.query.SearchFlags()
+                        .setFulltext(true)
+                        .setSkipAggregates(false)
+                        .setSkipHighlighting(true) // empty/wildcard
+                        .setMaxAggValues(20)
+                        .setSkipCache(false)
+        );
+    }
+
+    @Test
+    public void testOverrideSearchFlags() throws Exception {
+        EntityClient mockClient = initMockSearchEntityClient();
+        final SearchResolver resolver = new SearchResolver(mockClient);
+
+        final SearchFlags inputSearchFlags = new SearchFlags();
+        inputSearchFlags.setFulltext(false);
+        inputSearchFlags.setSkipAggregates(true);
+        inputSearchFlags.setSkipHighlighting(true);
+        inputSearchFlags.setMaxAggValues(10);
+        inputSearchFlags.setSkipCache(true);
+
+        final SearchInput testInput = new SearchInput(
+                EntityType.DATASET,
+                "",
+                1,
+                11,
+                null,
+                null,
+                inputSearchFlags
+        );
+        DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+        QueryContext mockContext = getMockAllowContext();
+        Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(testInput);
+        Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+
+        resolver.get(mockEnv).get();
+
+        verifyMockSearchEntityClient(
+                mockClient,
+                Constants.DATASET_ENTITY_NAME, // Verify that merged entity types were used.
+                "",
+                null,
+                null,
+                1,
+                11,
+                new com.linkedin.metadata.query.SearchFlags()
+                        .setFulltext(false)
+                        .setSkipAggregates(true)
+                        .setSkipHighlighting(true)
+                        .setMaxAggValues(10)
+                        .setSkipCache(true)
+        );
+    }
+
+    @Test
+    public void testNonWildCardSearchFlags() throws Exception {
+        EntityClient mockClient = initMockSearchEntityClient();
+        final SearchResolver resolver = new SearchResolver(mockClient);
+
+        final SearchInput testInput = new SearchInput(
+                EntityType.DATASET,
+                "not a wildcard",
+                0,
+                10,
+                null,
+                null,
+                null
+        );
+        DataFetchingEnvironment mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+        QueryContext mockContext = getMockAllowContext();
+        Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(testInput);
+        Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+
+        resolver.get(mockEnv).get();
+
+        verifyMockSearchEntityClient(
+                mockClient,
+                Constants.DATASET_ENTITY_NAME, // Verify that merged entity types were used.
+                "not a wildcard",
+                null, // Verify that view filter was used.
+                null,
+                0,
+                10,
+                new com.linkedin.metadata.query.SearchFlags()
+                        .setFulltext(true)
+                        .setSkipAggregates(false)
+                        .setSkipHighlighting(false) // empty/wildcard
+                        .setMaxAggValues(20)
+                        .setSkipCache(false)
+        );
+    }
+
+    private EntityClient initMockSearchEntityClient() throws Exception {
+        EntityClient client = Mockito.mock(EntityClient.class);
+        Mockito.when(client.search(
+                Mockito.anyString(),
+                Mockito.anyString(),
+                Mockito.any(),
+                Mockito.any(),
+                Mockito.anyInt(),
+                Mockito.anyInt(),
+                Mockito.any(Authentication.class),
+                Mockito.any()
+        )).thenReturn(
+                new SearchResult()
+                        .setEntities(new SearchEntityArray())
+                        .setNumEntities(0)
+                        .setFrom(0)
+                        .setPageSize(0)
+                        .setMetadata(new SearchResultMetadata())
+        );
+        return client;
+    }
+
+    private void verifyMockSearchEntityClient(
+            EntityClient mockClient,
+            String entityName,
+            String query,
+            Filter filter,
+            SortCriterion sortCriterion,
+            int start,
+            int limit,
+            com.linkedin.metadata.query.SearchFlags searchFlags
+    ) throws Exception {
+        Mockito.verify(mockClient, Mockito.times(1)).search(
+                Mockito.eq(entityName),
+                Mockito.eq(query),
+                Mockito.eq(filter),
+                Mockito.eq(sortCriterion),
+                Mockito.eq(start),
+                Mockito.eq(limit),
+                Mockito.any(Authentication.class),
+                Mockito.eq(searchFlags)
+        );
+    }
+
+    private SearchResolverTest() {
+    }
+}
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java
index b7574e1b43..7dd1d24f5d 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/elasticsearch/query/request/SearchRequestHandler.java
@@ -186,7 +186,7 @@ public class SearchRequestHandler {
   @WithSpan
   public SearchRequest getSearchRequest(@Nonnull String input, @Nullable Filter filter,
                                         @Nullable SortCriterion sortCriterion, int from, int size,
-                                        @Nonnull SearchFlags searchFlags) {
+                                        @Nullable SearchFlags searchFlags) {
     SearchFlags finalSearchFlags = applyDefaultSearchFlags(searchFlags, input, DEFAULT_SERVICE_SEARCH_FLAGS);
 
     SearchRequest searchRequest = new SearchRequest();
diff --git a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java
index c003504c8d..38bdef5bd3 100644
--- a/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java
+++ b/metadata-io/src/main/java/com/linkedin/metadata/search/utils/SearchUtils.java
@@ -170,9 +170,10 @@ public class SearchUtils {
     return listResult;
   }
 
+  @SneakyThrows
   public static SearchFlags applyDefaultSearchFlags(@Nullable SearchFlags inputFlags, @Nullable String query,
                                                     @Nonnull SearchFlags defaultFlags) {
-    SearchFlags finalSearchFlags = inputFlags != null ? inputFlags : defaultFlags;
+    SearchFlags finalSearchFlags = inputFlags != null ? inputFlags : defaultFlags.copy();
     if (!finalSearchFlags.hasFulltext() || finalSearchFlags.isFulltext() == null) {
       finalSearchFlags.setFulltext(defaultFlags.isFulltext());
     }
diff --git a/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java
index 2b080e68b8..7aa3bb19f0 100644
--- a/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java
+++ b/metadata-io/src/test/java/com/linkedin/metadata/search/utils/SearchUtilsTest.java
@@ -54,4 +54,22 @@ public class SearchUtilsTest {
                 "Expected all default values except skipHighlighting");
     }
 
+    @Test
+    public void testImmutableDefaults() throws CloneNotSupportedException {
+        SearchFlags defaultFlags = new SearchFlags()
+                .setFulltext(true)
+                .setSkipCache(true)
+                .setSkipAggregates(true)
+                .setMaxAggValues(1)
+                .setSkipHighlighting(true);
+        SearchFlags copyFlags = defaultFlags.copy();
+
+        assertEquals(SearchUtils.applyDefaultSearchFlags(new SearchFlags().setFulltext(false).setSkipCache(false)
+                        .setSkipAggregates(false).setMaxAggValues(2).setSkipHighlighting(false), "not empty", defaultFlags),
+                new SearchFlags().setFulltext(false).setSkipAggregates(false).setSkipCache(false).setMaxAggValues(2).setSkipHighlighting(false),
+                "Expected no default values");
+
+        assertEquals(defaultFlags, copyFlags, "Expected defaults to be unmodified");
+    }
+
 }