[ADD] pg.query_ids: query large numbers of ids memory-safely

This is mainly the code that has been recently added to `orm.recompute_fields`,
here we're making it re-usasble.

Author: cawo-odoo

src/base/tests/test_util.py

@@ -18,6 +18,11 @@
 from odoo import modules
 from odoo.tools import mute_logger
 
+try:
+    from odoo.sql_db import db_connect
+except ImportError:
+    from openerp.sql_db import db_connect
+
 from odoo.addons.base.maintenance.migrations import util
 from odoo.addons.base.maintenance.migrations.testing import UnitTestCase, parametrize
 from odoo.addons.base.maintenance.migrations.util import snippets
@@ -1481,6 +1486,27 @@ def test_iter(self):
         self.assertEqual(result, expected)
 
 
+class TestQueryIds(UnitTestCase):
+    def test_straight(self):
+        result = list(util.query_ids(self.env.cr, "SELECT * FROM (VALUES (1), (2)) AS x(x)", itersize=2))
+        self.assertEqual(result, [1, 2])
+
+    def test_chunks(self):
+        with util.query_ids(self.env.cr, "SELECT * FROM (VALUES (1), (2)) AS x(x)") as ids:
+            result = list(util.chunks(ids, 100, fmt=list))
+        self.assertEqual(result, [[1, 2]])
+
+    def test_destructor(self):
+        ids = util.query_ids(self.env.cr, "SELECT id from res_users")
+        del ids
+
+    def test_pk_violation(self):
+        with db_connect(self.env.cr.dbname).cursor() as cr, mute_logger("odoo.sql_db"), self.assertRaises(
+            ValueError
+        ), util.query_ids(cr, "SELECT * FROM (VALUES (1), (1)) AS x(x)") as ids:
+            list(ids)
+
+
 class TestRecords(UnitTestCase):
     def test_rename_xmlid(self):
         cr = self.env.cr

src/util/pg.py

@@ -1932,3 +1932,83 @@ def bulk_update_table(cr, table, columns, mapping, key_col="id"):
         key_col=key_col,
     )
     cr.execute(query, [Json(mapping)])
+
+
+class query_ids(object):
+    """
+    Iterator over ids returned by a query.
+
+    This iterator can memory efficiently query a potentially huge number of ids.
+
+    :param str query: the query that returns the ids. It can be DML, e.g. `UPDATE table WHERE ... RETURNING id`.
+    :param int itersize: passed to a named_cursor, determines the number of rows fetched from PG at once.
+    """
+
+    def __init__(self, cr, query, itersize=None):
+        self._ncr = None
+        self._cr = cr
+        self._tmp_tbl = "_upgrade_query_ids_{}".format(uuid.uuid4().hex)
+        cr.execute(
+            format_query(
+                cr,
+                "CREATE UNLOGGED TABLE {}(id) AS (WITH query AS ({}) SELECT * FROM query)",
+                self._tmp_tbl,
+                SQLStr(query),
+            )
+        )
+        self._len = cr.rowcount
+        try:
+            cr.execute(
+                format_query(
+                    cr,
+                    "ALTER TABLE {} ADD CONSTRAINT {} PRIMARY KEY (id)",
+                    self._tmp_tbl,
+                    "pk_{}_id".format(self._tmp_tbl),
+                )
+            )
+        except psycopg2.IntegrityError as e:
+            if e.pgcode == errorcodes.UNIQUE_VIOLATION:
+                raise ValueError("The query for ids is producing duplicate values: {}".format(query))
+            raise
+        self._ncr = named_cursor(cr, itersize)
+        self._ncr.execute(format_query(cr, "SELECT id FROM {} ORDER BY id", self._tmp_tbl))
+        self._it = iter(self._ncr)
+
+    def _close(self):
+        if self._ncr:
+            if self._ncr.closed:
+                return
+            self._ncr.close()
+        try:
+            self._cr.execute(format_query(self._cr, "DROP TABLE IF EXISTS {}", self._tmp_tbl))
+        except psycopg2.InternalError as e:
+            if e.pgcode != errorcodes.IN_FAILED_SQL_TRANSACTION:
+                raise
+
+    def __len__(self):
+        return self._len
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._ncr.closed:
+            raise StopIteration
+        try:
+            return next(self._it)[0]
+        except StopIteration:
+            self._close()
+            raise
+
+    def next(self):
+        return self.__next__()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self._close()
+        return False
+
+    def __del__(self):
+        self._close()