[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