[FIX] spreadsheet: batch process spreadsheet_revision.commands

src/util/spreadsheet/misc.py

@@ -1,32 +1,58 @@
-from .. import json
+from .. import json, pg
+
+MEMORY_CAP = 2 * 10**8  # 200MB
 
 
 def iter_commands(cr, like_all=(), like_any=()):
     if not (bool(like_all) ^ bool(like_any)):
         raise ValueError("Please specify `like_all` or `like_any`, not both")
-    cr.execute(
-        """
-        SELECT id,
-               commands
-          FROM spreadsheet_revision
-         WHERE commands LIKE {}(%s::text[])
-        """.format("ALL" if like_all else "ANY"),
-        [list(like_all or like_any)],
-    )
-    for revision_id, data in cr.fetchall():
-        data_loaded = json.loads(data)
-        if "commands" not in data_loaded:
-            continue
-        data_old = json.dumps(data_loaded, sort_keys=True)
-
-        changed = yield data_loaded["commands"]
-        if changed is None:
-            changed = data_old != json.dumps(data_loaded, sort_keys=True)
-
-        if changed:
-            cr.execute(
-                "UPDATE spreadsheet_revision SET commands=%s WHERE id=%s", [json.dumps(data_loaded), revision_id]
+
+    with pg.named_cursor(cr, itersize=1) as ncr:
+        ncr.execute(
+            """
+            WITH filtered AS (
+                SELECT id,
+                       commands,
+                       LENGTH(commands) AS commands_length
+                  FROM spreadsheet_revision
+                 WHERE commands LIKE {condition} (%s::text[])
+            ), smaller AS (
+                SELECT id,
+                       commands,
+                       sum(commands_length) OVER (ORDER BY id) / %s AS num
+                  FROM filtered
+                 WHERE commands_length <= %s
             )
+            SELECT array_agg(id ORDER BY id),
+                   array_agg(commands ORDER BY id)
+              FROM smaller
+          GROUP BY num
+
+             UNION ALL
+
+            SELECT ARRAY[id],
+                   ARRAY[commands]
+              FROM filtered
+             WHERE commands_length > %s
+            """.format(condition=pg.SQLStr("ALL" if like_all else "ANY")),
+            [list(like_any or like_all), MEMORY_CAP, MEMORY_CAP, MEMORY_CAP],
+        )
+        for ids, commands in ncr:
+            for revision_id, data in zip(ids, commands):
+                data_loaded = json.loads(data)
+                if "commands" not in data_loaded:
+                    continue
+                data_old = json.dumps(data_loaded, sort_keys=True)
+
+                changed = yield data_loaded["commands"]
+                if changed is None:
+                    changed = data_old != json.dumps(data_loaded, sort_keys=True)
+
+                if changed:
+                    cr.execute(
+                        "UPDATE spreadsheet_revision SET commands=%s WHERE id=%s",
+                        [json.dumps(data_loaded), revision_id],
+                    )
 
 
 def process_commands(cr, callback, *args, **kwargs):