Transactions within Client Extension Example

[dennemark]

Question

( Prisma Client 5.14.0 )

Hi,

I have seen this topic too often - transactions within client extensions. It has been bothering me for RLS and now it is bothering me in development of my own client extension prisma-extension-casl
I have seen others like zenstack avoiding sequential transactions and only using interactive transactions. And this seems to be currently the only way to go.

Here is an example, where we check our queried results afterwards. And if it errors, it should revert the interactive transaction. There is also a RLS function. (I haven't tested it for the 'itx' case. Let me know if I should revise it)

Prisma.defineExtension((client) => {
        return client.$extends({
            query: {
                $allModels: {
                    async $allOperations<T>({ args, query, model, operation, ...rest }: { args: any, query: any, model: any, operation: any }) {

                        const setRLS = async (client: any)=>{
                            await client.$executeRaw('some rls logic...')
                        }

                        const checkResults = (result: any) => {
                            if(result === 'not-valid'){
                                throw new Error('This is just a an example check function...')
                            }
                            return result
                        }
                         // check if this query already runs within a transaction
                        if ((rest as any).__internalParams.transaction) {
                            const transaction = (rest as any).__internalParams.transaction
                            if (transaction.kind === 'itx') {
                                // we are in an interactive transaction and create a client for it
                                const transactionClient = (client as any)._createItxClient(transaction)
                                setRLS(transactionClient)
                                return transactionClient[model][operation](args).then(checkResults)
                            } else if (transaction.kind === 'batch') {
                                throw new Error('Sequential transactions are not supported in prisma client extensions.')
                            }
                        } else {
                            // if we are not in a transaction, we can run our own
                            return client.$transaction(async (tx) => {
                                setRLS(tx)
                                // we place our check results function in an interactive transaction
                                // so that our query is rolled back if check resul
                                return tx[model][op](args).then(checkResults)
                            })
                        }
                      
                    },
                },
            }
        })
    })

I could collect a lot of issues relating to this, since I went through a lot...
This one was especially helpful: #20016 (reply in thread)
Some struggling with similar issues:
#17948

[ludralph]

Hi @dennemark 👋

You code looks good for the most parts. You can include error handling for the setRLS function. This can prevent the transaction from proceeding if the RLS logic fails.

Prisma.defineExtension((client) => {
    return client.$extends({
        query: {
            $allModels: {
                async $allOperations<T>({ args, query, model, operation, ...rest }: { args: any, query: any, model: any, operation: any }) {

                    const setRLS = async (client: any) => {
                        try {
                            await client.$executeRaw('some rls logic...')
                        } catch (error) {
                            console.error('RLS logic failed:', error)
                            throw new Error('Failed to set RLS')
                        }
                    }

                    const checkResults = (result: any) => {
                        if (result === 'not-valid') {
                            throw new Error('This is just an example check function...')
                        }
                        return result
                    }

                    // Check if this query already runs within a transaction
                    if ((rest as any).__internalParams.transaction) {
                        const transaction = (rest as any).__internalParams.transaction
                        if (transaction.kind === 'itx') {
                            // We are in an interactive transaction and create a client for it
                            const transactionClient = (client as any)._createItxClient(transaction)
                            await setRLS(transactionClient)
                            return transactionClient[model][operation](args).then(checkResults)
                        } else if (transaction.kind === 'batch') {
                            throw new Error('Sequential transactions are not supported in Prisma client extensions.')
                        }
                    } else {
                        // If we are not in a transaction, we can run our own
                        return client.$transaction(async (tx) => {
                            await setRLS(tx)
                            // We place our check results function in an interactive transaction
                            // so that our query is rolled back if check results fail
                            return tx[model][operation](args).then(checkResults)
                        })
                    }
                },
            },
        }
    })
})

[dennemark]

Thanks for the improvement :)

[matheusvellone]

Thanks for sharing this snippet!

I've tweaked on it a bit to also cover $queryRaw calls.

const prisma = new PrismaClient()

const setupRls = (config: string, value: string) => async ({
  query,
  args,
  model,
  operation,
  ...rest
}: any) => {
  const setRlsConfig = (client: PrismaClient) => {
    return client.$executeRaw`SELECT set_config(${config}, ${value}, TRUE)`
  }

  const transaction = (rest as any).__internalParams.transaction

  if (transaction) {
    if (transaction.kind !== 'itx') {
      throw new Error('Non itx transaction')
    }

    const transactionClient = (prisma as any)._createItxClient(transaction)
    await setRlsConfig(transactionClient)

    if (model) {
      return transactionClient[model][operation](args)
    }

    return transactionClient[operation](args)
  }

  const [, result] = await prisma.$transaction([
    setRlsConfig(prisma),
    query(args),
  ])

  return result
}

This way, I create a prisma client for a specific tenant whenever I need and RLS cover also raw queries

const PG_COMPANY_ID_RLS_CONFIG = 'app.current_company_id'

export const getCompanyPrisma = (companyId: Company['id']) => {
  const extension = Prisma.defineExtension((prisma) =>
    prisma.$extends({
      query: {
        // don't use it for $allModels, to cover $queryRaw, $queryRawUnsafe, $executeRaw and $executeRawUnsafe
        $allOperations: setupRls(PG_COMPANY_ID_RLS_CONFIG, companyId),
      },
    })
  )

  return prisma.$extends(extension)
}

(I couldn't get why the checkResults was necessary, so I removed it)

[dennemark]

It seems like my approch only works on the last prisma extension.

return client.$transaction(async (tx) => {
                                setRLS(tx)
                                // we place our check results function in an interactive transaction
                                // so that our query is rolled back if check resul
                                return tx[model][op](args).then(checkResults)
                            })

this part does not properly forward the query to the next extension...

[dennemark]

I do think this approach still has issues with batching in prisma.

[RTAndrew]

For now, it seems to fix my necessities.
In which edge cases have you found that it does not work?

[dennemark]

Hi,
i am using prisma with a graphql setup and when i use prisma findUnique in the relation resolvers it leads to huge performance issues, since prisma is not properly doing batching. So I get N+1 issues.

I solved it for now by creating my own batching solution. It is taken from Yates, my implementation lives in prisma-extension-casl

[emilaos]

Thank you for this.
@dennemark I was wondering if you have seen the __internalParams.transaction hold stable in production?

Asking since documentation about seems quite obscure

[dennemark]

hi, i am not using this approach anymore. also i do not think it would be compatible with newest prisma client.

@emilaos

[emilaos]

Oh! We need to detect if we are in a transaction in an extension exactly as the case mentioned -> to better enforce RLS.

If by any chance could you point me in the right direction / what's your preferred approach I would really appreciate it.

[dennemark]

I adapted an approach of the yates library. But i am not sure anymore if their current version is doing it the same way. But they created their own kind of transaction and batching system that is close to prisma. Since otherwise i think it was difficult to know if you are in a transaction or not.
Tbh. the path you take is probably not sustainable, since the prisma library is changing a lot from v5.

https://github.com/dennemark/prisma-extension-casl/blob/a0f8d6f6a452a9ad973f37571c40d32e6b538dee/src/index.ts#L303

https://github.com/cerebruminc/yates/blob/d1c7853a1245f5b6337f5a0f7d85fc70b0ac91d7/src/index.ts#L672