85 lines
3.3 KiB
TypeScript
85 lines
3.3 KiB
TypeScript
import { PrismaClient } from '@prisma/client';
|
|
import { PrismaPg } from '@prisma/adapter-pg';
|
|
import { Pool } from 'pg';
|
|
import { UserSession } from './auth';
|
|
|
|
let pool: Pool;
|
|
let globalPrisma: PrismaClient;
|
|
|
|
const getPrismaClient = (): PrismaClient => {
|
|
if (!globalPrisma) {
|
|
const dbUrl = new URL(process.env.DATABASE_URL!);
|
|
pool = new Pool({
|
|
host: dbUrl.hostname,
|
|
port: dbUrl.port ? parseInt(dbUrl.port) : 5432,
|
|
user: decodeURIComponent(dbUrl.username),
|
|
password: decodeURIComponent(dbUrl.password),
|
|
database: dbUrl.pathname.substring(1).split('?')[0],
|
|
ssl: false
|
|
});
|
|
const adapter = new PrismaPg(pool);
|
|
globalPrisma = new PrismaClient({
|
|
adapter,
|
|
log: ['query', 'info', 'warn', 'error']
|
|
});
|
|
}
|
|
return globalPrisma;
|
|
};
|
|
|
|
export const getPrisma = (session?: UserSession | null) => {
|
|
const prisma = getPrismaClient();
|
|
|
|
if (!session) {
|
|
return prisma;
|
|
}
|
|
|
|
return prisma.$extends({
|
|
client: {
|
|
async $transaction(args: any, options?: any) {
|
|
if (typeof args === 'function') {
|
|
const originalFn = args;
|
|
return prisma.$transaction(async (tx) => {
|
|
await tx.$executeRawUnsafe(`SET LOCAL app.current_user_id = '${session.userId}';`);
|
|
await tx.$executeRawUnsafe(`SET LOCAL app.current_user_role = '${session.role}';`);
|
|
await tx.$executeRawUnsafe(`SET LOCAL app.current_hotel_id = '${session.hotelId}';`);
|
|
await tx.$executeRawUnsafe(`SET LOCAL app.current_region_id = '${session.regionId}';`);
|
|
return originalFn(tx);
|
|
}, options);
|
|
} else if (Array.isArray(args)) {
|
|
const rlsQueries = [
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_user_id = '${session.userId}';`),
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_user_role = '${session.role}';`),
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_hotel_id = '${session.hotelId}';`),
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_region_id = '${session.regionId}';`),
|
|
];
|
|
const results = await prisma.$transaction([...rlsQueries, ...args], options);
|
|
return results.slice(rlsQueries.length);
|
|
}
|
|
return prisma.$transaction(args, options);
|
|
}
|
|
},
|
|
query: {
|
|
$allModels: {
|
|
async $allOperations({ args, query, __internalParams }: any) {
|
|
// Check if we are inside a transaction (either batch or interactive transaction)
|
|
const isInsideTransaction = __internalParams?.transaction !== undefined;
|
|
|
|
if (isInsideTransaction) {
|
|
return query(args);
|
|
}
|
|
|
|
// Execute RLS parameter setting followed by the original query in a batch transaction
|
|
const results = await prisma.$transaction([
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_user_id = '${session.userId}';`),
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_user_role = '${session.role}';`),
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_hotel_id = '${session.hotelId}';`),
|
|
prisma.$executeRawUnsafe(`SET LOCAL app.current_region_id = '${session.regionId}';`),
|
|
query(args)
|
|
]);
|
|
// Return the results of the query block
|
|
return results[4];
|
|
}
|
|
}
|
|
}
|
|
});
|
|
};
|