Restore Supabase RLS with Drizzle using tRPC middlewares

[craxrev]

---

subtitle: A journey tring to restore supabase RLS when using Drizzle ORM using tRPC middlewares
lede: An experiment trying to fix Supabase RLS not working when using a direct db connection through an ORM, in a context of a tRPC, we can use middleware to create a wrapper around db calls and inject postgres set_config cmd to impersonate the correct roles and claims.

Edit:

This post has become obsolete with the current versions of the used services, though it was a good experiment that I've enjoyed, now Supabase have a good support for DrizzleORM, and with the Next.JS 15 server actions, you could basically replicate what tRPC have.

I have built a boilerplate template that make use of the latest version of DrizzleORM, Server Actions, a replicated middleware implementation similar to tRPC's and much more: Locastic/next-web3-hybrid-starter

Intro

First of all, let's go througt some basics, what are those:

• Supabase: it's a DaaS open-source project that aims to simplify running and hosting Postgres databases.
• RLS: Role-Level Security, it's a postgres security functionality that aims to protect postgres data on the row level, Supabase uses it as one of their selling points for their implementation.
• Drizzle: is a new typescript ORM that memicks real SQL syntaxt without any runtime and overheads (like prisma ORM).
• tRPC: is typescript typesafe api "builder" that seemlisely integrates with your frontend.

Issue

So I've been fighting this issue for months not only with Drizzle but with Prisma before, and my solution has been to either ignore RLS entirely or use the Supabase js library.

The issue is that Supabase RLS is only working when we use their JS SDK, if you choose otherwise, then you're on your own, it's the ORM job to figure out how roles and claims to be managed and that's what has been the case for Prisma, they now have their own plugin to handle that, for Drizzle, there have been some talks like drizzle-team/drizzle-orm/issues/594, porsager/postgres/issues/559 and also some tries like rphlmr/drizzle-supabase-rls.

So a more detailed summary of the issue: if you use an ORM, the default postgres role that will be used is postgres which is the role with admin privileges, that role will be ignored by Supabase RLS, so all your requests will go right through as if RLS is not enabled.

A quick overview about supabase RLS internals, supabase has default internal functions that will be used whenever you want to setup RLS in supabase dashboard, these postgres functions are: auth.email(), auth.jwt(), auth.role(), auth.uid(), calling them will check request.jwt.claim(s) setting and read off the jwt to extract information about the user that's making the postgres request.

That's one thing we need to manually inject, the next thing is the postgres role, we need to change the role from postgres to either authenticated or anon.

For this, there are several solutions proposed on the aforementioned discussions above, but IMHO the best one is to use tRPC middlewares as they're best suited for cases like this.

Solution

Let's go through this, in a default t3 stack setup, we have the following files:

import postgres from "postgres";
import { drizzle } from "drizzle-orm/postgres-js";

import { env } from "@/env";
import * as schema from "./schema";

const client = postgres(env.DATABASE_URL, { prepare: false });
export const db = drizzle(client, { schema });

import { db } from "@/server/db";

export const createTRPCContext = async (opts: { headers: Headers }) => {
  const session = await getServerAuthSession();

  return {
    db,
    session,
    ...opts,
  };
};

/*
...
*/

export const protectedProcedure = t.procedure.use(({ ctx, next }) => {
  if (!ctx.session) {
    throw new TRPCError({ code: "UNAUTHORIZED" });
  }
  return next({
    ctx: {
      // infers the `session` as non-nullable
      session: { ...ctx.session },
    },
  });
});

Now we need to wrap the next() callback with the needed logic to change the default postgres role, we need to do the following:

• Inject jwt claims in order to make supabase be able to use auth.* functions.
• Change role to either authenticated or anon.

export const protectedProcedure = t.procedure.use(({ ctx, next }) => {
  return ctx.db.transaction(async (tx) => {
    if (!ctx.session) {
      throw new TRPCError({ code: "UNAUTHORIZED" });
    }

    // add jwt claims, it should be JSON stringified with the need
    // information (uid, email, role, ..)
    await tx.execute(
      sql.raw(
        `SELECT set_config('request.jwt.claims', '{claims}', TRUE)`,
      ),
    );
    // Change default role to either 'authenticated' or 'anon'
    await tx.execute(sql.raw(`SET ROLE '{role}'`));

    const res = await next({
      ctx: {
        // override default `db` with `tx`
        db: tx,
        // infers the `session` as non-nullable
        session: { ...ctx.session },
      },
    });

    // reset default claims and role
    await tx.execute(sql`SELECT set_config('request.jwt.claims', NULL, TRUE)`);
    await tx.execute(sql`RESET ROLE`);

    return res;
  });
});

A more refurbished version could be like this:

/**
 * Admin (postgres) procedure
 */
export const adminProcedure = t.procedure;

/**
 * Public (anon | authenticated) procedure
 */
export const publicProcedure = t.procedure.use(({ ctx, next }) => {
  return ctx.db.transaction(async (tx) => {
    if (ctx.session) {
      await tx.execute(
        sql.raw(
          `SELECT set_config('request.jwt.claims', '${JSON.stringify(ctx.session)}', TRUE)`,
        ),
      );
      await tx.execute(sql.raw(`SET ROLE 'authenticated'`));
    } else {
      await tx.execute(sql.raw(`SET ROLE 'anon'`));
    }
    const res = await next({
      ctx: {
        // override default `db` with `tx`
        db: tx,
      },
    });

    if (ctx.session) {
      await tx.execute(
        sql`SELECT set_config('request.jwt.claims', NULL, TRUE)`,
      );
    }
    await tx.execute(sql`RESET ROLE`);

    return res;
  });
});

/**
 * Protected (authenticated) procedure
 */
export const protectedProcedure = t.procedure.use(({ ctx, next }) => {
  return db.transaction(async (tx) => {
    if (!ctx.session) {
      throw new TRPCError({ code: "UNAUTHORIZED" });
    }

    await tx.execute(
      sql.raw(
        `SELECT set_config('request.jwt.claims', '${JSON.stringify(ctx.session)}', TRUE)`,
      ),
    );
    await tx.execute(sql.raw(`SET ROLE 'authenticated'`));

    const res = await next({
      ctx: {
        // override default `db` with `tx`
        db: tx,
        // infers the `session` as non-nullable
        session: { ...ctx.session },
      },
    });

    await tx.execute(sql`SELECT set_config('request.jwt.claims', NULL, TRUE)`);
    await tx.execute(sql`RESET ROLE`);

    return res;
  });
});

Caveats

Now with this, one should be careful of edge cases, this should not be THE solution for this issue but rather, a workaround, a more legit solution should come from the ORM itself, hopefully it will come soon and this workaround will become absolete.

If you ever stambled on this, I hope it was helpful.