From 1079a19cd90bbe2b6e550b8f606430f18abdf602 Mon Sep 17 00:00:00 2001
From: Hongbo Wu <hongbo@omnivore.app>
Date: Tue, 26 Dec 2023 20:36:43 +0800
Subject: [PATCH] execute update query in batch to avoid slow bulk action query

---
 packages/api/src/resolvers/article/index.ts |  9 +++-
 packages/api/src/services/library_item.ts   | 47 +++++++++++++++++++--
 2 files changed, 51 insertions(+), 5 deletions(-)

diff --git a/packages/api/src/resolvers/article/index.ts b/packages/api/src/resolvers/article/index.ts
index ec7bd1e8e..5d698fc81 100644
--- a/packages/api/src/resolvers/article/index.ts
+++ b/packages/api/src/resolvers/article/index.ts
@@ -46,6 +46,7 @@ import {
   SearchErrorCode,
   SearchSuccess,
   SetBookmarkArticleError,
+  SetBookmarkArticleErrorCode,
   SetBookmarkArticleSuccess,
   SetFavoriteArticleError,
   SetFavoriteArticleErrorCode,
@@ -70,6 +71,7 @@ import {
   findOrCreateLabels,
 } from '../../services/labels'
 import {
+  batchUpdateLibraryItems,
   createLibraryItem,
   findLibraryItemById,
   findLibraryItemByUrl,
@@ -78,7 +80,6 @@ import {
   sortParamsToSort,
   updateLibraryItem,
   updateLibraryItemReadingProgress,
-  updateLibraryItems,
 } from '../../services/library_item'
 import { parsedContentToLibraryItem } from '../../services/save_page'
 import {
@@ -551,6 +552,10 @@ export const setBookmarkArticleResolver = authorized<
   SetBookmarkArticleError,
   MutationSetBookmarkArticleArgs
 >(async (_, { input: { articleID } }, { uid, log, pubsub }) => {
+  if (!articleID) {
+    return { errorCodes: [SetBookmarkArticleErrorCode.NotFound] }
+  }
+
   // delete the item and its metadata
   const deletedLibraryItem = await updateLibraryItem(
     articleID,
@@ -860,7 +865,7 @@ export const bulkActionResolver = authorized<
         labels = await findLabelsByIds(labelIds, uid)
       }
 
-      await updateLibraryItems(
+      await batchUpdateLibraryItems(
         action,
         {
           query,
diff --git a/packages/api/src/services/library_item.ts b/packages/api/src/services/library_item.ts
index 99f143f48..e97f27ba8 100644
--- a/packages/api/src/services/library_item.ts
+++ b/packages/api/src/services/library_item.ts
@@ -886,7 +886,7 @@ export const countByCreatedAt = async (
   )
 }
 
-export const updateLibraryItems = async (
+export const batchUpdateLibraryItems = async (
   action: BulkActionType,
   searchArgs: SearchArgs,
   userId: string,
@@ -902,7 +902,7 @@ export const updateLibraryItems = async (
   }
 
   // build the script
-  let values: QueryDeepPartialEntity<LibraryItem> = {}
+  let values: Record<string, any> = {}
   let addLabels = false
   switch (action) {
     case BulkActionType.Archive:
@@ -984,7 +984,48 @@ export const updateLibraryItems = async (
       return tx.getRepository(EntityLabel).save(labelsToAdd)
     }
 
-    return queryBuilder.update(LibraryItem).set(values).execute()
+    const countSql = queryBuilder.select('COUNT(1) INTO total_rows').getSql()
+    const [subQuery, params] = queryBuilder.select('id').getQueryAndParameters()
+    const valuesSql = Object.keys(values)
+      // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
+      .map((key) => `${key} = ${values[key]}`)
+      .join(', ')
+
+    const sql = `
+    -- Set batch size
+    DO $$ 
+    DECLARE 
+        batch_size INT := 100;
+        total_rows INT := 1000;
+        num_batches INT;
+        current_offset INT;
+    BEGIN
+        -- Get the total count of rows to be updated
+        ${countSql};
+
+        -- Calculate the number of batches
+        num_batches := CEIL(total_rows * 1.0 / batch_size);
+
+        -- Loop through batches
+        FOR i IN 0..num_batches-1 LOOP
+            -- Set the current offset
+            current_offset := i * batch_size;
+
+            -- Perform incremental update in batches using LIMIT and OFFSET
+            UPDATE omnivore.library_item
+            SET ${valuesSql}
+            FROM (
+                ${subQuery}
+                ORDER BY id
+                LIMIT batch_size
+                OFFSET current_offset
+            ) AS batch
+            WHERE library_item.id = batch.id;
+        END LOOP;
+    END $$
+    `
+
+    return tx.query(sql, params)
   })
 }