Improve SOR.find reference filter implementation (#121042) (#121177)

Co-authored-by: Pierre Gayvallet <pierre.gayvallet@elastic.co>

src/core/server/saved_objects/service/lib/search_dsl/query_params.test.ts

@@ -6,12 +6,15 @@
  * Side Public License, v 1.
  */
 
+import { getReferencesFilterMock } from './query_params.tests.mocks';
+
 import * as esKuery from '@kbn/es-query';
+
 type KueryNode = any;
 
 import { ALL_NAMESPACES_STRING, DEFAULT_NAMESPACE_STRING } from '../utils';
 import { SavedObjectTypeRegistry } from '../../../saved_objects_type_registry';
-import { getQueryParams, getClauseForReference } from './query_params';
+import { getQueryParams } from './query_params';
 
 const registerTypes = (registry: SavedObjectTypeRegistry) => {
   registry.registerType({
@@ -85,6 +88,12 @@ describe('#getQueryParams', () => {
   beforeEach(() => {
     registry = new SavedObjectTypeRegistry();
     registerTypes(registry);
+
+    getReferencesFilterMock.mockReturnValue({ references_filter: true });
+  });
+
+  afterEach(() => {
+    getReferencesFilterMock.mockClear();
   });
 
   const createTypeClause = (type: string, namespaces?: string[]) => {
@@ -185,102 +194,42 @@ describe('#getQueryParams', () => {
 
   describe('reference filter clause', () => {
     describe('`hasReference` parameter', () => {
-      const getReferencesFilter = (result: any) => {
+      it('does not call `getReferencesFilter` when `hasReference` is not specified', () => {
-        const filters = result.query.bool.filter;
+        getQueryParams({
-        return filters.find((filter: any) => {
-          const clauses = filter.bool?.must ?? filter.bool?.should;
-          if (!clauses) {
-            return false;
-          }
-          return clauses[0].nested?.path === 'references' ?? false;
-        });
-      };
-
-      it('does not include the clause when `hasReference` is not specified', () => {
-        const result = getQueryParams({
           registry,
           hasReference: undefined,
         });
 
-        expect(getReferencesFilter(result)).toBeUndefined();
+        expect(getReferencesFilterMock).not.toHaveBeenCalled();
       });
 
-      it('creates a should clause for specified reference when operator is `OR`', () => {
+      it('calls `getReferencesFilter` with the correct parameters', () => {
         const hasReference = { id: 'foo', type: 'bar' };
-        const result = getQueryParams({
+        getQueryParams({
           registry,
           hasReference,
-          hasReferenceOperator: 'OR',
+          hasReferenceOperator: 'AND',
         });
-        expect(getReferencesFilter(result)).toEqual({
+
-          bool: {
+        expect(getReferencesFilterMock).toHaveBeenCalledTimes(1);
-            should: [getClauseForReference(hasReference)],
+        expect(getReferencesFilterMock).toHaveBeenCalledWith({
-            minimum_should_match: 1,
+          references: [hasReference],
-          },
+          operator: 'AND',
         });
       });
 
-      it('creates a must clause for specified reference when operator is `AND`', () => {
+      it('includes the return of `getReferencesFilter` in the `filter` clause', () => {
+        getReferencesFilterMock.mockReturnValue({ references_filter: true });
+
         const hasReference = { id: 'foo', type: 'bar' };
         const result = getQueryParams({
           registry,
           hasReference,
           hasReferenceOperator: 'AND',
         });
-        expect(getReferencesFilter(result)).toEqual({
-          bool: {
-            must: [getClauseForReference(hasReference)],
-          },
-        });
-      });
 
-      it('handles multiple references when operator is `OR`', () => {
+        const filters: any[] = result.query.bool.filter;
-        const hasReference = [
+        expect(filters.some((filter) => filter.references_filter === true)).toBeDefined();
-          { id: 'foo', type: 'bar' },
-          { id: 'hello', type: 'dolly' },
-        ];
-        const result = getQueryParams({
-          registry,
-          hasReference,
-          hasReferenceOperator: 'OR',
-        });
-        expect(getReferencesFilter(result)).toEqual({
-          bool: {
-            should: hasReference.map(getClauseForReference),
-            minimum_should_match: 1,
-          },
-        });
-      });
-
-      it('handles multiple references when operator is `AND`', () => {
-        const hasReference = [
-          { id: 'foo', type: 'bar' },
-          { id: 'hello', type: 'dolly' },
-        ];
-        const result = getQueryParams({
-          registry,
-          hasReference,
-          hasReferenceOperator: 'AND',
-        });
-        expect(getReferencesFilter(result)).toEqual({
-          bool: {
-            must: hasReference.map(getClauseForReference),
-          },
-        });
-      });
-
-      it('defaults to `OR` when operator is not specified', () => {
-        const hasReference = { id: 'foo', type: 'bar' };
-        const result = getQueryParams({
-          registry,
-          hasReference,
-        });
-        expect(getReferencesFilter(result)).toEqual({
-          bool: {
-            should: [getClauseForReference(hasReference)],
-            minimum_should_match: 1,
-          },
-        });
       });
     });
   });

src/core/server/saved_objects/service/lib/search_dsl/query_params.tests.mocks.ts

@@ -0,0 +1,13 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+export const getReferencesFilterMock = jest.fn();
+
+jest.doMock('./references_filter', () => ({
+  getReferencesFilter: getReferencesFilterMock,
+}));

src/core/server/saved_objects/service/lib/search_dsl/query_params.ts

@@ -7,10 +7,12 @@
  */
 
 import * as esKuery from '@kbn/es-query';
+
 type KueryNode = any;
 
 import { ISavedObjectTypeRegistry } from '../../../saved_objects_type_registry';
 import { ALL_NAMESPACES_STRING, DEFAULT_NAMESPACE_STRING } from '../utils';
+import { getReferencesFilter } from './references_filter';
 
 /**
  * Gets the types based on the type. Uses mappings to support
@@ -139,50 +141,6 @@ interface QueryParams {
   kueryNode?: KueryNode;
 }
 
-function getReferencesFilter(
-  references: HasReferenceQueryParams[],
-  operator: SearchOperator = 'OR'
-) {
-  if (operator === 'AND') {
-    return {
-      bool: {
-        must: references.map(getClauseForReference),
-      },
-    };
-  } else {
-    return {
-      bool: {
-        should: references.map(getClauseForReference),
-        minimum_should_match: 1,
-      },
-    };
-  }
-}
-
-export function getClauseForReference(reference: HasReferenceQueryParams) {
-  return {
-    nested: {
-      path: 'references',
-      query: {
-        bool: {
-          must: [
-            {
-              term: {
-                'references.id': reference.id,
-              },
-            },
-            {
-              term: {
-                'references.type': reference.type,
-              },
-            },
-          ],
-        },
-      },
-    },
-  };
-}
-
 // A de-duplicated set of namespaces makes for a more efficient query.
 const uniqNamespaces = (namespacesToNormalize?: string[]) =>
   namespacesToNormalize ? Array.from(new Set(namespacesToNormalize)) : undefined;
@@ -215,7 +173,14 @@ export function getQueryParams({
   const bool: any = {
     filter: [
       ...(kueryNode != null ? [esKuery.toElasticsearchQuery(kueryNode)] : []),
-      ...(hasReference?.length ? [getReferencesFilter(hasReference, hasReferenceOperator)] : []),
+      ...(hasReference?.length
+        ? [
+            getReferencesFilter({
+              references: hasReference,
+              operator: hasReferenceOperator,
+            }),
+          ]
+        : []),
       {
         bool: {
           should: types.map((shouldType) => {

src/core/server/saved_objects/service/lib/search_dsl/references_filter.test.ts

@@ -0,0 +1,164 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { getReferencesFilter } from './references_filter';
+
+describe('getReferencesFilter', () => {
+  const nestedRefMustClauses = (nestedMustClauses: unknown[]) => ({
+    nested: {
+      path: 'references',
+      query: {
+        bool: {
+          must: nestedMustClauses,
+        },
+      },
+    },
+  });
+
+  describe('when using the `OR` operator', () => {
+    it('generates one `should` clause per type of reference', () => {
+      const references = [
+        { type: 'foo', id: 'foo-1' },
+        { type: 'foo', id: 'foo-2' },
+        { type: 'foo', id: 'foo-3' },
+        { type: 'bar', id: 'bar-1' },
+        { type: 'bar', id: 'bar-2' },
+      ];
+      const clause = getReferencesFilter({
+        references,
+        operator: 'OR',
+      });
+
+      expect(clause).toEqual({
+        bool: {
+          should: [
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['foo-1', 'foo-2', 'foo-3'] } },
+              { term: { 'references.type': 'foo' } },
+            ]),
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['bar-1', 'bar-2'] } },
+              { term: { 'references.type': 'bar' } },
+            ]),
+          ],
+          minimum_should_match: 1,
+        },
+      });
+    });
+
+    it('does not include mode than `maxTermsPerClause` per `terms` clauses', () => {
+      const references = [
+        { type: 'foo', id: 'foo-1' },
+        { type: 'foo', id: 'foo-2' },
+        { type: 'foo', id: 'foo-3' },
+        { type: 'foo', id: 'foo-4' },
+        { type: 'foo', id: 'foo-5' },
+        { type: 'bar', id: 'bar-1' },
+        { type: 'bar', id: 'bar-2' },
+        { type: 'bar', id: 'bar-3' },
+        { type: 'dolly', id: 'dolly-1' },
+      ];
+      const clause = getReferencesFilter({
+        references,
+        operator: 'OR',
+        maxTermsPerClause: 2,
+      });
+
+      expect(clause).toEqual({
+        bool: {
+          should: [
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['foo-1', 'foo-2'] } },
+              { term: { 'references.type': 'foo' } },
+            ]),
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['foo-3', 'foo-4'] } },
+              { term: { 'references.type': 'foo' } },
+            ]),
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['foo-5'] } },
+              { term: { 'references.type': 'foo' } },
+            ]),
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['bar-1', 'bar-2'] } },
+              { term: { 'references.type': 'bar' } },
+            ]),
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['bar-3'] } },
+              { term: { 'references.type': 'bar' } },
+            ]),
+            nestedRefMustClauses([
+              { terms: { 'references.id': ['dolly-1'] } },
+              { term: { 'references.type': 'dolly' } },
+            ]),
+          ],
+          minimum_should_match: 1,
+        },
+      });
+    });
+  });
+
+  describe('when using the `AND` operator', () => {
+    it('generates one `must` clause per reference', () => {
+      const references = [
+        { type: 'foo', id: 'foo-1' },
+        { type: 'foo', id: 'foo-2' },
+        { type: 'bar', id: 'bar-1' },
+      ];
+
+      const clause = getReferencesFilter({
+        references,
+        operator: 'AND',
+      });
+
+      expect(clause).toEqual({
+        bool: {
+          must: references.map((ref) => ({
+            nested: {
+              path: 'references',
+              query: {
+                bool: {
+                  must: [
+                    { term: { 'references.id': ref.id } },
+                    { term: { 'references.type': ref.type } },
+                  ],
+                },
+              },
+            },
+          })),
+        },
+      });
+    });
+  });
+
+  it('defaults to using the `OR` operator', () => {
+    const references = [
+      { type: 'foo', id: 'foo-1' },
+      { type: 'bar', id: 'bar-1' },
+    ];
+    const clause = getReferencesFilter({
+      references,
+    });
+
+    expect(clause).toEqual({
+      bool: {
+        should: [
+          nestedRefMustClauses([
+            { terms: { 'references.id': ['foo-1'] } },
+            { term: { 'references.type': 'foo' } },
+          ]),
+          nestedRefMustClauses([
+            { terms: { 'references.id': ['bar-1'] } },
+            { term: { 'references.type': 'bar' } },
+          ]),
+        ],
+        minimum_should_match: 1,
+      },
+    });
+  });
+});

src/core/server/saved_objects/service/lib/search_dsl/references_filter.ts

@@ -0,0 +1,107 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import type { HasReferenceQueryParams, SearchOperator } from './query_params';
+
+export function getReferencesFilter({
+  references,
+  operator = 'OR',
+  maxTermsPerClause = 1000,
+}: {
+  references: HasReferenceQueryParams[];
+  operator?: SearchOperator;
+  maxTermsPerClause?: number;
+}) {
+  if (operator === 'AND') {
+    return {
+      bool: {
+        must: references.map(getNestedTermClauseForReference),
+      },
+    };
+  } else {
+    return {
+      bool: {
+        should: getAggregatedTermsClauses(references, maxTermsPerClause),
+        minimum_should_match: 1,
+      },
+    };
+  }
+}
+
+const getAggregatedTermsClauses = (
+  references: HasReferenceQueryParams[],
+  maxTermsPerClause: number
+) => {
+  const refTypeToIds = references.reduce((map, { type, id }) => {
+    const ids = map.get(type) ?? [];
+    map.set(type, [...ids, id]);
+    return map;
+  }, new Map<string, string[]>());
+
+  // we create chunks per type to avoid generating `terms` clauses with too many terms
+  const typeIdChunks = [...refTypeToIds.entries()].flatMap(([type, ids]) => {
+    return createChunks(ids, maxTermsPerClause).map((chunkIds) => ({ type, ids: chunkIds }));
+  });
+
+  return typeIdChunks.map(({ type, ids }) => getNestedTermsClausesForReferences(type, ids));
+};
+
+const createChunks = <T>(array: T[], chunkSize: number): T[][] => {
+  const chunks: T[][] = [];
+  for (let i = 0, len = array.length; i < len; i += chunkSize)
+    chunks.push(array.slice(i, i + chunkSize));
+  return chunks;
+};
+
+export const getNestedTermClauseForReference = (reference: HasReferenceQueryParams) => {
+  return {
+    nested: {
+      path: 'references',
+      query: {
+        bool: {
+          must: [
+            {
+              term: {
+                'references.id': reference.id,
+              },
+            },
+            {
+              term: {
+                'references.type': reference.type,
+              },
+            },
+          ],
+        },
+      },
+    },
+  };
+};
+
+const getNestedTermsClausesForReferences = (type: string, ids: string[]) => {
+  return {
+    nested: {
+      path: 'references',
+      query: {
+        bool: {
+          must: [
+            {
+              terms: {
+                'references.id': ids,
+              },
+            },
+            {
+              term: {
+                'references.type': type,
+              },
+            },
+          ],
+        },
+      },
+    },
+  };
+};