RLS drizzle client that hooks into `transaction` to set variables for supabase RLS policies

supabase-rls-drizzle-client.ts

// rlsPgClient is a separate database connection using a postgres user that has 
// permissions only to assume 'authenticated' and 'anon' roles

const rlsDrizzleBase = drizzle(rlsPgClient, { schema });
export function createRlsDrizzle({ session }: { session: Session | undefined | null }) {
  let claims = "";
  if (session) {
    const accessToken = session.access_token;
    const parsedJwt = JSON.parse(
      Buffer.from(z.string().parse(accessToken.split(".")[1]), "base64").toString(),
    );
    claims = JSON.stringify(parsedJwt ?? {});
  }
  return new Proxy<typeof rlsDrizzleBase>(rlsDrizzleBase, {
    get(target, prop, receiver) {
      if (prop === "transaction") {
        return async (
          first: DrizzleTransactionFunctionFirst,
          ...rest: DrizzleTransactionFunctionRest
        ) => {
          return target.transaction(
            async (tx) => {
              if (claims) {
                await tx.execute(
                  sql.raw(
                    `SET LOCAL ROLE authenticated; SELECT set_config('request.jwt.claims', '${claims}', TRUE);`,
                  ),
                );
              } else {
                await tx.execute(sql.raw("SET LOCAL ROLE anon;"));
              }
              const result = await first(tx);
              // make sure to reset the role and permissions to avoid a supabase connection pooler bug
              // that may reuse transactions. See:
              // https://github.com/drizzle-team/drizzle-orm/issues/594#issuecomment-1917830225
              await tx.execute(
                sql.raw("SELECT set_config('request.jwt.claims', NULL, true); RESET ROLE;"),
              );
              return result;
            },
            ...rest,
          );
        };
      }
      return Reflect.get(target, prop, receiver);
    },
  });
}