[Prisma plugin] Update Prisma client from context during subscription

[jgnieuwhof]

First off - thank you for the great work! We've been using Pothos on this project for a couple of years now, and have been enjoying it.

Background

This project I'm on uses RLS policies and Prisma transactions to control what data the resolvers (... and user) have access to. An Envelop plugin wraps query execution in a user-scoped PrismaTransaction that is placed in context for the resolvers and Pothos.

The flow for subscriptions is similar, but modified so that context is updated with a new PrismaTransaction on every iteration of the subscription.

The Issue

This works fine for fields with custom resolvers, but fails with a 'transaction already committed' error for fields resolved by Prisma plugin:

// builder.ts
import SchemaBuilder from "@pothos/core";

export const builder = new SchemaBuilder({
  plugins: [
    PrismaPlugin,
    SmartSubscriptionsPlugin
  ],
  prisma: {
    // The 'prisma' here is actually a `PrismaTransaction`, that is set:
    // - before execution of each query or mutation
    // - before each iteration of a subscription
    client({ prisma }) {
      return prisma as unknown as PrismaClient;
    },
    dmmf: Prisma.dmmf,
  },
});

// Resolving user.tokens fails with 'transaction already committed' error on subsequent subscription events
const User = builder.prismaNode("User", {
  id: { field: "id" },
  smartSubscriptions: true,
  fields: (t) => ({
    tokens: t.relatedConnection("tokens", {
      cursor: "id",
      totalCount: true,
    }),
  }),
});

// Resolving user.tokens succeeds on all subsequent subscription events
const User = builder.prismaNode("User", {
  id: { field: "id" },
  fields: (t) => ({
    tokens: t.relatedConnection("tokens", {
      cursor: "id",
      totalCount: true,
      resolve(query, { id }, { prisma }) {
        return prisma.token.findMany({
          ...query,
          where: {
            userId: id,
          }
        })
      }
    }),
  }),
});

I believe this is because the plugin's resolvers are using a cached version of the PrismaClient, which in this case would be the transaction that is placed in context for the first iteration of a subscription:

• https://github.com/hayes/pothos/blob/main/packages/plugin-prisma/src/util/get-client.ts#L57
• https://github.com/hayes/pothos/blob/main/packages/plugin-prisma/src/model-loader.ts#L56

So.... Is there a way to invalidate or disable the Prisma client context cache? Or a way to ask the plugin to re-initialize its model loaders / delegates after a new transaction has been placed in context?

I've played around with passing in a Proxy object in place of the PrismaClient, but that felt rather hacky and I thought I'd reach out!

[jgnieuwhof]

Here's the Envelop plugin that wraps query execution / subscription iterations in case some more context is helpful.

// envelop.ts
import { Plugin } from '@envelop/core';

export function prismaTransactionPlugin(): Plugin {
  return {
    // Wrap query resolution in a transaction with the user's id
    // set. The Postgres RLS policies use the user_id to
    // determine what records the user is allowed to read/write.
    onExecute({ executeFn, setExecuteFn, extendContext }) {
      setExecuteFn(async function executor(args) {
        const { token, prisma } = args.contextValue;
        // Start a transaction that will span the execution
        return prisma.$transaction(async (transaction: PrismaTransaction) => {
          // Set the user's ID for this transaction
          await transaction.$executeRaw`
            SELECT set_config('jwt.claims.user_id', ${token.sub}, true)
          `;
          // Replace the PrismaClient in context with this user-scoped transaction
          extendContext({ prisma: transaction });
          return await executeFn(args);
        });
      });
    },
    // Similar idea as above - wrap each iteration of the subscription
    // in its own user-scoped transaction.
    onSubscribe({ subscribeFn, setSubscribeFn, extendContext }) {
      setSubscribeFn(async (args) => {
        const subscriber = subscribeFn(args);
        return {
          [Symbol.asyncIterator]() {
            return {
              async next() {
                const { token, prisma } = args.contextValue;
                // Start a transaction in which to run the next event, with
                // an extended timeout since there is no guarantee on when
                // the next event will arrive.
                const result = await prisma.$transaction(
                  async (transaction: PrismaTransaction) => {
                    await transaction.$executeRaw`
                      SELECT set_config('jwt.claims.user_id', ${token.sub}, true)
                    `;
                    extendContext({ prisma: transaction });
                    return await subscriber.next(args);
                  },
                  { timeout: Number.MAX_SAFE_INTEGER }
                );
                // Reset 'prisma' context to its original PrismaClient for
                // the next iteration.
                extendContext({ prisma });
                return result;
              },
              async return() {
                return await subscriber.return();
              },
              async throw(error: unknown) {
                return await subscriber.throw(error);
              },
            };
          },
        };
      });
    },
  };
}

[hayes]

I think the easiest option here is to invalidate the whole context cache by re-initializing it as described here: https://pothos-graphql.dev/docs/guide/context#initialize-context-cache

I haven't tried to replicate your set up, so there might be other complications, but give that a try and see if it works

[hayes]

I think re-initializing the whole context cache will probably break the smart subscriptions plugin.

Context caches have a delete method, so I think if we export the prismaContextCache from the Prisma plugin, you could delete the cache for the current context object when you create the transaction

[hayes]

Published a new version, if you update, you should be able to import prismaClientCache from the Prisma plugin, and then call prismaClientCache.delete(context)

[jgnieuwhof]

Thanks for such a fast response!

Managed to get this working with the following changes:

• Invalidate the prisma client cache via prismaClientCache(builder)(context)
• Patch the model-loader module to re-obtain its PrismaDelegate before running a query

What are your thoughts on the model-loader change? It does incur a slight runtime cost per query, however that could be minimized by holding the delegate name in state, rather than a reference to the delegate itself.

// @pothos/plugin-prisma/src/model-loader.ts
export class ModelLoader {
  ...

  async initLoad(...) {
    ...
    this.tick.then(() => {
      this.staged.delete(entry);
      // added this line to refresh the delegate prior to each query, in case the prisma client cache
      // has been modified.
      this.delegate = getDelegateFromModel(getClient(this.builder, this.context), this.modelName);
      for (const [model, { resolve, reject }] of entry.models) {
        ...
      }
    });
  }
}

// envelop.ts
import { Plugin } from '@envelop/core';
import { prismaClientCache } from "@pothos/plugin-prisma";

import { builder } from "./schema/builder";

export function prismaTransactionPlugin(): Plugin {
  return {
    onExecute({ executeFn, setExecuteFn, extendContext }) {
      ...
    },
    onSubscribe({ subscribeFn, setSubscribeFn, extendContext }) {
      setSubscribeFn(async (args) => {
        const subscriber = subscribeFn(args);
        return {
          [Symbol.asyncIterator]() {
            return {
              async next() {
                const { token, prisma } = args.contextValue;
                const result = await prisma.$transaction(
                  async (transaction: PrismaTransaction) => {
                    await transaction.$executeRaw`
                      SELECT set_config('jwt.claims.user_id', ${token.sub}, true)
                    `;
                    // Delete the pothos/prisma client cache so that it is re-obtained
                    // before resolving next
                    prismaClientCache(builder as any).delete(args.contextValue);
                    extendContext({ prisma: transaction });
                    return await subscriber.next(args);
                  },
                  { timeout: Number.MAX_SAFE_INTEGER }
                );
                extendContext({ prisma });
                return result;
              },
              async return() {
                return await subscriber.return();
              },
              async throw(error: unknown) {
                return await subscriber.throw(error);
              },
            };
          },
        };
      });
    },
  };
}

[hayes]

This seems reasonable, I think a little extra overhead here is okay, since it is in the process of kicking off a much more expensive db query.

I'm on vacation right now, but I'll try to take a look at this in the next couple days

[hayes]

If you wanted to submit a pr, I would probably just remove the delegate prop and just store it in a variable in initLoad

It probably doesn't matter, but I'd also think through of the delegate should be loaded when initLoad is called, or just before trigging the query.

I'm not 100% which is more correct, my initial instinct is that the client should be locked in when initLoad is called.

[jgnieuwhof]

Sounds good! PR here to get us started:

#1331

I'm off here for a couple days as well, no rush here. Enjoy your vacation!