[Saved Object Service] Adds Repository Factory Provider (#4149)

* Adds Repository Factory Provider

Signed-off-by: Bandini Bhopi <bandinib@amazon.com>
(cherry picked from commit 5c5de03)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

# Conflicts:
#	CHANGELOG.md

src/core/server/legacy/legacy_service.ts

@@ -275,6 +275,7 @@ export class LegacyService implements CoreService {
         addClientWrapper: setupDeps.core.savedObjects.addClientWrapper,
         registerType: setupDeps.core.savedObjects.registerType,
         getImportExportObjectLimit: setupDeps.core.savedObjects.getImportExportObjectLimit,
+        setRepositoryFactoryProvider: setupDeps.core.savedObjects.setRepositoryFactoryProvider,
       },
       status: {
         isStatusPageAnonymous: setupDeps.core.status.isStatusPageAnonymous,

src/core/server/plugins/plugin_context.ts

@@ -204,6 +204,7 @@ export function createPluginSetupContext<TPlugin, TPluginDependencies>(
       addClientWrapper: deps.savedObjects.addClientWrapper,
       registerType: deps.savedObjects.registerType,
       getImportExportObjectLimit: deps.savedObjects.getImportExportObjectLimit,
+      setRepositoryFactoryProvider: deps.savedObjects.setRepositoryFactoryProvider,
     },
     status: {
       core$: deps.status.core$,

src/core/server/saved_objects/saved_objects_service.mock.ts

@@ -79,6 +79,7 @@ const createSetupContractMock = () => {
     addClientWrapper: jest.fn(),
     registerType: jest.fn(),
     getImportExportObjectLimit: jest.fn(),
+    setRepositoryFactoryProvider: jest.fn(),
   };
 
   setupContract.getImportExportObjectLimit.mockReturnValue(100);

src/core/server/saved_objects/saved_objects_service.test.ts

@@ -41,14 +41,15 @@ import { errors as opensearchErrors } from '@opensearch-project/opensearch';
 import { SavedObjectsService } from './saved_objects_service';
 import { mockCoreContext } from '../core_context.mock';
 import { Env } from '../config';
-import { configServiceMock } from '../mocks';
+import { configServiceMock, savedObjectsRepositoryMock } from '../mocks';
 import { opensearchServiceMock } from '../opensearch/opensearch_service.mock';
 import { opensearchClientMock } from '../opensearch/client/mocks';
 import { httpServiceMock } from '../http/http_service.mock';
 import { httpServerMock } from '../http/http_server.mocks';
 import { SavedObjectsClientFactoryProvider } from './service/lib';
 import { NodesVersionCompatibility } from '../opensearch/version_check/ensure_opensearch_version';
 import { SavedObjectsRepository } from './service/lib/repository';
+import { SavedObjectRepositoryFactoryProvider } from './service/lib/scoped_client_provider';
 
 jest.mock('./service/lib/repository');
 
@@ -169,6 +170,27 @@ describe('SavedObjectsService', () => {
         expect(typeRegistryInstanceMock.registerType).toHaveBeenCalledWith(type);
       });
     });
+
+    describe('#setRepositoryFactoryProvider', () => {
+      it('throws error if a repository is already registered', async () => {
+        const coreContext = createCoreContext();
+        const soService = new SavedObjectsService(coreContext);
+        const setup = await soService.setup(createSetupDeps());
+
+        const firstRepository: SavedObjectRepositoryFactoryProvider = () =>
+          savedObjectsRepositoryMock.create();
+        const secondRepository: SavedObjectRepositoryFactoryProvider = () =>
+          savedObjectsRepositoryMock.create();
+
+        setup.setRepositoryFactoryProvider(firstRepository);
+
+        expect(() => {
+          setup.setRepositoryFactoryProvider(secondRepository);
+        }).toThrowErrorMatchingInlineSnapshot(
+          `"custom repository factory is already set, and can only be set once"`
+        );
+      });
+    });
   });
 
   describe('#start()', () => {
@@ -281,6 +303,15 @@ describe('SavedObjectsService', () => {
       }).toThrowErrorMatchingInlineSnapshot(
         `"cannot call \`registerType\` after service startup."`
       );
+
+      const customRpository: SavedObjectRepositoryFactoryProvider = () =>
+        savedObjectsRepositoryMock.create();
+
+      expect(() => {
+        setup.setRepositoryFactoryProvider(customRpository);
+      }).toThrowErrorMatchingInlineSnapshot(
+        '"cannot call `setRepositoryFactoryProvider` after service startup."'
+      );
     });
 
     describe('#getTypeRegistry', () => {
@@ -368,6 +399,36 @@ describe('SavedObjectsService', () => {
 
         expect(includedHiddenTypes).toEqual(['someHiddenType']);
       });
+
+      it('Should not create SavedObjectsRepository when custom repository is registered ', async () => {
+        const coreContext = createCoreContext({ skipMigration: false });
+        const soService = new SavedObjectsService(coreContext);
+        const coreSetup = createSetupDeps();
+        const setup = await soService.setup(coreSetup);
+
+        const customRpository: SavedObjectRepositoryFactoryProvider = () =>
+          savedObjectsRepositoryMock.create();
+        setup.setRepositoryFactoryProvider(customRpository);
+
+        const coreStart = createStartDeps();
+        const { createInternalRepository } = await soService.start(coreStart);
+        createInternalRepository();
+
+        expect(SavedObjectsRepository.createRepository as jest.Mocked<any>).not.toHaveBeenCalled();
+      });
+
+      it('Should create SavedObjectsRepository when no custom repository is registered ', async () => {
+        const coreContext = createCoreContext({ skipMigration: false });
+        const soService = new SavedObjectsService(coreContext);
+        const coreSetup = createSetupDeps();
+        await soService.setup(coreSetup);
+
+        const coreStart = createStartDeps();
+        const { createInternalRepository } = await soService.start(coreStart);
+        createInternalRepository();
+
+        expect(SavedObjectsRepository.createRepository as jest.Mocked<any>).toHaveBeenCalled();
+      });
     });
   });
 });

src/core/server/saved_objects/saved_objects_service.ts

@@ -56,6 +56,7 @@ import { ISavedObjectsRepository, SavedObjectsRepository } from './service/lib/r
 import {
   SavedObjectsClientFactoryProvider,
   SavedObjectsClientWrapperFactory,
+  SavedObjectRepositoryFactoryProvider,
 } from './service/lib/scoped_client_provider';
 import { Logger } from '../logging';
 import { SavedObjectTypeRegistry, ISavedObjectTypeRegistry } from './saved_objects_type_registry';
@@ -166,6 +167,14 @@ export interface SavedObjectsServiceSetup {
    * Returns the maximum number of objects allowed for import or export operations.
    */
   getImportExportObjectLimit: () => number;
+
+  /**
+   * Set the default {@link SavedObjectRepositoryFactoryProvider | factory provider} for creating Saved Objects repository.
+   * Only one repository can be set, subsequent calls to this method will fail.
+   */
+  setRepositoryFactoryProvider: (
+    respositoryFactoryProvider: SavedObjectRepositoryFactoryProvider
+  ) => void;
 }
 
 /**
@@ -291,6 +300,8 @@ export class SavedObjectsService
   private typeRegistry = new SavedObjectTypeRegistry();
   private started = false;
 
+  private respositoryFactoryProvider?: SavedObjectRepositoryFactoryProvider;
+
   constructor(private readonly coreContext: CoreContext) {
     this.logger = coreContext.logger.get('savedobjects-service');
   }
@@ -348,6 +359,15 @@ export class SavedObjectsService
         this.typeRegistry.registerType(type);
       },
       getImportExportObjectLimit: () => this.config!.maxImportExportSize,
+      setRepositoryFactoryProvider: (repositoryProvider) => {
+        if (this.started) {
+          throw new Error('cannot call `setRepositoryFactoryProvider` after service startup.');
+        }
+        if (this.respositoryFactoryProvider) {
+          throw new Error('custom repository factory is already set, and can only be set once');
+        }
+        this.respositoryFactoryProvider = repositoryProvider;
+      },
     };
   }
 
@@ -422,13 +442,21 @@ export class SavedObjectsService
       opensearchClient: OpenSearchClient,
       includedHiddenTypes: string[] = []
     ) => {
-      return SavedObjectsRepository.createRepository(
-        migrator,
-        this.typeRegistry,
-        opensearchDashboardsConfig.index,
-        opensearchClient,
-        includedHiddenTypes
-      );
+      if (this.respositoryFactoryProvider) {
+        return this.respositoryFactoryProvider({
+          migrator,
+          typeRegistry: this.typeRegistry,
+          includedHiddenTypes,
+        });
+      } else {
+        return SavedObjectsRepository.createRepository(
+          migrator,
+          this.typeRegistry,
+          opensearchDashboardsConfig.index,
+          opensearchClient,
+          includedHiddenTypes
+        );
+      }
     };
 
     const repositoryFactory: SavedObjectsRepositoryFactory = {

src/core/server/saved_objects/service/lib/scoped_client_provider.ts

@@ -31,8 +31,31 @@
 import { PriorityCollection } from './priority_collection';
 import { SavedObjectsClientContract } from '../../types';
 import { SavedObjectsRepositoryFactory } from '../../saved_objects_service';
-import { ISavedObjectTypeRegistry } from '../../saved_objects_type_registry';
+import {
+  ISavedObjectTypeRegistry,
+  SavedObjectTypeRegistry,
+} from '../../saved_objects_type_registry';
 import { OpenSearchDashboardsRequest } from '../../../http';
+import { ISavedObjectsRepository } from './repository';
+import { IOpenSearchDashboardsMigrator } from '../../migrations';
+
+/**
+ * Options passed to each SavedObjectRepositoryFactoryProvider to aid in creating the repository instance.
+ * @public
+ */
+export interface SavedObjectsRepositoryOptions {
+  migrator: IOpenSearchDashboardsMigrator;
+  typeRegistry: SavedObjectTypeRegistry;
+  includedHiddenTypes: string[];
+}
+
+/**
+ * Provider to invoke to a factory function for creating ISavedObjectRepository {@link ISavedObjectRepository} instances.
+ * @public
+ */
+export type SavedObjectRepositoryFactoryProvider = (
+  options: SavedObjectsRepositoryOptions
+) => ISavedObjectsRepository;
 
 /**
  * Options passed to each SavedObjectsClientWrapperFactory to aid in creating the wrapper instance.