diff --git a/src/common/index.ts b/src/common/index.ts index c76f212..ec0dc34 100644 --- a/src/common/index.ts +++ b/src/common/index.ts @@ -1,5 +1,5 @@ import { createNamespace, getNamespace, Namespace } from 'cls-hooked'; -import { DataSource, EntityManager } from 'typeorm'; +import { DataSource, EntityManager, Repository } from 'typeorm'; import { NAMESPACE_NAME, TYPEORM_DATA_SOURCE_NAME, @@ -145,39 +145,51 @@ export const getTransactionalOptions = () => data.options; export const initializeTransactionalContext = (options?: Partial) => { setTransactionalOptions(options); - const originalGetRepository = EntityManager.prototype.getRepository; + const patchManager = (repositoryType: unknown) => { + Object.defineProperty(repositoryType, 'manager', { + get() { + return ( + getEntityManagerInContext( + this[TYPEORM_ENTITY_MANAGER_NAME].connection[ + TYPEORM_DATA_SOURCE_NAME + ] as DataSourceName, + ) || this[TYPEORM_ENTITY_MANAGER_NAME] + ); + }, + set(manager: EntityManager | undefined) { + this[TYPEORM_ENTITY_MANAGER_NAME] = manager; + }, + }); + }; - EntityManager.prototype.getRepository = function (...args: unknown[]) { - const repository = originalGetRepository.apply(this, args); - - if (!(TYPEORM_ENTITY_MANAGER_NAME in repository)) { - /** - * Store current manager - */ - repository[TYPEORM_ENTITY_MANAGER_NAME] = repository.manager; - - /** - * Patch repository object - */ - Object.defineProperty(repository, 'manager', { - get() { - return ( - getEntityManagerInContext( - this[TYPEORM_ENTITY_MANAGER_NAME].connection[ - TYPEORM_DATA_SOURCE_NAME - ] as DataSourceName, - ) || this[TYPEORM_ENTITY_MANAGER_NAME] - ); - }, - set(manager: EntityManager | undefined) { - this[TYPEORM_ENTITY_MANAGER_NAME] = manager; - }, - }); - } + const getRepository = (originalFn: (args: unknown) => unknown) => { + return function patchRepository(...args: unknown[]) { + const repository = originalFn.apply(this, args); + + if (!(TYPEORM_ENTITY_MANAGER_NAME in repository)) { + /** + * Store current manager + */ + repository[TYPEORM_ENTITY_MANAGER_NAME] = repository.manager; + + /** + * Patch repository object + */ + patchManager(repository); + } - return repository; + return repository; + }; }; + const originalGetRepository = EntityManager.prototype.getRepository; + const originalExtend = Repository.prototype.extend; + + EntityManager.prototype.getRepository = getRepository(originalGetRepository); + Repository.prototype.extend = getRepository(originalExtend); + + patchManager(Repository.prototype); + return createNamespace(NAMESPACE_NAME) || getNamespace(NAMESPACE_NAME); };