Skip to content

Commit baafcc7

Browse files
committed
[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.
1 parent 9a1be7d commit baafcc7

File tree

2 files changed

+108
-0
lines changed

2 files changed

+108
-0
lines changed

src/base/tests/test_util.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@
1818
from odoo import modules
1919
from odoo.tools import mute_logger
2020

21+
try:
22+
from odoo.sql_db import db_connect
23+
except ImportError:
24+
from openerp.sql_db import db_connect
25+
2126
from odoo.addons.base.maintenance.migrations import util
2227
from odoo.addons.base.maintenance.migrations.testing import UnitTestCase, parametrize
2328
from odoo.addons.base.maintenance.migrations.util import snippets
@@ -1494,6 +1499,31 @@ def test_iter(self):
14941499
self.assertEqual(result, expected)
14951500

14961501

1502+
class TestQueryIds(UnitTestCase):
1503+
def test_straight(self):
1504+
with db_connect(self.env.cr.dbname).cursor() as cr:
1505+
result = list(util.query_ids(cr, "SELECT * FROM (VALUES (1), (2)) AS x(x)", itersize=2))
1506+
self.assertEqual(result, [1, 2])
1507+
1508+
def test_chunks(self):
1509+
with db_connect(self.env.cr.dbname).cursor() as cr, util.query_ids(
1510+
cr, "SELECT * FROM (VALUES (1), (2)) AS x(x)"
1511+
) as ids:
1512+
result = list(util.chunks(ids, 100, fmt=list))
1513+
self.assertEqual(result, [[1, 2]])
1514+
1515+
def test_destructor(self):
1516+
with db_connect(self.env.cr.dbname).cursor() as cr:
1517+
ids = util.query_ids(cr, "SELECT id from res_users")
1518+
del ids
1519+
1520+
def test_pk_violation(self):
1521+
with db_connect(self.env.cr.dbname).cursor() as cr, mute_logger("odoo.sql_db"), self.assertRaises(
1522+
ValueError
1523+
), util.query_ids(cr, "SELECT * FROM (VALUES (1), (1)) AS x(x)") as ids:
1524+
list(ids)
1525+
1526+
14971527
class TestRecords(UnitTestCase):
14981528
def test_rename_xmlid(self):
14991529
cr = self.env.cr

src/util/pg.py

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1932,3 +1932,81 @@ def bulk_update_table(cr, table, columns, mapping, key_col="id"):
19321932
key_col=key_col,
19331933
)
19341934
cr.execute(query, [Json(mapping)])
1935+
1936+
1937+
class query_ids(object):
1938+
"""
1939+
Iterator over ids returned by a query.
1940+
1941+
This iterator can memory efficiently query a potentially huge number of ids.
1942+
1943+
:param str query: the query that returns the ids. It can be DML, e.g. `UPDATE table WHERE ... RETURNING id`.
1944+
:param int itersize: passed to a named_cursor, determines the number of rows fetched from PG at once.
1945+
"""
1946+
1947+
def __init__(self, cr, query, itersize=None):
1948+
self._ncr = None
1949+
self._cr = cr
1950+
self._tmp_tbl = "_upgrade_query_ids_{}".format(uuid.uuid4().hex)
1951+
cr.commit()
1952+
cr.execute(
1953+
format_query(
1954+
cr,
1955+
"CREATE UNLOGGED TABLE {}(id) AS (WITH query AS ({}) SELECT * FROM query)",
1956+
self._tmp_tbl,
1957+
SQLStr(query),
1958+
)
1959+
)
1960+
self._len = cr.rowcount
1961+
try:
1962+
cr.execute(
1963+
format_query(
1964+
cr,
1965+
"ALTER TABLE {} ADD CONSTRAINT {} PRIMARY KEY (id)",
1966+
self._tmp_tbl,
1967+
"pk_{}_id".format(self._tmp_tbl),
1968+
)
1969+
)
1970+
except psycopg2.IntegrityError as e:
1971+
if e.pgcode == errorcodes.UNIQUE_VIOLATION:
1972+
cr.rollback()
1973+
raise ValueError("The query for ids is producing duplicate values: {}".format(query))
1974+
raise
1975+
self._ncr = named_cursor(cr, itersize)
1976+
self._ncr.execute(format_query(cr, "SELECT id FROM {} ORDER BY id", self._tmp_tbl))
1977+
self._it = iter(self._ncr)
1978+
1979+
def _close(self):
1980+
if self._ncr:
1981+
if self._ncr.closed:
1982+
return
1983+
self._ncr.close()
1984+
self._cr.execute(format_query(self._cr, "DROP TABLE IF EXISTS {}", self._tmp_tbl))
1985+
1986+
def __len__(self):
1987+
return self._len
1988+
1989+
def __iter__(self):
1990+
return self
1991+
1992+
def __next__(self):
1993+
if self._ncr.closed:
1994+
raise StopIteration
1995+
try:
1996+
return next(self._it)[0]
1997+
except StopIteration:
1998+
self._close()
1999+
raise
2000+
2001+
def next(self):
2002+
return self.__next__()
2003+
2004+
def __enter__(self):
2005+
return self
2006+
2007+
def __exit__(self, exc_type, exc_value, traceback):
2008+
self._close()
2009+
return False
2010+
2011+
def __del__(self):
2012+
self._close()

0 commit comments

Comments
 (0)