Supporting Postgres' `SET` across queries of a request

Problem

In Postgres you can set a user for a connection on the database:

await prisma.$executeRaw(`SET current_user_id = ${currentUser.id}`)

This SET can then be used in combination with row-level security (RLS) to issue queries like this:

select * from messages

that only give you back messages from that user. This technique is used by Postgraphile and Postgrest and really takes advantage of what Postgres offers.

Without access to the connection pool, there’s no way to guarantee you’ll get the same connection each query. Since SET values are bound to the query, subsequent queries may be missing the SET or may even override another request’s SET.

I’m not sure what the best approach is. Tying a connection to the lifecycle of a request has its own performance implications.

A point of reference potentially worth investigating. Go’s standard SQL library uses a connection pool under the hood that’s transparent to developers. Do they run into this problem too? If so, do they or how do they deal with it?

Originally from: https://github.com/prisma/prisma/issues/4303#issuecomment-756157408

[wladiston]

TLDR;

Wait the official support for the RLS or go there and contribute on how Prisma works. Don’t try the “quick solution“ thing because you’ll ended up putting more work in than needed. If you want to change Prisma, go ahead and read it.

Idk how you guys have implemented this whole ‘SET’ stuff but I have been implementing the full support to RLS over this entire week for an app that we have using Remix + Prisma. And until this point, I really don’t know if this week was worth the work and if I had actually made a good decision to implement it. So I’ve decided to make this comment and leave it here in case someone is thinking about implementing it and to be 100% sure of what you are stepping into.

I feel like Prisma kinda has lost its point.

First of all, I need to clarify the reasons why I chose to work with Prisma in the first place and probably the main reasons why I feel like this. As a full stack developer, Prisma has increased my productivity in 2-3x at least simplify by replacing the hard work of handling SQL stuff and another 10x just by allowing me to do that in TYPESCRIPT. The ability to capture errors before even thinking that you could be doing something wrong is GOLD (or BTC for some 😅). The model map and the migration tool… shut up and take my money!!! Now, after I had spend this much time on this I feel like this whole RLS support has added too much complexity to setup, maintain the code and at the end I’m spending more time in the migration.sql file than the actual typescript.

So, what have I learned?

1. You need two DB users

As mentioned by @AllanOricil You’ll need to set up a new postgres/db user with limited permission so it doens’t bypass the RLS. Since I need to guarantee that developers can pull the project and things keep working, I needed to add the following script into my migrate file.

-- Create the prisma user only if the user isn't there already
DO
$$
BEGIN
  IF NOT EXISTS (SELECT 1 FROM pg_user WHERE usename = 'prisma') THEN
     CREATE USER prisma WITH PASSWORD 'prism$1234'; -- this password must be changed in production
  END IF;
END
$$;

-- `prisma migrate reset` tearsdown the schema with the permissions so we need to re-grant it
GRANT USAGE ON SCHEMA public TO prisma;
grant select, insert, update, delete on all tables in schema public to prisma;

2. It needs to have two separate connection strings

One for the migrate script with full permission and one with the connection pooling using the prisma user you just created. If you are using supabase, you probably have two connection strings anyway so no big deal.

// package.json
{
  "scripts": {
    "db:deploy": "DB_CONNECTION_STR=$DB_CONNECTION_STR_MIGRATE prisma migrate deploy",
    "db:reset": "DB_CONNECTION_STR=$DB_CONNECTION_STR_MIGRATE prisma migrate reset -f",
    ...
  }
}

See discuttion: #6485

3. You need to manually create policies and enable RLS

Here it is where things started getting overly complicated. So let’s go through this together so you can understand my point.

First you add the following on your migration file which depending on the number of tables your db has, you might have a bit of work to do and depending on the actual complexity of the policies you are screwed. There’s no debug here and you basically figure out that something is wrong only during runtime.

-- RLS
alter table "_CuisineToVenue" enable row level security;
alter table "_DietaryLabelToModifierOption" enable row level security;
alter table "_DietaryLabelToProduct" enable row level security;
alter table "Cart" enable row level security;
alter table "CartItem" enable row level security;
alter table "Cuisine" enable row level security;
alter table "Customization" enable row level security;
alter table "DietaryLabel" enable row level security;
alter table "Location" enable row level security;
alter table "Menu" enable row level security;
alter table "MenuCategory" enable row level security;
alter table "MenuItem" enable row level security;
alter table "MenuModifierGroup" enable row level security;
alter table "MenuModifierOption" enable row level security;
alter table "ModifiersOnItem" enable row level security;
...

-- Policies
CREATE POLICY "Public profiles are viewable by everyone."
  on "Profile" for select
  using ( true );

CREATE POLICY "Users can insert their own profile."
  on "Profile" for insert
  with check ( auth.uid() = id );

CREATE POLICY "Users can update own profile."
  on "Profile" for update
  using ( auth.uid() = id );

CREATE POLICY "Locations are viewable by everyone." ON "public"."Location" FOR SELECT USING (true);
CREATE POLICY "Users can insert locations in their own venues" ON "public"."Location" FOR INSERT WITH CHECK (EXISTS (SELECT 1 FROM "Venue" AS v WHERE v."id" = "venueId" and auth.uid() = v."ownerId" ));
CREATE POLICY "Users can update locations in their own venues" ON "public"."Location" FOR UPDATE USING (EXISTS (SELECT 1 FROM "Venue" AS v WHERE v."id" = "venueId" and auth.uid() = v."ownerId"));

CREATE POLICY "Menus are viewable by everyone." ON "public"."Menu" USING (true);
CREATE POLICY "Users can insert menus in their own venues" ON "public"."Menu" FOR INSERT WITH CHECK (EXISTS (SELECT 1 FROM "Venue" AS v WHERE v."id" = "venueId" and auth.uid() = v."ownerId" ));
CREATE POLICY "Users can update menus in their own venues" ON "public"."Menu" FOR UPDATE USING (EXISTS (SELECT 1 FROM "Venue" AS v WHERE v."id" = "venueId" and auth.uid() = v."ownerId"));

-- too complicated policy
CREATE POLICY "Categories are viewable by everyone." ON "public"."MenuCategory" USING (true);
CREATE POLICY "Users can insert categories in their own venues" ON "public"."MenuCategory" FOR INSERT WITH CHECK (EXISTS (SELECT 1 FROM "Venue" AS v INNER JOIN "Menu" AS m ON v."id" = m."venueId" WHERE m."id" = "menuId" AND v."ownerId" =  auth.uid()));
CREATE POLICY "Users can update categories in their own venues" ON "public"."MenuCategory" FOR UPDATE USING (EXISTS (SELECT 1 FROM "Venue" AS v INNER JOIN "Menu" AS m ON v."id" = m."venueId" WHERE m."id" = "menuId" AND v."ownerId" =  auth.uid()));
...

Now, I’m no SQL expert and I’m pretty sure there are hundreds of ways to do it better but the thing that scratches my head is what comes after few iterations of this. After few migrations created how do you tell what’s actually being applied to this model?

model Cart {
  id         String     @id @default(uuid())
  createdAt  DateTime   @default(now())
  customer   Profile?   @relation(fields: [customerId], references: [id])
  customerId String?    @db.Uuid()
  items      CartItem[]
}

There’s no simple answer. You have to go looking around inside migration files to see what was applied last… in supabase there’s even a bug that after you teardown the schema, you have to turn the RLS off and on again to able to see what policies are applied to each table on their UI.

It should be something that we can maintain like a single source of true. To see and to edit everything through the schema with something like: Just as an example, I haven’t actually put too much thought on how this API would be.

-- defines what should be sent throght 
define session {
  id         String
  role       Role
}

enum Role {
  USER
  ADMIN
}

model Cart {
  id         String     @id @default(uuid())
  customer   Profile   @relation(fields: [customerId], references: [id])
  customerId String    
	
  @rls(read: ({row, session}) => row.customerId === session.id)
}

model Profile {
  id         String     @id @default(uuid())
  cart       Cart[]
	
  @rls(read: ({row, session}) => session.role === 'ADMIN' || row.id === session.id)
}

--
-- or maybe even going further and defining a whole model as the session
-- this way users could manage their session
--
define session {
  @inherit Session
}

model Session {
  id         String     @id @default(uuid())
  profile    Profile    @relation(fields: [profileId], references: [id])
  profileId  String    
  device     String
  expiresAt  DateTime
}

model Cart {
  ...
  customer   Profile   @relation(fields: [customerId], references: [id])
  customerId String    
	
  @rls(read-use: session) // would enforce that a prisma.cart.find() needs a definition of session
  @rls(read: ({row, session}) => row.customerId === session.profileId || session.profile.role === 'ADMIN')

  -- At this point, I don’t know if a multiple policies strategy is better or necessary...
  @rls(read: [
    ({row, session}) => row.customerId === session.profileId,
    ({session}) => session.profile.role === 'ADMIN'
  ])
}

Which brings me to my last consideration…

4. Queries need to be ran inside a transaction

Since it sets a specific variable only for the specific query, you need to wrap all queries around a transaction.

Following examples listed here by many contributors, I managed to create a middleware that puts everything in a transaction and also recreate the $transaction function so I didn’t need to go through the code replacing anything. BUT, long story short, it doesn’t work! The way supabase retrieves the session ends up getting mixed between requests made in parallel so I needed a way pass the session down the tree that’s where I started questioning this whole solution and where I gave up on this thing.

Here it is my implementation

/**
 * This middleware wrap simple query around a transaction so it can respect the RLS
 *
 * @param prisma
 * @returns
 */
const createRlsMiddleware =
  (prisma: PrismaClient) =>
  async (
    params: Prisma.MiddlewareParams,
    next: (params: Prisma.MiddlewareParams) => Promise<any>,
  ) => {
    if (params.runInTransaction) return next(params)
    if (params.model == null) return next(params)

    // Generate model class name from model params (PascalCase to camelCase)
    const modelName =
      params.model.charAt(0).toLowerCase() + params.model.slice(1)

    // @ts-ignore
    const prismaAction = prisma[modelName][params.action](params.args)

    const [results] = await prisma.$transaction([prismaAction])
    return results
  }

/**
 * Set the current user for a transaction
 *
 * @param prisma
 * @returns
 */
export function setCurrentJWT(prisma: PrismaTransaction): PrismaPromise<any> {
  // FIX: the session is not kept separately for each request
  const session = supabase.auth.session()
  const claim = session?.access_token ? jwtDecode(session.access_token) : ''

  // warning: set local does not work completely as you expect
  // https://www.postgresql.org/message-id/flat/[e-mail]
  return prisma.$executeRawUnsafe(
    `SET LOCAL request.jwt.claims = '${JSON.stringify(claim)}'`,
  )
}

type TransactionType = typeof prisma.$transaction
type PrismaTransaction = Omit<
  PrismaClient,
  '$connect' | '$disconnect' | '$on' | '$transaction' | '$use'
>

/**
 * Creates a customized $transaction that injects a jwt in the beginning
 *
 * @param prisma
 * @returns
 */
function createTransaction(prisma: PrismaClient): TransactionType {
  return async <P extends PrismaPromise<any>[]>(
    arg: unknown,
    options?: unknown,
  ) => {
    if (typeof arg === 'function') {
      // @ts-expect-error
      return await prisma.$originalTransaction(async prismaTransaction => {
        await setCurrentJWT(prismaTransaction)
        return arg(prismaTransaction)
      }, options)
    }

    // @ts-expect-error
    const [, ...results] = await prisma.$originalTransaction([
      setCurrentJWT(prisma),
      ...(arg as [...P]),
    ])

    return results
  }
}

/**
 * Creates the prisma cliente
 */
export const prisma = getClient(() => {
  const client = new PrismaClient({
    datasources: {
      db: {
        url: dbUrl.toString(),
      },
    },
  })
  
  // @ts-expect-error
  client.$originalTransaction = client.$transaction
  client.$transaction = createTransaction(client)
  client.$use(createRlsMiddleware(client))

  return client
})

Conclusion

The SET isn’t actually the hardest part to get to work. On the other hand, managing the policies is exponentially harder and making sure that the specific query has the specific variable for the policies is just throwing the best of prisma out of the window.

Ideally you want the query complaining on what policies you breaking

await prisma.cart.findMany({
  define: {session},
  data: {...},
})

await prisma.cart.update({
  define: {role: 'ANON'}, // only USERs & ADMINs can update
  data: {...},
})

await prisma.$transaction([
  ...,
  await prisma.cart.findMany({
    // define: {session}, should not exist
    data: {...},
  })
], {define: {session})

Well, that’s it. I’m feeling a bit frustrated for losing the time invested. I need to ship something that pays the bills for now otherwise I’d gladly open a PR. I hope this is useful to someone.

[hyusetiawan]

it’d be ideal if we can get first-class support for row-level security instead of hopping around the unknown landmine that is reusing sessions, especially for the purpose of authentication where one slip could be disastrous