From 02dd5087ca674557548136ba605dee52cb99fd8f Mon Sep 17 00:00:00 2001
From: exAspArk <exaspark@gmail.com>
Date: Thu, 24 Oct 2024 09:16:56 -0400
Subject: [PATCH] Enable passing application context in non-public schemas

---
 CHANGELOG.md |  4 ++++
 package.json |  2 +-
 src/cli.ts   | 15 ++++++++-------
 src/index.ts |  6 ++++--
 4 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 59127ec..6e9487e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # Changelog
 
+#### [v0.7.0](https://github.com/BemiHQ/bemi-prisma/compare/v0.6.0...v0.7.0) - 2024-10-24
+
+- Enable passing application context in non-public schemas.
+
 #### [v0.6.0](https://github.com/BemiHQ/bemi-prisma/compare/v0.5.0...v0.6.0) - 2024-10-14
 
 - Fix compatibility with Prisma v5.20+
diff --git a/package.json b/package.json
index 641da04..903c65d 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@bemi-db/prisma",
-  "version": "0.6.0",
+  "version": "0.7.0",
   "description": "Automatic data change tracking for Prisma",
   "main": "dist/index.js",
   "module": "./dist/index.mjs",
diff --git a/src/cli.ts b/src/cli.ts
index 2005f95..4998aae 100644
--- a/src/cli.ts
+++ b/src/cli.ts
@@ -26,19 +26,20 @@ $$ LANGUAGE plpgsql;
 CREATE OR REPLACE PROCEDURE _bemi_create_triggers()
 AS $$
 DECLARE
+  current_schemaname TEXT;
   current_tablename TEXT;
 BEGIN
-  FOR current_tablename IN
-    SELECT tablename FROM pg_tables
+  FOR current_schemaname, current_tablename IN
+    SELECT schemaname, tablename FROM pg_tables
     LEFT JOIN information_schema.triggers ON tablename = event_object_table AND schemaname = trigger_schema AND trigger_name LIKE '_bemi_row_trigger_%'
-    WHERE schemaname = 'public' AND trigger_name IS NULL
-    GROUP BY tablename
+    WHERE schemaname NOT IN ('information_schema', 'pg_catalog') AND trigger_name IS NULL
+    GROUP BY schemaname, tablename
   LOOP
     EXECUTE format(
-      'CREATE OR REPLACE TRIGGER _bemi_row_trigger_%s
-      BEFORE INSERT OR UPDATE OR DELETE ON %I FOR EACH ROW
+      'CREATE OR REPLACE TRIGGER _bemi_row_trigger_%s_%s
+      BEFORE INSERT OR UPDATE OR DELETE ON %I.%I FOR EACH ROW
       EXECUTE FUNCTION _bemi_row_trigger_func()',
-      current_tablename, current_tablename
+      current_schemaname, current_tablename, current_schemaname, current_tablename
     );
   END LOOP;
 END;
diff --git a/src/index.ts b/src/index.ts
index 35cddee..73e3f3d 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -8,7 +8,8 @@ import { isContextComment, isWriteQuery, contextToSqlComment } from './pg-utils'
 import { logger } from './logger'
 
 const WRITE_OPERATIONS = ["create", "update", "upsert", "delete", "createMany", "updateMany", "deleteMany"]
-const EXECUTE_OPERATIONS = ["$executeRaw", "$executeRawUnsafe"]
+const EXECUTE_RAW_UNSAFE_OPERATION = ["$executeRawUnsafe"]
+const EXECUTE_OPERATIONS = ["$executeRaw", EXECUTE_RAW_UNSAFE_OPERATION]
 const ASYNC_LOCAL_STORAGE = new AsyncLocalStorage();
 const MAX_CONTEXT_SIZE = 1000000 // ~ 1MB
 
@@ -21,6 +22,7 @@ export const withPgAdapter = <PrismaClientType>(
   const prisma = (originalPrisma as any).$extends({
     query: {
       async $allOperations({ args, query, operation, model }: any) {
+
         // Not included model
         if (model && includeModels && !includeModels.includes(model)) {
           return query(args)
@@ -38,7 +40,7 @@ export const withPgAdapter = <PrismaClientType>(
         }
 
         // Injected context query
-        if (operation === '$executeRawUnsafe' && args[0] && isContextComment(args[0])) {
+        if (operation === EXECUTE_RAW_UNSAFE_OPERATION && args[0] && isContextComment(args[0])) {
           return query(args)
         }