Skip to content

Commit f67f74d

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 09d18f5 commit f67f74d

File tree

1 file changed

+46
-0
lines changed

1 file changed

+46
-0
lines changed

src/util/pg.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1932,3 +1932,49 @@ 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+
def __init__(self, cr, query, itersize=None):
1939+
self._cr = cr
1940+
self._tmp_tbl = "_upgrade_query_ids_{}".format(uuid.uuid4().hex)
1941+
cr.execute(
1942+
format_query(
1943+
cr,
1944+
"CREATE UNLOGGED TABLE {}(id) AS (WITH query AS ({}) SELECT * FROM query)",
1945+
self._tmp_tbl,
1946+
SQLStr(query),
1947+
)
1948+
)
1949+
self._len = cr.rowcount
1950+
cr.execute(
1951+
format_query(
1952+
cr, "ALTER TABLE {} ADD CONSTRAINT {} PRIMARY KEY (id)", self._tmp_tbl, "pk_{}_id".format(self._tmp_tbl)
1953+
)
1954+
)
1955+
self._ncr = named_cursor(cr, itersize)
1956+
self._ncr.execute(format_query(cr, "SELECT id FROM {} ORDER BY id", self._tmp_tbl))
1957+
1958+
def _drop_tmp_tbl(self):
1959+
if self._cr:
1960+
self._cr.execute(format_query(self._cr, "DROP TABLE IF EXISTS {}", self._tmp_tbl))
1961+
self._cr = None
1962+
1963+
def __len__(self):
1964+
return self._len
1965+
1966+
def __iter__(self):
1967+
return (id_ for (id_,) in self._ncr)
1968+
1969+
def __enter__(self):
1970+
return self
1971+
1972+
def __exit__(self, exc_type, exc_value, traceback):
1973+
self._ncr.__exit__(exc_type, exc_value, traceback)
1974+
self._drop_tmp_tbl()
1975+
return False
1976+
1977+
def __del__(self):
1978+
self._ncr.close()
1979+
del self._ncr
1980+
self._drop_tmp_tbl()

0 commit comments

Comments
 (0)