Design Question: Storing Prisma transaction client in AsyncLocalStorage

[fuvidani]

Hi @Papooch !

I am considering implementing decorator-based database transactions in our NestJs application using AsyncLocalStorage. I'm reaching out to you, because looking at your CLS implementation, you seem to have extensive experience with this.

I'd like to merge the concept from this issue for Prisma, but using AsyncLocalStorage instead of cls-hooked. The idea is to introduce a @Transaction decorator for service methods, that would store a transaction client in the storage and wrap the function s.t. the repositories within that call chain could retrieve it and execute their operations with it.

A given repository, when invoked, would try to first retrieve the transaction client from the store and use it if it's available. If not available, it would just use its own db-client (i.e. there is no composite transaction happening).

Do you think this is a sound approach? I'm not entirely sure about where the AsyncLocalStorage should be instantiated. Does it need to happen at bootstrap time or can it be spontaneous (when decorator is called)?

Any feedback would be much appreciated! 😊

PS.: I also considered to use your CLSService, but AFAIK injecting a custom service into a decorator would be challenging, so I'm not sure it'd be applicable, or would it? 🤔

[Papooch]

Hi there, it is absolutely doable! (It is basically what I have suggested before: prisma/prisma#5729 (comment), and you can use pure AsyncLocalStorage in practically the same way).

In my library, the ClsService is instantiated outside of DI and can be retrieved outside of it, if needed.

The AsyncLocalStorage object can be instantiated anywhere, but to enter the context, you need to wrap the method call with the .run(store, ()=>{...}) call. This can be done in a middleware or an enhancer, or by replacing the method implementation in the decorator (similar to what I'm aiming for here).

[fuvidani]

Thank you @Papooch for the quick and detailed response! 🤩
I still have some doubts on how to reliably access the context if it's not wrapped into a .run() call. Using a @Transaction decorator is clear to me I think - I'd just need to wrap the the whole method into a run and everything inside would have access to the context.

But now let's consider a repository implementation that would look like this:

export abstract class BaseRepository {
    constructor(
          private readonly prisma: PrismaService, 
          private readonly cls: ClsService,) {}

    protected get prisma(): PrismaClient | Prisma.TransactionClient {
        const transactionClient = this.cls.get('transactionClient');
        return transactionClient ?? this.prisma;
    }

}

And a concrete repository:

export class PhotoRepository extends BaseRepository {

    public async create(photo: Photo): Promise<Photo> {
        // Persisting the photo ...
        return Photo.fromPrismaToEntity(
           await this.prisma.create({ data: photo })
        );
    }

}

Now I fear that this is not that simple, because if I understood the concept right, then the invocation should've already been wrapped into a run in order for the the line const transactionClient = this.cls.get('transactionClient'); to work reliably. My goal is to have encapsulation and isolation s.t. if the repository is simply called without an explicit transaction decorator, then it falls back to its own Prisma instance. If my suspicion is right, then actually every invocation of the repository would require a run wrap somewhere, regardless of an existing transaction or not. Am I wrong? 😬

[Papooch]

Yes, that's right. As I said, you can do that inside the @Transactional decorator by replacing the decorated method's implementatioj with a wrapped version. You would also have to modify the class to have the Prisma service injected so you could also wrap it in a transaction.

Otherwise you'd have to do all this in an enhancer or mideware, but that wouldn't work with invocations that don't go through a controller (because enhancers aren't available in that case).

EDIT: Actually now that I re-read the question, you if you call cls.get outside of a run context, it should just return undefined, so it would still work ourside of the context. But yes, the transaction will only work if the context has been set up before entering the transaction.

[fuvidani]

Thanks for the clarification! I came up with the following rough decorator implementation (haven't tested it yet, it's just a PoC). What do you think of it? I'm particularly interested in the order of wrapping. Should I start the transaction first and then do a cls.run() or the other way around (like in the code snippet below)?

Another question: prisma.$transaction() returns a promise. Does cls.run() takes care of resolving the Promise or do I need to await it explicitly? 🤔

import { Inject } from '@nestjs/common';
import { PrismaService } from 'our-own-custom-lib';
import { ClsServiceManager } from 'nestjs-cls';

export function Transaction() {
  const injectPrisma = Inject(PrismaService);

  return (target: any, propertyKey: string, descriptor: PropertyDescriptor) => {
    if (this.prisma) {
      console.log('Object already contains a PrismaService dep');
    } else {
      // this is equivalent to have a constructor like constructor(yourservice: YourServiceClass)
      // note that this will injected to the instance, while your decorator runs for the class constructor
      injectPrisma(target, 'prisma');
    }

    // we use a ref here so we can type it
    const originalMethod = descriptor.value;
    descriptor.value = function (...args: any[]) {
      const cls = ClsServiceManager.getClsService();
      if (cls.get('TRANSACTION_CLIENT')) {
        // A transaction has already been started. We just call the function.
        return originalMethod.apply(this, [...args]);
      }
      cls.run(async () => {
        const prisma: PrismaService = this.prisma;
        return prisma.$transaction(async (transactionClient) => {
          cls.set('TRANSACTION_CLIENT', transactionClient);
          try {
            // We can now call the function that had been decorated with the Transaction decorator.
            const result = await originalMethod.apply(this, [...args]);
            return result;
          } finally {
            // We just finished working with our transaction: we remove it rom the current context.
            cls.set('TRANSACTION_CLIENT', null);
          }
        });
      });
    };
  };
}

[Papooch]

Withouh having run any code, what you have there seems reasonable.

One more check I would do is run cls.isActive() and don't wrap the call if we're already running inside a cls context, so we're not nesting contexts, in which case you'd have to copy the contents of the ClsStore to the new context if other code relied on it.

cls.run only wraps the call and does not perform any awaiting or error handling itself. The return type of the wrapped callback is the same as the original one.

[fuvidani]

Thank you @Papooch!! I'll try it out tomorrow, I'm very excited if this can work reliably. We could finally get rid of leaking Prisma outside of the data access layer.

[Papooch]

@fuvidani did you come to a satisfactory solution?

[fuvidani]

@fuvidani did you come to a satisfactory solution?

Hi @Papooch! I've been working on it as a side-task and it's in a pretty good state. I'm just concerned about the multi-module usage in Nestjs. Let's say we have module A and module B, they are both importing your CLS implementation via ClsModule.register({global: false, middleware: { mount: false },}). Now consider the situation where module A is also importing module B in a way that module A calls a service method of module B. Which CLS store is going to be used? Is it guaranteed that calls to ClsServiceManager.getClsService(); will target the right store?

To give you more context, I created a separate BaseRepository module, which practically provides a base repository class to wrap the retrieval of the Prisma instance:

import { Injectable } from '@nestjs/common';
import { ClsService } from 'nestjs-cls';
import { Prisma, PrismaClient } from '@prisma/client';
// some imports omitted 

@Injectable()
export class BaseRepository {
  constructor(
    private readonly prismaClient: PrismaService, // this service is just a wrapper around PrismaClient
    private readonly cls: ClsService
  ) {}

  protected get prisma(): PrismaClient | Prisma.TransactionClient {
    const maybeTransactionClient = this.cls.get(
      TRANSACTION_CLIENT_KEY
    ) as PrismaService;
    if (this.cls.isActive() && maybeTransactionClient) {
      logger.info('Using transaction client from ALS');
      return maybeTransactionClient;
    }
    logger.info('Using basic Prisma client');
    return this.prismaClient;
  }
}

• Now the idea is that other modules in the codebase can simply import BaseRepositoryModule and write their own repository implementations.
• There is also a @Transactional decorator I created separately, that also uses ClsServiceManager.getClsService(); to check if there is a transaction client stored. Its implementation is based on my previous comment.
• At the moment I have ClsModule imported in other modules like this:

@Module({
  controllers: [],
  imports: [
   //...
    ClsModule.register({
      global: false,
      middleware: { mount: false },
    }),
  ],
  providers: [/* providers*/],
  exports: [/* exports*/,
})
export class MyModuleA {}

Again, if there are 2 decorated methods in 2 different modules and I start a call chain from the first module, will the decorator in the second module get the same CLSService instance or the one from its own module? Do you understand what I'm trying to say? My fear is that in runtime the ClsServiceManager.getClsService() invocation in the decorator returns two different instances of the service, s.t. transactionality across two modules could not be guaranteed. Am I completely misunderstanding something?
If you could help me out here with the theory, I'd be very grateful. 🙇 🙏

[Papooch]

Since ClsServiceManager.getClsService() is a static method, it returns the same each time. Actually, it's the same instance as the injected ClsService that lives in the DI container. Even if they weren't, they share the same AsyncLocalStorage underneath.

Btw are you actually mounting the ClsMiddleware manually somewhere? If not, then it's redundant in the module registration.

[fuvidani]

Makes sense @Papooch, this is quite reassuring. Huge thank you @Papooch, really appreciate the support! 💪

[fuvidani]

Hi @Papooch! Do you know how this behaves in terms of concurrency? If the event loop switches between two asynchronous call chains, wouldn't ClsServiceManager.getClsService().get('transactionClient') return the same value to both chains? That would potentially mean that:

• the two chains would work with the same tx client, although each of them should have one for themselves
• the first chain to finish would remove the tx client from the CLSService causing the second chain being unable to retrieve the client (and falling back to basic Prisma, which breaks the whole transactionality).

What do you think about that?

[Papooch]

That's the nice thing about AsyncLocalStorage. Each execution chain gets an unique instance of the store and it gets propoagated through promise chains and callbacks (it even works with asynchronous event emitters and nestjs/cqrs).

The only case in which there is a chance that 2 requests would access the wame store is when you use the enterWith function to set up the context, but otherwise all code that is invoked from wtihin the run callback gets exclusive access to an unique instance of the store.