[getSavedObjectsCount] Use soClient instead of .kibana searches

packages/core/saved-objects/core-saved-objects-api-server-internal/src/lib/aggregations/validation.test.ts

@@ -522,4 +522,21 @@ describe('validateAndConvertAggregations', () => {
       '"[aggName.cardinality.field] Invalid attribute path: alert.alert.actions.group"'
     );
   });
+
+  it('allows aggregations for root fields', () => {
+    const aggregations: AggsMap = {
+      types: {
+        terms: {
+          field: 'type',
+        },
+      },
+    };
+    expect(validateAndConvertAggregations(['foo'], aggregations, mockMappings)).toEqual({
+      types: {
+        terms: {
+          field: 'type',
+        },
+      },
+    });
+  });
 });

packages/core/saved-objects/core-saved-objects-api-server-internal/src/lib/aggregations/validation.ts

@@ -18,6 +18,7 @@ import {
   rewriteRootLevelAttribute,
 } from './validation_utils';
 import { aggregationSchemas } from './aggs_types';
+import { getRootFields } from '../included_fields';
 
 const aggregationKeys = ['aggs', 'aggregations'];
 
@@ -226,6 +227,10 @@ const isAttributeValue = (fieldName: string, fieldValue: unknown): boolean => {
   return attributeFields.includes(fieldName) && typeof fieldValue === 'string';
 };
 
+const isRootField = (fieldName: string): boolean => {
+  return getRootFields().includes(fieldName);
+};
+
 const validateAndRewriteAttributePath = (
   attributePath: string,
   { allowedTypes, indexMapping, currentPath }: ValidationContext
@@ -236,5 +241,8 @@ const validateAndRewriteAttributePath = (
   if (isObjectTypeAttribute(attributePath, indexMapping, allowedTypes)) {
     return rewriteObjectTypeAttribute(attributePath);
   }
+  if (isRootField(attributePath)) {
+    return attributePath;
+  }
   throw new Error(`[${currentPath.join('.')}] Invalid attribute path: ${attributePath}`);
 };

src/plugins/kibana_usage_collection/server/collectors/saved_objects_counts/get_saved_object_counts/get_saved_object_counts.test.ts

@@ -6,71 +6,54 @@
  * Side Public License, v 1.
  */
 
-import { elasticsearchServiceMock } from '@kbn/core/server/mocks';
+import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server';
+import { createCollectorFetchContextMock } from '@kbn/usage-collection-plugin/server/mocks';
 import { getSavedObjectsCounts } from './get_saved_object_counts';
 
-function mockGetSavedObjectsCounts<TBody>(params: TBody) {
-  const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
-  // @ts-expect-error arbitrary type
-  esClient.search.mockResponse(params);
-  return esClient;
-}
+const soEmptyResponse = { total: 0, saved_objects: [], per_page: 0, page: 1 };
 
 describe('getSavedObjectsCounts', () => {
+  const fetchContextMock = createCollectorFetchContextMock();
+  const soClient = fetchContextMock.soClient as jest.Mocked<SavedObjectsClientContract>;
+
+  beforeEach(() => {
+    soClient.find.mockReset();
+  });
+
   test('should not fail if no body returned', async () => {
-    const esClient = mockGetSavedObjectsCounts({});
+    soClient.find.mockResolvedValueOnce(soEmptyResponse);
 
-    const results = await getSavedObjectsCounts(esClient, '.kibana', ['type-a']);
+    const results = await getSavedObjectsCounts(soClient, ['type-a']);
     // Make sure ES.search is triggered (we'll test the actual params in other specific tests)
-    expect(esClient.search).toHaveBeenCalledTimes(1);
+    expect(soClient.find).toHaveBeenCalledTimes(1);
     expect(results).toStrictEqual({ total: 0, per_type: [], non_expected_types: [], others: 0 });
   });
 
   test('should match all and request the `missing` bucket (size + 1) when `exclusive === false`', async () => {
-    const esClient = mockGetSavedObjectsCounts({});
-    await getSavedObjectsCounts(esClient, '.kibana', ['type-a', 'type_2']);
-    expect(esClient.search).toHaveBeenCalledWith({
-      index: '.kibana',
-      ignore_unavailable: true,
-      filter_path: [
-        'aggregations.types.buckets',
-        'aggregations.types.sum_other_doc_count',
-        'hits.total',
-      ],
-      body: {
-        size: 0,
-        track_total_hits: true,
-        query: { match_all: {} },
-        aggs: {
-          types: {
-            terms: {
-              field: 'type',
-              size: 3,
-              missing: 'missing_so_type',
-            },
+    soClient.find.mockResolvedValueOnce(soEmptyResponse);
+    await getSavedObjectsCounts(soClient, ['type-a', 'type_2']);
+    expect(soClient.find).toHaveBeenCalledWith({
+      type: ['type-a', 'type_2'],
+      perPage: 0,
+      aggs: {
+        types: {
+          terms: {
+            field: 'type',
+            size: 3,
+            missing: 'missing_so_type',
           },
         },
       },
     });
   });
 
   test('should apply the terms query and aggregation with the size matching the length of the list when `exclusive === true`', async () => {
-    const esClient = mockGetSavedObjectsCounts({});
-    await getSavedObjectsCounts(esClient, '.kibana', ['type_one', 'type_two'], true);
-    expect(esClient.search).toHaveBeenCalledWith({
-      index: '.kibana',
-      ignore_unavailable: true,
-      filter_path: [
-        'aggregations.types.buckets',
-        'aggregations.types.sum_other_doc_count',
-        'hits.total',
-      ],
-      body: {
-        size: 0,
-        track_total_hits: true,
-        query: { terms: { type: ['type_one', 'type_two'] } },
-        aggs: { types: { terms: { field: 'type', size: 2 } } },
-      },
+    soClient.find.mockResolvedValueOnce(soEmptyResponse);
+    await getSavedObjectsCounts(soClient, ['type_one', 'type_two'], true);
+    expect(soClient.find).toHaveBeenCalledWith({
+      type: ['type_one', 'type_two'],
+      perPage: 0,
+      aggs: { types: { terms: { field: 'type', size: 2 } } },
     });
   });
 
@@ -80,39 +63,13 @@ describe('getSavedObjectsCounts', () => {
       { key: 'type-two', doc_count: 2 },
     ];
 
-    const esClient = mockGetSavedObjectsCounts({
-      hits: { total: { value: 13 } },
-      aggregations: { types: { buckets, sum_other_doc_count: 10 } },
-    });
-
-    const results = await getSavedObjectsCounts(esClient, '.kibana', [
-      'type_one',
-      'type-two',
-      'type-3',
-    ]);
-    expect(results).toStrictEqual({
+    soClient.find.mockResolvedValueOnce({
+      ...soEmptyResponse,
       total: 13,
-      per_type: [
-        { key: 'type_one', doc_count: 1 },
-        { key: 'type-two', doc_count: 2 },
-      ],
-      non_expected_types: [],
-      others: 10,
-    });
-  });
-
-  test('supports ES returning total as a number (just in case)', async () => {
-    const buckets = [
-      { key: 'type_one', doc_count: 1 },
-      { key: 'type-two', doc_count: 2 },
-    ];
-
-    const esClient = mockGetSavedObjectsCounts({
-      hits: { total: 13 },
       aggregations: { types: { buckets, sum_other_doc_count: 10 } },
     });
 
-    const results = await getSavedObjectsCounts(esClient, '.kibana', ['type_one', 'type-two']);
+    const results = await getSavedObjectsCounts(soClient, ['type_one', 'type-two', 'type-3']);
     expect(results).toStrictEqual({
       total: 13,
       per_type: [
@@ -132,12 +89,13 @@ describe('getSavedObjectsCounts', () => {
       { key: 'type-four', doc_count: 2 },
     ];
 
-    const esClient = mockGetSavedObjectsCounts({
-      hits: { total: { value: 13 } },
+    soClient.find.mockResolvedValueOnce({
+      ...soEmptyResponse,
+      total: 13,
       aggregations: { types: { buckets, sum_other_doc_count: 6 } },
     });
 
-    const results = await getSavedObjectsCounts(esClient, '.kibana', ['type_one', 'type-two']);
+    const results = await getSavedObjectsCounts(soClient, ['type_one', 'type-two']);
     expect(results).toStrictEqual({
       total: 13,
       per_type: [

src/plugins/kibana_usage_collection/server/collectors/saved_objects_counts/get_saved_object_counts/get_saved_object_counts.ts

@@ -7,7 +7,7 @@
  */
 
 import { estypes } from '@elastic/elasticsearch';
-import { ElasticsearchClient } from '@kbn/core/server';
+import type { SavedObjectsClientContract } from '@kbn/core/server';
 
 const MISSING_TYPE_KEY = 'missing_so_type';
 
@@ -39,40 +39,28 @@ export interface SavedObjectsCounts {
  * It also returns a break-down of the document count for all the built-in SOs in Kibana (or the types specified in `soTypes`).
  * Finally, it completes the information with an `others` counter, that indicates the number of documents that do not match the SO type breakdown.
  *
- * @param esClient The {@link ElasticsearchClient} to use when performing the aggregation.
- * @param kibanaIndex The index where SOs are stored. Typically '.kibana'.
+ * @param soClient The {@link SavedObjectsClientContract} to use when performing the aggregation.
  * @param soTypes The SO types we want to know about.
  * @param exclusive If `true`, the results will only contain the breakdown for the specified `soTypes`. Otherwise, it'll also return `missing` and `others` bucket.
  * @returns {@link SavedObjectsCounts}
  */
 export async function getSavedObjectsCounts(
-  esClient: ElasticsearchClient,
-  kibanaIndex: string, // Typically '.kibana'. We might need a way to obtain it from the SavedObjects client (or the SavedObjects client to provide a way to run aggregations?)
+  soClient: SavedObjectsClientContract,
   soTypes: string[],
   exclusive: boolean = false
 ): Promise<SavedObjectsCounts> {
-  const body = await esClient.search<void, { types: estypes.AggregationsStringTermsAggregate }>({
-    index: kibanaIndex,
-    ignore_unavailable: true,
-    filter_path: [
-      'aggregations.types.buckets',
-      'aggregations.types.sum_other_doc_count',
-      'hits.total',
-    ],
-    body: {
-      size: 0,
-      track_total_hits: true,
-      query: exclusive ? { terms: { type: soTypes } } : { match_all: {} },
-      aggs: {
-        types: {
-          terms: {
-            field: 'type',
-            // If `exclusive == true`, we only care about the strict length of the provided SO types.
-            // Otherwise, we want to account for the `missing` bucket (size and missing option).
-            ...(exclusive
-              ? { size: soTypes.length }
-              : { missing: MISSING_TYPE_KEY, size: soTypes.length + 1 }),
-          },
+  const body = await soClient.find<void, { types: estypes.AggregationsStringTermsAggregate }>({
+    type: soTypes,
+    perPage: 0,
+    aggs: {
+      types: {
+        terms: {
+          field: 'type',
+          // If `exclusive == true`, we only care about the strict length of the provided SO types.
+          // Otherwise, we want to account for the `missing` bucket (size and missing option).
+          ...(exclusive
+            ? { size: soTypes.length }
+            : { missing: MISSING_TYPE_KEY, size: soTypes.length + 1 }),
         },
       },
     },
@@ -93,7 +81,7 @@ export async function getSavedObjectsCounts(
   });
 
   return {
-    total: (typeof body.hits?.total === 'number' ? body.hits?.total : body.hits?.total?.value) ?? 0,
+    total: body.total,
     per_type: perType,
     non_expected_types: nonExpectedTypes,
     others: body.aggregations?.types?.sum_other_doc_count ?? 0,

src/plugins/kibana_usage_collection/server/collectors/saved_objects_counts/kibana_usage_collector.test.ts

@@ -6,7 +6,7 @@
  * Side Public License, v 1.
  */
 
-import { loggingSystemMock, elasticsearchServiceMock } from '@kbn/core/server/mocks';
+import { loggingSystemMock } from '@kbn/core/server/mocks';
 import {
   Collector,
   createCollectorFetchContextMock,
@@ -52,7 +52,8 @@ describe('kibana_usage', () => {
 });
 
 describe('getKibanaSavedObjectCounts', () => {
-  const esClient = elasticsearchServiceMock.createClusterClient().asInternalUser;
+  const fetchContextMock = createCollectorFetchContextMock();
+  const soClient = fetchContextMock.soClient;
 
   test('Get all the saved objects equal to 0 because no results were found', async () => {
     getSavedObjectsCountsMock.mockResolvedValueOnce({
@@ -61,7 +62,7 @@ describe('getKibanaSavedObjectCounts', () => {
       non_expected_types: [],
       others: 0,
     });
-    const results = await getKibanaSavedObjectCounts(esClient, '.kibana');
+    const results = await getKibanaSavedObjectCounts(soClient);
     expect(results).toStrictEqual({
       dashboard: { total: 0 },
       visualization: { total: 0 },
@@ -83,7 +84,7 @@ describe('getKibanaSavedObjectCounts', () => {
       others: 0,
     });
 
-    const results = await getKibanaSavedObjectCounts(esClient, '.kibana');
+    const results = await getKibanaSavedObjectCounts(soClient);
     expect(results).toStrictEqual({
       dashboard: { total: 1 },
       visualization: { total: 0 },
@@ -93,8 +94,7 @@ describe('getKibanaSavedObjectCounts', () => {
     });
 
     expect(getSavedObjectsCountsMock).toHaveBeenCalledWith(
-      esClient,
-      '.kibana',
+      soClient,
       ['dashboard', 'visualization', 'search', 'index-pattern', 'graph-workspace'],
       true
     );

src/plugins/kibana_usage_collection/server/collectors/saved_objects_counts/kibana_usage_collector.ts

@@ -6,9 +6,9 @@
  * Side Public License, v 1.
  */
 
-import type { ElasticsearchClient } from '@kbn/core/server';
 import type { UsageCollectionSetup } from '@kbn/usage-collection-plugin/server';
 import { snakeCase } from 'lodash';
+import { SavedObjectsClientContract } from '@kbn/core/server';
 import { getSavedObjectsCounts } from './get_saved_object_counts';
 
 interface KibanaSavedObjectCounts {
@@ -26,10 +26,9 @@ interface KibanaUsage extends KibanaSavedObjectCounts {
 const TYPES = ['dashboard', 'visualization', 'search', 'index-pattern', 'graph-workspace'];
 
 export async function getKibanaSavedObjectCounts(
-  esClient: ElasticsearchClient,
-  kibanaIndex: string
+  soClient: SavedObjectsClientContract
 ): Promise<KibanaSavedObjectCounts> {
-  const { per_type: buckets } = await getSavedObjectsCounts(esClient, kibanaIndex, TYPES, true);
+  const { per_type: buckets } = await getSavedObjectsCounts(soClient, TYPES, true);
 
   const allZeros = Object.fromEntries(
     TYPES.map((type) => [snakeCase(type), { total: 0 }])
@@ -80,10 +79,10 @@ export function registerKibanaUsageCollector(
           },
         },
       },
-      async fetch({ esClient }) {
+      async fetch({ soClient }) {
         return {
           index: kibanaIndex,
-          ...(await getKibanaSavedObjectCounts(esClient, kibanaIndex)),
+          ...(await getKibanaSavedObjectCounts(soClient)),
         };
       },
     })

src/plugins/kibana_usage_collection/server/collectors/saved_objects_counts/saved_objects_count_collector.test.ts

@@ -17,10 +17,8 @@ describe('saved_objects_count_collector', () => {
   const usageCollectionMock = createUsageCollectionSetupMock();
   const fetchContextMock = createCollectorFetchContextMock();
 
-  const kibanaIndex = '.kibana-tests';
-
   beforeAll(() =>
-    registerSavedObjectsCountUsageCollector(usageCollectionMock, kibanaIndex, () =>
+    registerSavedObjectsCountUsageCollector(usageCollectionMock, () =>
       Promise.resolve(['type_one', 'type_two', 'type-three', 'type-four'])
     )
   );
@@ -81,8 +79,7 @@ describe('saved_objects_count_collector', () => {
     });
 
     expect(getSavedObjectsCountsMock).toHaveBeenCalledWith(
-      fetchContextMock.esClient,
-      kibanaIndex,
+      fetchContextMock.soClient,
       ['type_one', 'type_two', 'type-three', 'type-four'],
       false
     );