fix: use aggregation search instead of suggest query (#13861)

* fix: use searchAggregates instead of suggest query

* fix: aggregation fields

openmetadata-ui/src/main/resources/ui/src/constants/AdvancedSearch.constants.ts

@@ -12,7 +12,6 @@
  */
 
 import { t } from 'i18next';
-import { isUndefined, uniq } from 'lodash';
 import {
   BasicConfig,
   Fields,
@@ -24,7 +23,6 @@ import AntdConfig from 'react-awesome-query-builder/lib/config/antd';
 import { EntityFields, SuggestionField } from '../enums/AdvancedSearch.enum';
 import { SearchIndex } from '../enums/search.enum';
 import { getAggregateFieldOptions } from '../rest/miscAPI';
-import { suggestQuery } from '../rest/searchAPI';
 import { renderAdvanceSearchButtons } from '../utils/AdvancedSearchUtils';
 import { getCombinedQueryFilterObject } from '../utils/ExplorePage/ExplorePageUtils';
 
@@ -220,64 +218,28 @@ export const emptyJsonTree: JsonTree = {
  */
 export const autocomplete: (args: {
   searchIndex: SearchIndex | SearchIndex[];
-  entitySearchIndex: SearchIndex | SearchIndex[];
   entityField: EntityFields;
   suggestField?: SuggestionField;
-}) => SelectFieldSettings['asyncFetch'] = ({
+}) => SelectFieldSettings['asyncFetch'] = ({ searchIndex, entityField }) => {
-  searchIndex,
-  suggestField,
-  entitySearchIndex,
-  entityField,
-}) => {
-  const isUserAndTeamSearchIndex =
-    searchIndex.includes(SearchIndex.USER) ||
-    searchIndex.includes(SearchIndex.TEAM);
-
   return (search) => {
-    if (search) {
+    return getAggregateFieldOptions(
-      return suggestQuery({
+      searchIndex,
-        query: search ?? '*',
+      entityField,
-        searchIndex: searchIndex,
+      search ?? '',
-        field: suggestField,
+      JSON.stringify(getCombinedQueryFilterObject())
-        // fetch source if index is type of user or team and both
+    ).then((response) => {
-        fetchSource: isUserAndTeamSearchIndex,
+      const buckets =
-      }).then((resp) => {
+        response.data.aggregations[`sterms#${entityField}`].buckets;
-        return {
-          values: uniq(resp).map(({ text, _source }) => {
-            // set displayName or name if index is type of user or team and both.
-            // else set the text
-            const name =
-              isUserAndTeamSearchIndex && !isUndefined(_source)
-                ? _source?.displayName || _source.name
-                : text;
 
-            return {
+      return {
-              value: name,
+        values: buckets.map((bucket) => ({
-              title: name,
+          value: bucket.key,
-            };
+          title: bucket.label ?? bucket.key,
-          }),
+        })),
-          hasMore: false,
+        hasMore: false,
-        };
+      };
-      });
+    });
-    } else {
+    // }
-      return getAggregateFieldOptions(
-        entitySearchIndex,
-        entityField,
-        '',
-        JSON.stringify(getCombinedQueryFilterObject())
-      ).then((response) => {
-        const buckets =
-          response.data.aggregations[`sterms#${entityField}`].buckets;
-
-        return {
-          values: buckets.map((bucket) => ({
-            value: bucket.key,
-            title: bucket.label ?? bucket.key,
-          })),
-          hasMore: false,
-        };
-      });
-    }
   };
 };
 
@@ -306,8 +268,10 @@ const getCommonQueryBuilderFields = (
 
       fieldSettings: {
         asyncFetch: autocomplete({
-          searchIndex: [SearchIndex.USER, SearchIndex.TEAM],
+          searchIndex: entitySearchIndex ?? [
-          entitySearchIndex: [SearchIndex.USER, SearchIndex.TEAM],
+            SearchIndex.USER,
+            SearchIndex.TEAM,
+          ],
           entityField: EntityFields.OWNER,
         }),
         useAsyncSearch: true,
@@ -320,8 +284,10 @@ const getCommonQueryBuilderFields = (
       mainWidgetProps,
       fieldSettings: {
         asyncFetch: autocomplete({
-          searchIndex: [SearchIndex.TAG, SearchIndex.GLOSSARY],
+          searchIndex: entitySearchIndex ?? [
-          entitySearchIndex,
+            SearchIndex.TAG,
+            SearchIndex.GLOSSARY,
+          ],
           entityField: EntityFields.TAG,
         }),
         useAsyncSearch: true,
@@ -334,8 +300,7 @@ const getCommonQueryBuilderFields = (
       mainWidgetProps,
       fieldSettings: {
         asyncFetch: autocomplete({
-          searchIndex: [SearchIndex.TAG, SearchIndex.GLOSSARY],
+          searchIndex: entitySearchIndex ?? [SearchIndex.TAG],
-          entitySearchIndex,
           entityField: EntityFields.TIER,
         }),
         useAsyncSearch: true,
@@ -364,9 +329,7 @@ const getServiceQueryBuilderFields = (index: SearchIndex) => {
       fieldSettings: {
         asyncFetch: autocomplete({
           searchIndex: index,
-          entitySearchIndex: index,
           entityField: EntityFields.SERVICE,
-          suggestField: SuggestionField.SERVICE,
         }),
         useAsyncSearch: true,
       },
@@ -387,9 +350,7 @@ const tableQueryBuilderFields: Fields = {
     fieldSettings: {
       asyncFetch: autocomplete({
         searchIndex: SearchIndex.TABLE,
-        entitySearchIndex: SearchIndex.TABLE,
         entityField: EntityFields.DATABASE,
-        suggestField: SuggestionField.DATABASE,
       }),
       useAsyncSearch: true,
     },
@@ -402,9 +363,7 @@ const tableQueryBuilderFields: Fields = {
     fieldSettings: {
       asyncFetch: autocomplete({
         searchIndex: SearchIndex.TABLE,
-        entitySearchIndex: SearchIndex.TABLE,
         entityField: EntityFields.DATABASE_SCHEMA,
-        suggestField: SuggestionField.SCHEMA,
       }),
       useAsyncSearch: true,
     },
@@ -417,9 +376,7 @@ const tableQueryBuilderFields: Fields = {
     fieldSettings: {
       asyncFetch: autocomplete({
         searchIndex: SearchIndex.TABLE,
-        entitySearchIndex: SearchIndex.TABLE,
         entityField: EntityFields.COLUMN,
-        suggestField: SuggestionField.COLUMN,
       }),
       useAsyncSearch: true,
     },

openmetadata-ui/src/main/resources/ui/src/enums/AdvancedSearch.enum.ts

@@ -38,10 +38,10 @@ export enum EntityFields {
   OWNER = 'displayName.keyword',
   TAG = 'tags.tagFQN',
   TIER = 'tier.tagFQN',
-  SERVICE = 'service.name',
+  SERVICE = 'service.name.keyword',
-  DATABASE = 'database.name',
+  DATABASE = 'database.name.keyword',
-  DATABASE_SCHEMA = 'databaseSchema.name',
+  DATABASE_SCHEMA = 'databaseSchema.name.keyword',
-  COLUMN = 'columns.name',
+  COLUMN = 'columns.name.keyword',
   CHART = 'charts.displayName.keyword',
   TASK = 'tasks.displayName.keyword',
 }

openmetadata-ui/src/main/resources/ui/src/pages/DataInsightPage/DataInsightProvider.tsx

@@ -51,7 +51,6 @@ export const DataInsightContext = createContext<DataInsightContextType>(
 );
 const fetchTeamSuggestions = autocomplete({
   searchIndex: SearchIndex.TEAM,
-  entitySearchIndex: SearchIndex.TEAM,
   entityField: EntityFields.OWNER,
 });
 

openmetadata-ui/src/main/resources/ui/src/rest/searchAPI.test.ts

@@ -83,113 +83,9 @@ const mockTableSearchResponse = {
   },
 };
 
-const mockSuggestUserResponse = {
-  suggest: {
-    'metadata-suggest': [
-      {
-        text: 'a',
-        offset: 0,
-        length: 1,
-        options: [
-          {
-            text: 'Aaron Johnson',
-            _index: 'user_search_index',
-            _type: '_doc',
-            _id: '2cae227c-e2c4-487c-b52c-a96ae242d90d',
-            _score: 10.0,
-            _source: {
-              id: '2cae227c-e2c4-487c-b52c-a96ae242d90d',
-              name: 'aaron_johnson0',
-              fullyQualifiedName: 'aaron_johnson0',
-              displayName: 'Aaron Johnson',
-              version: 0.1,
-              updatedAt: 1661336540995,
-              updatedBy: 'anonymous',
-              email: 'aaron_johnson0@gmail.com',
-              href: 'http://localhost:8585/api/v1/users/2cae227c-e2c4-487c-b52c-a96ae242d90d',
-              isAdmin: false,
-              deleted: false,
-              roles: [],
-              inheritedRoles: [],
-              entityType: 'user',
-              teams: null,
-              some: {
-                nested: {
-                  nullValue: null,
-                },
-              },
-            },
-          },
-        ],
-      },
-    ],
-  },
-};
-
 describe('searchAPI tests', () => {
   beforeEach(() => jest.resetModules());
 
-  it('suggestQuery should return object and aggregations', async () => {
-    jest.mock('./index', () => ({
-      get: jest
-        .fn()
-        .mockImplementation(() =>
-          Promise.resolve({ data: mockTableSearchResponse })
-        ),
-    }));
-    // eslint-disable-next-line @typescript-eslint/no-var-requires
-    const { searchQuery } = require('./searchAPI');
-    const res = await searchQuery({ searchIndex: SearchIndex.TABLE });
-
-    expect(res.hits.total.value).toBe(10_000);
-
-    expect(res.hits.hits).toHaveLength(1);
-    expect(res.hits.hits[0]._index).toEqual(SearchIndex.TABLE);
-    expect(res.hits.hits[0]._source).toEqual(
-      expect.objectContaining({
-        id: '9b30a945-239a-4cb7-93b0-f1b7425aed41',
-        name: 'raw_product_catalog',
-        fullyQualifiedName:
-          'sample_data.ecommerce_db.shopify.raw_product_catalog',
-        description:
-          'This is a raw product catalog table contains the product listing, price, seller etc.. represented in our online DB. ',
-        version: 0.1,
-        updatedAt: 1661336543968,
-        updatedBy: 'anonymous',
-        href: 'http://localhost:8585/api/v1/tables/9b30a945-239a-4cb7-93b0-f1b7425aed41',
-        tableType: 'Regular',
-      })
-    );
-
-    expect(res.aggregations).toEqual(
-      expect.objectContaining({
-        EntityType: {
-          buckets: expect.arrayContaining([
-            {
-              key: 'table',
-              doc_count: 10960,
-            },
-          ]),
-        },
-        ServiceName: {
-          buckets: expect.arrayContaining([
-            {
-              key: 'trino',
-              doc_count: 10924,
-            },
-            {
-              key: 'sample_data',
-              doc_count: 36,
-            },
-          ]),
-        },
-        Tags: {
-          buckets: [],
-        },
-      })
-    );
-  });
-
   it('searchQuery should not return nulls', async () => {
     jest.mock('./index', () => ({
       get: jest
@@ -224,70 +120,4 @@ describe('searchAPI tests', () => {
 
     expect(res.hits.hits[0]._source.type).toBe('table');
   });
-
-  it('suggestQuery should return object and text', async () => {
-    jest.mock('./index', () => ({
-      get: jest
-        .fn()
-        .mockImplementation(() =>
-          Promise.resolve({ data: mockSuggestUserResponse })
-        ),
-    }));
-    // eslint-disable-next-line @typescript-eslint/no-var-requires
-    const { suggestQuery } = require('./searchAPI');
-    const res = await suggestQuery({ searchIndex: SearchIndex.USER });
-
-    expect(res).toEqual([
-      expect.objectContaining({
-        _index: SearchIndex.USER,
-        _source: expect.objectContaining({
-          id: '2cae227c-e2c4-487c-b52c-a96ae242d90d',
-          name: 'aaron_johnson0',
-          fullyQualifiedName: 'aaron_johnson0',
-          displayName: 'Aaron Johnson',
-          version: 0.1,
-          updatedAt: 1661336540995,
-          updatedBy: 'anonymous',
-          email: 'aaron_johnson0@gmail.com',
-          href: 'http://localhost:8585/api/v1/users/2cae227c-e2c4-487c-b52c-a96ae242d90d',
-          isAdmin: false,
-          deleted: false,
-          roles: [],
-          inheritedRoles: [],
-          entityType: 'user',
-        }),
-      }),
-    ]);
-  });
-
-  it('suggestQuery should not return nulls', async () => {
-    jest.mock('./index', () => ({
-      get: jest
-        .fn()
-        .mockImplementation(() =>
-          Promise.resolve({ data: mockSuggestUserResponse })
-        ),
-    }));
-    // eslint-disable-next-line @typescript-eslint/no-var-requires
-    const { suggestQuery } = require('./searchAPI');
-    const res = await suggestQuery({ searchIndex: SearchIndex.USER });
-
-    // Deep checking for null values
-    expect(flatten(res[0]._source).filter(isNull)).toHaveLength(0);
-  });
-
-  it('suggestQuery should have type field', async () => {
-    jest.mock('./index', () => ({
-      get: jest
-        .fn()
-        .mockImplementation(() =>
-          Promise.resolve({ data: mockSuggestUserResponse })
-        ),
-    }));
-    // eslint-disable-next-line @typescript-eslint/no-var-requires
-    const { suggestQuery } = require('./searchAPI');
-    const res = await suggestQuery({ searchIndex: SearchIndex.USER });
-
-    expect(res[0]._source.type).toBe('user');
-  });
 });

openmetadata-ui/src/main/resources/ui/src/rest/searchAPI.ts

@@ -19,12 +19,9 @@ import {
   Aggregations,
   DataInsightSearchResponse,
   KeysOfUnion,
-  RawSuggestResponse,
   SearchIndexSearchSourceMapping,
   SearchRequest,
   SearchResponse,
-  SuggestRequest,
-  SuggestResponse,
 } from '../interface/search.interface';
 import { omitDeep } from '../utils/APIUtils';
 import { getQueryWithSlash } from '../utils/SearchUtils';
@@ -290,139 +287,3 @@ export const searchQueryDataInsight = async (
     res.data
   ) as unknown as DataInsightSearchResponse;
 };
-
-/**
- * Formats a response from {@link rawSuggestQuery}
- *
- * Warning: avoid this pattern unless applying custom transformation to the raw response!
- * ```ts
- * const response = await rawSuggestQuery(req);
- * const data = formatSuggestQueryResponse(response.data);
- * ```
- *
- * Instead use the shorthand {@link suggestQuery}
- * ```ts
- * const data = suggestQuery(req);
- * ```
- *
- * @param data
- */
-export const formatSuggestQueryResponse = <
-  SI extends SearchIndex | SearchIndex[],
-  TIncludeFields extends KeysOfUnion<
-    SearchIndexSearchSourceMapping[SI extends Array<SearchIndex>
-      ? SI[number]
-      : SI]
-  >
->(
-  data: RawSuggestResponse<
-    SI extends Array<SearchIndex> ? SI[number] : SI,
-    TIncludeFields
-  >
-): SuggestResponse<
-  SI extends Array<SearchIndex> ? SI[number] : SI,
-  TIncludeFields
-> => {
-  let _data;
-
-  _data = data;
-
-  // Elasticsearch responses use 'null' for missing values, we want undefined
-  _data = omitDeep<
-    SuggestResponse<
-      SI extends Array<SearchIndex> ? SI[number] : SI,
-      TIncludeFields
-    >
-  >(_data.suggest['metadata-suggest'][0].options, isNil);
-
-  /* Elasticsearch objects use `entityType` to track their type, but the EntityReference interface uses `type`
-      This copies `entityType` into `type` (if `entityType` exists) so responses implement EntityReference */
-  _data = _data.map((datum) =>
-    '_source' in datum
-      ? 'entityType' in datum._source
-        ? {
-            ...datum,
-            _source: {
-              ...(datum._source as SearchIndexSearchSourceMapping[SI extends Array<SearchIndex>
-                ? SI[number]
-                : SI]),
-              type: (
-                datum._source as SearchIndexSearchSourceMapping[SI extends Array<SearchIndex>
-                  ? SI[number]
-                  : SI]
-              ).entityType,
-            },
-          }
-        : datum
-      : datum
-  );
-
-  return _data;
-};
-
-/**
- * Executes a request to /search/suggest, returning the raw response.
- * Warning: Only call this function directly in special cases. Otherwise use {@link suggestQuery}
- *
- * @param req Request object
- */
-export const rawSuggestQuery = <
-  SI extends SearchIndex | SearchIndex[],
-  TIncludeFields extends KeysOfUnion<
-    SearchIndexSearchSourceMapping[SI extends Array<SearchIndex>
-      ? SI[number]
-      : SI]
-  >
->(
-  req: SuggestRequest<SI, TIncludeFields>
-): Promise<
-  AxiosResponse<
-    RawSuggestResponse<
-      SI extends Array<SearchIndex> ? SI[number] : SI,
-      TIncludeFields
-    >
-  >
-> => {
-  const { query, searchIndex, field, fetchSource } = req;
-
-  return APIClient.get<
-    RawSuggestResponse<
-      SI extends Array<SearchIndex> ? SI[number] : SI,
-      TIncludeFields
-    >
-  >('/search/suggest', {
-    params: {
-      q: query,
-      field,
-      index: getSearchIndexParam(searchIndex),
-      fetch_source: fetchSource,
-      include_source_fields: req.fetchSource ? req.includeFields : undefined,
-    },
-  });
-};
-
-/**
- * Access point for the Suggestion API.
- * Executes a request to /search/suggest, returning a formatted response.
- *
- * @param req Request object
- */
-export const suggestQuery = async <
-  SI extends SearchIndex | SearchIndex[],
-  TIncludeFields extends KeysOfUnion<
-    SearchIndexSearchSourceMapping[SI extends Array<SearchIndex>
-      ? SI[number]
-      : SI]
-  >
->(
-  req: SuggestRequest<SI, TIncludeFields>
-): Promise<
-  SuggestResponse<
-    SI extends Array<SearchIndex> ? SI[number] : SI,
-    TIncludeFields
-  >
-> => {
-  const res = await rawSuggestQuery(req);
-
-  return formatSuggestQueryResponse(res.data);
-};