| 1 | import type { Injectable } from '@nestjs/common/interfaces' |
| 2 | import type { InstanceWrapper } from '@nestjs/core/injector/instance-wrapper' |
| 3 | import { AsyncLocalStorage } from 'node:async_hooks' |
| 4 | import { Injectable as InjectableDec, Logger, OnModuleInit } from '@nestjs/common' |
| 5 | import { MetadataScanner, ModulesContainer } from '@nestjs/core' |
| 6 | import { InjectConnection } from '@nestjs/mongoose' |
| 7 | import { TransactionOptions } from 'mongodb' |
| 8 | import { Connection } from 'mongoose' |
| 9 | import { DB_CONNECTION_NAME } from './common' |
| 10 | import { TRANSACTIONAL_METADATA } from './decorators/transactional.decorator' |
| 11 | |
| 12 | interface TransactionContext { |
| 13 | inTransaction: boolean |
| 14 | } |
| 15 | |
| 16 | @InjectableDec() |
| 17 | export class TransactionalInjector implements OnModuleInit { |
| 18 | private readonly logger = new Logger(TransactionalInjector.name) |
| 19 | private readonly metadataScanner: MetadataScanner = new MetadataScanner() |
| 20 | private readonly transactionContext = new AsyncLocalStorage<TransactionContext>() |
| 21 | |
| 22 | constructor( |
| 23 | private readonly modulesContainer: ModulesContainer, |
| 24 | @InjectConnection(DB_CONNECTION_NAME) private readonly connection: Connection, |
| 25 | ) {} |
| 26 | |
| 27 | async onModuleInit() { |
| 28 | for (const provider of this.getProviders()) { |
| 29 | this.injectToProvider(provider) |
| 30 | } |
| 31 | } |
| 32 | |
| 33 | private* getProviders(): Generator<InstanceWrapper<Injectable>> { |
| 34 | for (const module of this.modulesContainer.values()) { |
| 35 | for (const provider of module.providers.values()) { |
| 36 | if (provider && provider.metatype?.prototype) { |
| 37 | yield provider as InstanceWrapper<Injectable> |
| 38 | } |
| 39 | } |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | private injectToProvider(wrapper: InstanceWrapper<Injectable>): void { |
| 44 | const { metatype } = wrapper |
| 45 | if (!metatype) |
| 46 | return |
| 47 | |
| 48 | const prototype = metatype.prototype |
| 49 | const methodNames = this.metadataScanner.getAllMethodNames(prototype) |
| 50 | |
| 51 | for (const methodName of methodNames) { |
| 52 | const method = prototype[methodName] |
| 53 | if (this.isDecorated(method)) { |
| 54 | const options = this.getDecoratorOptions(method) |
| 55 | const wrappedMethod = this.wrapMethod(method, methodName, prototype.constructor.name, options) |
| 56 | this.reDecorate(method, wrappedMethod) |
| 57 | prototype[methodName] = wrappedMethod |
| 58 | this.logger.log(`Injected transaction to ${prototype.constructor.name}.${methodName}`) |
| 59 | } |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | private isDecorated(target: object): boolean { |
| 64 | return Reflect.hasMetadata(TRANSACTIONAL_METADATA, target) |
| 65 | } |
| 66 | |
| 67 | private getDecoratorOptions(target: object): TransactionOptions { |
| 68 | return Reflect.getMetadata(TRANSACTIONAL_METADATA, target) |
| 69 | } |
| 70 | |
| 71 | private reDecorate(source: object, destination: object): void { |
| 72 | const keys = Reflect.getMetadataKeys(source) |
| 73 | for (const key of keys) { |
| 74 | const meta = Reflect.getMetadata(key, source) |
| 75 | Reflect.defineMetadata(key, meta, destination) |
| 76 | } |
| 77 | } |
| 78 | |
| 79 | private wrapMethod( |
| 80 | originalMethod: (...args: unknown[]) => unknown, |
| 81 | methodName: string, |
| 82 | className: string, |
| 83 | options: TransactionOptions, |
| 84 | ): (...args: unknown[]) => unknown { |
| 85 | return new Proxy(originalMethod, { |
| 86 | apply: async (target, thisArg, args: unknown[]) => { |
| 87 | const fullMethodName = `${className}.${methodName}` |
| 88 | |
| 89 | const currentContext = this.transactionContext.getStore() |
| 90 | if (currentContext?.inTransaction) { |
| 91 | this.logger.debug(`Skipping nested transaction for ${fullMethodName}`) |
| 92 | return Reflect.apply(target, thisArg, args) |
| 93 | } |
| 94 | |
| 95 | this.logger.debug(`Executing transactional method: ${fullMethodName}`) |
| 96 | |
| 97 | return this.transactionContext.run( |
| 98 | { inTransaction: true }, |
| 99 | () => this.connection.transaction( |
| 100 | () => Reflect.apply(target, thisArg, args) as Promise<unknown>, |
| 101 | options, |
| 102 | ), |
| 103 | ) |
| 104 | }, |
| 105 | }) |
| 106 | } |
| 107 | } |
| 108 |