Ensure correct batch context for queries (#23)

src/queries.ts

@@ -83,6 +83,10 @@ export const getSyntaxProxy = (config?: {
         // all queries within it think they are running inside a batch transaction,
         // in order to retrieve their serialized values.
         if (typeof value === 'function') {
+          // Temporarily store the original value of `IN_BATCH`, so that we can resume it
+          // after the nested function has been called.
+          const ORIGINAL_IN_BATCH = IN_BATCH;
+
           // Since `value()` is synchronous, `IN_BATCH` should not affect any
           // other queries somewhere else in the app, even if those are run inside
           // an asynchronous function, so we don't need to use `IN_BATCH_ASYNC`,
@@ -117,7 +121,8 @@ export const getSyntaxProxy = (config?: {
             value = wrapExpressions(value);
           }
 
-          IN_BATCH = false;
+          // Restore the original value of `IN_BATCH`.
+          IN_BATCH = ORIGINAL_IN_BATCH;
         }
 
         // If the function call is happening after an existing function call in the

tests/queries.test.ts

@@ -208,6 +208,51 @@ describe('syntax proxy', () => {
     ]);
   });
 
+  test('using nested function as argument in batch', async () => {
+    // It's important to define the `callback` here, in order to guarantee that the
+    // queries are executed standalone if no batch context is detected.
+    const callback = () => undefined;
+
+    const addProxy = getSyntaxProxy({ rootProperty: 'add', callback });
+    const getProxy = getSyntaxProxy({ rootProperty: 'get', callback });
+    const alterProxy = getSyntaxProxy({ rootProperty: 'alter', callback });
+
+    const queryList = getBatchProxy(() => [
+      addProxy.newUsers.to(() => getProxy.oldUsers()),
+      // Assert whether function chaining is still possible after a nested function call,
+      // since the latter is able to manipulate the batch context.
+      alterProxy
+        .model('newUsers')
+        .to({ slug: 'accounts' }),
+    ]);
+
+    expect(queryList).toEqual([
+      {
+        structure: {
+          add: {
+            newUsers: {
+              to: {
+                __RONIN_QUERY: {
+                  get: { oldUsers: {} },
+                },
+              },
+            },
+          },
+        },
+      },
+      {
+        structure: {
+          alter: {
+            model: 'newUsers',
+            to: {
+              slug: 'accounts',
+            },
+          },
+        },
+      },
+    ]);
+  });
+
   test('using schema query types', async () => {
     const callback = () => undefined;