fix query autofetch

OpenFn · Oct 31, 2024 · ff1081b · ff1081b
1 parent 3f56c87
commit ff1081b
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 44 deletions.
diff --git a/packages/salesforce/ast.json b/packages/salesforce/ast.json
@@ -464,7 +464,8 @@
       "name": "query",
       "params": [
         "qs",
-        "options"
+        "options",
+        "callback"
       ],
       "docs": {
         "description": "Execute an SOQL query.\nNote that in an event of a query error,\nerror logs will be printed but the operation will not throw the error.\n\nThe Salesforce query API is subject to rate limits, {@link https://sforce.co/3W9zyaQ See for more details}.",
@@ -519,6 +520,15 @@
             "name": "options.autoFetch",
             "default": "false"
           },
+          {
+            "title": "param",
+            "description": "A callback to execute once the record is retrieved",
+            "type": {
+              "type": "NameExpression",
+              "name": "function"
+            },
+            "name": "callback"
+          },
           {
             "title": "returns",
             "description": null,

diff --git a/packages/salesforce/src/Adaptor.js b/packages/salesforce/src/Adaptor.js
@@ -458,63 +458,75 @@ export function post(path, data, options = {}) {
  * @param {string} qs - A query string. Must be less than `4000` characters in WHERE clause
  * @param {object} options - Options passed to the bulk api.
  * @param {boolean} [options.autoFetch=false] - Fetch next records if available.
+ * @param {function} callback - A callback to execute once the record is retrieved
  * @returns {Operation}
  */
-export function query(qs, options = {}) {
+export function query(qs, options = {}, callback = s => s) {
   return async state => {
-    let done = false;
-    let qResult = null;
-    let result = [];
-
     const { connection } = state;
     const [resolvedQs, resolvedOptions] = expandReferences(state, qs, options);
-    const { autoFetch = false } = resolvedOptions;
-
     console.log(`Executing query: ${resolvedQs}`);
+    const autoFetch = resolvedOptions.autoFetch || resolvedOptions.autofetch;
+
+    if (autoFetch) {
+      console.log('autoFetch is enabled: all records will be downloaded');
+    }
+
+    const result = {
+      done: true,
+      totalSize: 0,
+      records: [],
+    };
+
+    const processRecords = async res => {
+      const { done, totalSize, records, nextRecordsUrl } = res;
+
+      result.done = done;
+      result.totalSize = totalSize;
+      result.records.push(...records);
+
+      if (!done && !autoFetch && nextRecordsUrl) {
+        result.nextRecordsUrl = nextRecordsUrl;
+      }
+      if (!done && autoFetch) {
+        console.log('Fetched records so far:', result.records.length);
+        console.log('Fetching next records...');
+
+        try {
+          const newResult = await connection.request({ url: nextRecordsUrl });
+          await processRecords(newResult);
+        } catch (err) {
+          const { message, errorCode } = err;
+          console.error(`Error ${errorCode}: ${message}`);
+          throw err;
+        }
+      }
+    };
+
     try {
-      qResult = await connection.query(resolvedQs);
+      const qResult = await connection.query(resolvedQs);
+      if (qResult.totalSize > 0) {
+        console.log('Total records:', qResult.totalSize);
+        await processRecords(qResult);
+        console.log('Done ✔ retrieved records:', result.records.length);
+      } else {
+        console.log('No records found.');
+      }
     } catch (err) {
       const { message, errorCode } = err;
       console.log(`Error ${errorCode}: ${message}`);
       throw err;
     }
 
-    if (qResult.totalSize > 0) {
-      console.log('Total records', qResult.totalSize);
-
-      while (!done) {
-        result.push(qResult);
-
-        if (qResult.done) {
-          done = true;
-        } else if (autoFetch) {
-          console.log(
-            'Fetched records so far',
-            result.map(ref => ref.records).flat().length
-          );
-          console.log('Fetching next records...');
-          try {
-            qResult = await connection.request({ url: qResult.nextRecordsUrl });
-          } catch (err) {
-            const { message, errorCode } = err;
-            console.log(`Error ${errorCode}: ${message}`);
-            throw err;
-          }
-        } else {
-          done = true;
-        }
-      }
-
-      console.log(
-        'Done ✔ retrieved records',
-        result.map(ref => ref.records).flat().length
-      );
-    } else {
-      result.push(qResult);
-      console.log('No records found.');
-    }
+    console.log(
+      'Results retrieved and pushed to position [0] of the references array.'
+    );
 
-    return composeNextState(state, result, result?.records);
+    const nextState = {
+      ...state,
+      references: [result, ...state.references],
+    };
+    return callback(nextState);
   };
 }