fix: bug in enhancement proxy for detecting nested transactions

zenstackhq · Jan 9, 2024 · ba763be · ba763be
1 parent 9e840ea
commit ba763be
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 2 deletions.
diff --git a/packages/runtime/src/enhancements/proxy.ts b/packages/runtime/src/enhancements/proxy.ts
@@ -183,6 +183,10 @@ export function makeProxy<T extends PrismaProxyHandler>(
     errorTransformer?: ErrorTransformer
 ) {
     const models = Object.keys(modelMeta.fields).map((k) => k.toLowerCase());
+
+    // a store for saving fields that belong to the proxy (not the target)
+    const proxyStorage: Record<string, unknown> = {};
+
     const proxy = new Proxy(prisma, {
         get: (target: any, prop: string | symbol, receiver: any) => {
             // enhancer metadata
@@ -194,6 +198,10 @@ export function makeProxy<T extends PrismaProxyHandler>(
                 return () => `$zenstack_${name}[${target.toString()}]`;
             }
 
+            if (typeof prop === 'string' && prop in proxyStorage) {
+                return proxyStorage[prop];
+            }
+
             if (prop === '$transaction') {
                 // for interactive transactions, we need to proxy the transaction function so that
                 // when it runs the callback, it provides a proxy to the Prisma client wrapped with
@@ -213,8 +221,14 @@ export function makeProxy<T extends PrismaProxyHandler>(
 
                         const txFunc = input;
                         return $transaction.bind(target)((tx: any) => {
+                            // create a proxy for the transaction function
                             const txProxy = makeProxy(tx, modelMeta, makeHandler, name + '$tx');
+
+                            // record in-transaction flag on the proxy (not the target)
+                            // see logic in "set" handler below
                             txProxy[PRISMA_TX_FLAG] = true;
+
+                            // call the transaction function with the proxy
                             return txFunc(txProxy);
                         }, ...rest);
                     };
@@ -235,6 +249,17 @@ export function makeProxy<T extends PrismaProxyHandler>(
 
             return createHandlerProxy(makeHandler(target, prop), propVal, errorTransformer);
         },
+
+        set: (target: any, prop: string | symbol, value: any) => {
+            if (prop === PRISMA_TX_FLAG) {
+                // set to the proxy store
+                proxyStorage[prop] = value;
+            } else {
+                // pass through to the original target
+                target[prop] = value;
+            }
+            return true;
+        },
     });
 
     return proxy;

diff --git a/tests/integration/tests/enhancements/with-policy/postgres.test.ts b/tests/integration/tests/enhancements/with-policy/postgres.test.ts
@@ -16,14 +16,14 @@ describe('With Policy: with postgres', () => {
 
     beforeEach(async () => {
         dbUrl = await createPostgresDb(DB_NAME);
-        const { prisma: _prisma, withPolicy } = await loadSchemaFromFile(
+        const { prisma: _prisma, enhance } = await loadSchemaFromFile(
             path.join(__dirname, '../../schema/todo-pg.zmodel'),
             {
                 provider: 'postgresql',
                 dbUrl,
             }
         );
-        getDb = withPolicy;
+        getDb = enhance;
         prisma = _prisma;
     });