collections: batch set calls into pages

OpenFn · Nov 15, 2024 · a21441b · a21441b
1 parent 055b679
commit a21441b
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 18 deletions.
diff --git a/packages/collections/src/collections.js b/packages/collections/src/collections.js
@@ -136,6 +136,8 @@ export function get(name, query = {}) {
  */
 export function set(name, keyGen, values) {
   return async state => {
+    const batchSize = 1000;
+
     const [resolvedName, resolvedValues] = expandReferences(
       state,
       name,
@@ -155,28 +157,36 @@ export function set(name, keyGen, values) {
       value: JSON.stringify(value),
     }));
 
-    const response = await request(state, getClient(state), resolvedName, {
-      method: 'POST',
-      body: JSON.stringify({ items: pairs }),
-      headers: {
-        'content-type': 'application/json',
-      },
-    });
+    while (pairs.length) {
+      const batch = pairs.splice(0, batchSize);
 
-    if (response.statusCode >= 400) {
       console.log(
-        `Collections: Error setting ${pairs.length} values in "${name}"`
+        `Collections: uploading batch of ${batch.length} values to "${name}"...`
       );
-      const text = await response.body.text();
-      const e = new Error('ERROR from collections server:' + 400);
-      e.body = text;
-      throw e;
-    }
+      const response = await request(state, getClient(state), resolvedName, {
+        method: 'POST',
+        body: JSON.stringify({ items: batch }),
+        headers: {
+          'content-type': 'application/json',
+        },
+      });
 
-    const result = await response.body.json();
-    console.log(`Collections: set ${result.upserted} values in "${name}"`);
-    if (result.error) {
-      console.log(`Collections: errors reported on set:`, result.error);
+      if (response.statusCode >= 400) {
+        console.log(
+          `Collections: Error setting ${batch.length} values in "${name}"`
+        );
+        const text = await response.body.text();
+        const e = new Error('ERROR from collections server:' + 400);
+        e.body = text;
+        throw e;
+      }
+
+      const result = await response.body.json();
+      console.log(`Collections: set ${result.upserted} values in "${name}"`);
+
+      if (result.error) {
+        console.log(`Collections: errors reported on set:`, result.error);
+      }
     }
 
     return state;

diff --git a/packages/collections/src/mock.js b/packages/collections/src/mock.js
@@ -73,6 +73,9 @@ export function API() {
   const asJSON = (name, key) => {
     return JSON.parse(collections[name][key]);
   };
+  const count = name => {
+    return Object.keys(collections[name]).length;
+  };
 
   // TODO strictly speaking this should support patterns
   // but keeping it super simple in the mock for now
@@ -103,6 +106,7 @@ export function API() {
     remove,
     byKey,
     asJSON,
+    count,
   };
 
   return api;

diff --git a/packages/collections/test/Adaptor.test.js b/packages/collections/test/Adaptor.test.js
@@ -370,6 +370,23 @@ describe('set', () => {
     const y = api.asJSON(COLLECTION, items[1].id);
     expect(y).to.eql(items[1]);
   });
+
+  // TODO: there's no actual test of pagination here, save the logs
+  it.only('should set several batches of items', async () => {
+    const { state } = init();
+
+    // the collection has one item by default
+    expect(api.count(COLLECTION)).to.equal(1);
+
+    const items = new Array(2499).fill(1).map((_item, idx) => ({
+      id: `${idx}`,
+    }));
+    const keygen = item => item.id;
+
+    await collections.set(COLLECTION, keygen, items)(state);
+
+    expect(api.count(COLLECTION)).to.equal(2500);
+  });
 });
 
 describe('remove', () => {