[Files] Adds bulk get method (elastic#155636)

## Summary Closes elastic#155369 ### Checklist Delete any items that are not applicable to this PR. - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios ### For maintainers - [x] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) --------- Co-authored-by: kibanamachine <[email protected]>
kibanamachine · Apr 27, 2023 · a0cd724 · a0cd724
1 parent 1146d03
commit a0cd724
Show file tree

Hide file tree

Showing 9 changed files with 251 additions and 7 deletions.
diff --git a/src/plugins/files/server/file_client/file_metadata_client/adapters/es_index.ts b/src/plugins/files/server/file_client/file_metadata_client/adapters/es_index.ts
@@ -12,20 +12,24 @@ import { Logger } from '@kbn/core/server';
 import { toElasticsearchQuery } from '@kbn/es-query';
 import { ElasticsearchClient } from '@kbn/core-elasticsearch-server';
 import { MappingProperty, SearchTotalHits } from '@elastic/elasticsearch/lib/api/types';
+import pLimit from 'p-limit';
+
 import type { FilesMetrics, FileMetadata, Pagination } from '../../../../common';
 import type { FindFileArgs } from '../../../file_service';
 import type {
   DeleteArg,
   FileDescriptor,
   FileMetadataClient,
   GetArg,
+  BulkGetArg,
   GetUsageMetricsArgs,
   UpdateArgs,
 } from '../file_metadata_client';
 import { filterArgsToKuery } from './query_filters';
 import { fileObjectType } from '../../../saved_objects/file';
 
 const filterArgsToESQuery = pipe(filterArgsToKuery, toElasticsearchQuery);
+const bulkGetConcurrency = pLimit(10);
 
 const fileMappings: MappingProperty = {
   dynamic: false,
@@ -120,6 +124,22 @@ export class EsIndexFilesMetadataClient<M = unknown> implements FileMetadataClie
     };
   }
 
+  async bulkGet(arg: { ids: string[]; throwIfNotFound?: true }): Promise<FileDescriptor[]>;
+  async bulkGet({ ids, throwIfNotFound }: BulkGetArg): Promise<Array<FileDescriptor | null>> {
+    const promises = ids.map((id) =>
+      bulkGetConcurrency(() =>
+        this.get({ id }).catch((e) => {
+          if (throwIfNotFound) {
+            throw e;
+          }
+          return null;
+        })
+      )
+    );
+    const result = await Promise.all(promises);
+    return result;
+  }
+
   async delete({ id }: DeleteArg): Promise<void> {
     await this.esClient.delete({ index: this.index, id });
   }

diff --git a/src/plugins/files/server/file_client/file_metadata_client/adapters/saved_objects.ts b/src/plugins/files/server/file_client/file_metadata_client/adapters/saved_objects.ts
@@ -21,6 +21,8 @@ import type { FileMetadata, FilesMetrics, FileStatus } from '../../../../common/
 import type {
   FileMetadataClient,
   UpdateArgs,
+  GetArg,
+  BulkGetArg,
   FileDescriptor,
   GetUsageMetricsArgs,
 } from '../file_metadata_client';
@@ -54,14 +56,33 @@ export class SavedObjectsFileMetadataClient implements FileMetadataClient {
       metadata: result.attributes as FileDescriptor['metadata'],
     };
   }
-  async get({ id }: { id: string }): Promise<FileDescriptor> {
+
+  async get({ id }: GetArg): Promise<FileDescriptor> {
     const result = await this.soClient.get(this.soType, id);
     return {
       id: result.id,
       metadata: result.attributes as FileDescriptor['metadata'],
     };
   }
 
+  async bulkGet(arg: { ids: string[]; throwIfNotFound?: true }): Promise<FileDescriptor[]>;
+  async bulkGet({ ids, throwIfNotFound }: BulkGetArg): Promise<Array<FileDescriptor | null>> {
+    const result = await this.soClient.bulkGet(ids.map((id) => ({ id, type: this.soType })));
+    return result.saved_objects.map((so) => {
+      if (so.error) {
+        if (throwIfNotFound) {
+          throw new Error(`File [${so.id}] not found`);
+        }
+        return null;
+      }
+
+      return {
+        id: so.id,
+        metadata: so.attributes as FileDescriptor['metadata'],
+      };
+    });
+  }
+
   async find({ page, perPage, ...filterArgs }: FindFileArgs = {}): Promise<{
     total: number;
     files: Array<FileDescriptor<unknown>>;

diff --git a/src/plugins/files/server/file_client/file_metadata_client/file_metadata_client.ts b/src/plugins/files/server/file_client/file_metadata_client/file_metadata_client.ts
@@ -63,6 +63,21 @@ export interface GetArg {
   id: string;
 }
 
+/**
+ * Bulk get files
+ */
+export interface BulkGetArg {
+  /**
+   * Unique IDs of file metadata
+   */
+  ids: string[];
+  /**
+   * Flag to indicate if an Error is thrown if any of the file id is not found. If set to `false`, "null" will be returned.
+   * @default true
+   */
+  throwIfNotFound?: boolean;
+}
+
 export interface DeleteArg {
   /**
    * Unique ID of file metadata to delete
@@ -98,6 +113,16 @@ export interface FileMetadataClient {
    */
   get(arg: GetArg): Promise<FileDescriptor>;
 
+  /**
+   * Bulk get file metadata
+   *
+   * @param arg - Arguments to bulk retrieve file metadata
+   */
+  bulkGet(arg: { ids: string[]; throwIfNotFound?: true }): Promise<FileDescriptor[]>;
+  bulkGet(
+    arg: BulkGetArg | { ids: string[]; throwIfNotFound: false }
+  ): Promise<Array<FileDescriptor | null>>;
+
   /**
    * The file metadata to update
    *

diff --git a/src/plugins/files/server/file_service/file_action_types.ts b/src/plugins/files/server/file_service/file_action_types.ts
@@ -82,6 +82,26 @@ export interface GetByIdArgs {
   id: string;
 }
 
+/**
+ * Arguments to bulk get multiple files by their IDs.
+ */
+export interface BulkGetByIdArgs {
+  /**
+   * File IDs.
+   */
+  ids: string[];
+  /**
+   * Flag to indicate if an Error is thrown if any of the file id is not found. If set to `false`, "null" will be returned.
+   * @default true
+   */
+  throwIfNotFound?: boolean;
+  /**
+   * Format of the response, either a list of File[] (sorted by id passed) or a map `{[fileId: string]: File}`
+   * @default "array"
+   */
+  format?: 'array' | 'map';
+}
+
 /**
  * Arguments to filter for files.
  *

diff --git a/src/plugins/files/server/file_service/file_service.ts b/src/plugins/files/server/file_service/file_service.ts
@@ -14,6 +14,7 @@ import type {
   DeleteFileArgs,
   BulkDeleteFilesArgs,
   GetByIdArgs,
+  BulkGetByIdArgs,
   FindFileArgs,
 } from './file_action_types';
 
@@ -58,6 +59,23 @@ export interface FileServiceStart {
    */
   getById<M>(args: GetByIdArgs): Promise<File<M>>;
 
+  /**
+   * Bulk get files by IDs. Will throw if any of the files fail to load (set `throwIfNotFound` to `false` to not throw and return `null` instead)
+   *
+   * @param args - bulk get files by IDs args
+   */
+  bulkGetById<M>(args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound?: true }): Promise<File[]>;
+  bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound?: true; format: 'map' }
+  ): Promise<{ [id: string]: File }>;
+  bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound: false }
+  ): Promise<Array<File | null>>;
+  bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound: false; format: 'map' }
+  ): Promise<{ [id: string]: File | null }>;
+  bulkGetById<M>(args: BulkGetByIdArgs): Promise<Array<File<M> | { [id: string]: File | null }>>;
+
   /**
    * Find files given a set of parameters.
    *

diff --git a/src/plugins/files/server/file_service/file_service_factory.ts b/src/plugins/files/server/file_service/file_service_factory.ts
@@ -20,7 +20,13 @@ import { fileObjectType, fileShareObjectType, hiddenTypes } from '../saved_objec
 import { BlobStorageService } from '../blob_storage_service';
 import { FileClientImpl } from '../file_client/file_client';
 import { InternalFileShareService } from '../file_share_service';
-import { CreateFileArgs, FindFileArgs, GetByIdArgs, UpdateFileArgs } from './file_action_types';
+import {
+  CreateFileArgs,
+  FindFileArgs,
+  GetByIdArgs,
+  BulkGetByIdArgs,
+  UpdateFileArgs,
+} from './file_action_types';
 import { InternalFileService } from './internal_file_service';
 import { FileServiceStart } from './file_service';
 import { FileKindsRegistry } from '../../common/file_kinds_registry';
@@ -86,6 +92,26 @@ export class FileServiceFactoryImpl implements FileServiceFactory {
       this.logger
     );
 
+    function bulkGetById<M>(
+      args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound?: true }
+    ): Promise<Array<File<M>>>;
+    function bulkGetById<M>(
+      args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound?: true; format: 'map' }
+    ): Promise<{ [id: string]: File<M> }>;
+    function bulkGetById<M>(
+      args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound: false }
+    ): Promise<Array<File<M> | null>>;
+    function bulkGetById<M>(
+      args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound: false; format: 'map' }
+    ): Promise<{ [id: string]: File<M> | null }>;
+    function bulkGetById<M>(
+      args: BulkGetByIdArgs
+    ): Promise<Array<File<M> | null> | { [id: string]: File<M> | null }> {
+      return internalFileService.bulkGetById<M>(
+        args as Parameters<InternalFileService['bulkGetById']>[0]
+      );
+    }
+
     return {
       async create<M>(args: CreateFileArgs<M>) {
         return internalFileService.createFile(args) as Promise<File<M>>;
@@ -102,6 +128,7 @@ export class FileServiceFactoryImpl implements FileServiceFactory {
       async getById<M>(args: GetByIdArgs) {
         return internalFileService.getById(args) as Promise<File<M>>;
       },
+      bulkGetById,
       async find<M>(args: FindFileArgs) {
         return internalFileService.findFilesJSON(args) as Promise<{
           files: Array<FileJSON<M>>;

diff --git a/src/plugins/files/server/file_service/internal_file_service.ts b/src/plugins/files/server/file_service/internal_file_service.ts
@@ -24,6 +24,7 @@ import type {
   BulkDeleteFilesArgs,
   FindFileArgs,
   GetByIdArgs,
+  BulkGetByIdArgs,
 } from './file_action_types';
 import { createFileClient, FileClientImpl } from '../file_client/file_client';
 
@@ -93,10 +94,71 @@ export class InternalFileService {
     }
   }
 
+  private async bulkGet<M>(
+    ids: string[],
+    {
+      throwIfNotFound = true,
+      format = 'array',
+    }: { throwIfNotFound?: boolean; format?: BulkGetByIdArgs['format'] } = {}
+  ): Promise<Array<IFile<M> | null> | { [id: string]: IFile<M> | null }> {
+    try {
+      const metadatas = await this.metadataClient.bulkGet({ ids, throwIfNotFound: false });
+      const result = metadatas.map((fileMetadata, i) => {
+        const notFound = !fileMetadata || !fileMetadata.metadata;
+        const deleted = fileMetadata?.metadata?.Status === 'DELETED';
+
+        if (notFound || deleted) {
+          if (!throwIfNotFound) {
+            return null;
+          }
+          throw new FileNotFoundError(
+            deleted ? 'File has been deleted' : `File [${fileMetadata?.id ?? ids[i]}] not found`
+          );
+        }
+
+        const { id, metadata } = fileMetadata;
+        return this.toFile<M>(id, metadata as FileMetadata<M>, metadata.FileKind);
+      });
+
+      return format === 'array'
+        ? result
+        : ids.reduce<{ [id: string]: IFile<M> | null }>(
+            (acc, id, i) => ({
+              ...acc,
+              [id]: result[i],
+            }),
+            {}
+          );
+    } catch (e) {
+      this.logger.error(`Could not retrieve files: ${e}`);
+      throw e;
+    }
+  }
+
   public async getById({ id }: GetByIdArgs): Promise<IFile> {
     return await this.get(id);
   }
 
+  public async bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound?: true }
+  ): Promise<Array<IFile<M>>>;
+  public async bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound?: true; format: 'map' }
+  ): Promise<{ [id: string]: IFile<M> }>;
+  public async bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound: false }
+  ): Promise<Array<IFile<M> | null>>;
+  public async bulkGetById<M>(
+    args: Pick<BulkGetByIdArgs, 'ids'> & { throwIfNotFound: false; format: 'map' }
+  ): Promise<{ [id: string]: IFile<M> | null }>;
+  public async bulkGetById<M>({
+    ids,
+    throwIfNotFound,
+    format,
+  }: BulkGetByIdArgs): Promise<Array<IFile<M> | null> | { [id: string]: IFile<M> | null }> {
+    return await this.bulkGet<M>(ids, { throwIfNotFound, format });
+  }
+
   public getFileKind(id: string): FileKind {
     return this.fileKindRegistry.get(id);
   }
@@ -122,15 +184,15 @@ export class InternalFileService {
     return this.get(fileId);
   }
 
-  private toFile(
+  private toFile<M>(
     id: string,
-    fileMetadata: FileMetadata,
+    fileMetadata: FileMetadata<M>,
     fileKind: string,
     fileClient?: FileClientImpl
-  ): IFile {
-    return new File(
+  ): IFile<M> {
+    return new File<M>(
       id,
-      toJSON(id, fileMetadata),
+      toJSON<M>(id, fileMetadata),
       fileClient ?? this.createFileClient(fileKind),
       this.logger.get(`file-${id}`)
     );

diff --git a/src/plugins/files/server/integration_tests/file_service.test.ts b/src/plugins/files/server/integration_tests/file_service.test.ts
@@ -140,6 +140,56 @@ describe('FileService', () => {
     expect(myFile?.id).toMatch(id);
   });
 
+  it('retrieves a file using the bulk method', async () => {
+    const { id } = await createDisposableFile({ fileKind, name: 'test' });
+    const [myFile] = await fileService.bulkGetById({ ids: [id] });
+    expect(myFile?.id).toMatch(id);
+  });
+
+  it('retrieves multiple files using the bulk method', async () => {
+    const file1 = await createDisposableFile({ fileKind, name: 'test' });
+    const file2 = await createDisposableFile({ fileKind, name: 'test' });
+    const [myFile1, myFile2] = await fileService.bulkGetById({ ids: [file1.id, file2.id] });
+    expect(myFile1?.id).toMatch(file1.id);
+    expect(myFile2?.id).toMatch(file2.id);
+  });
+
+  it('throws if one of the file does not exists', async () => {
+    const file1 = await createDisposableFile({ fileKind, name: 'test' });
+    const unknownID = 'foo';
+
+    expect(async () => {
+      await fileService.bulkGetById({ ids: [file1.id, unknownID] });
+    }).rejects.toThrowError(`File [${unknownID}] not found`);
+  });
+
+  it('does not throw if one of the file does not exists', async () => {
+    const file1 = await createDisposableFile({ fileKind, name: 'test' });
+    const unknownID = 'foo';
+
+    const [myFile1, myFile2] = await fileService.bulkGetById({
+      ids: [file1.id, unknownID],
+      throwIfNotFound: false,
+    });
+
+    expect(myFile1?.id).toBe(file1?.id);
+    expect(myFile2).toBe(null);
+  });
+
+  it('returns the files under a map of id/File', async () => {
+    const file1 = await createDisposableFile({ fileKind, name: 'test' });
+    const unknownID = 'foo';
+
+    const myFiles = await fileService.bulkGetById({
+      ids: [file1.id, unknownID],
+      throwIfNotFound: false,
+      format: 'map',
+    });
+
+    expect(myFiles[file1?.id]?.id).toBe(file1?.id);
+    expect(myFiles[unknownID]).toBe(null);
+  });
+
   it('lists files', async () => {
     await Promise.all([
       createDisposableFile({ fileKind, name: 'test-1' }),