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