Skip to content

Commit 20f04c9

Browse files
PirolsPirols
committed
[ADD] util/update_table_from_dict
A recurrent challenge in writing upgrade scripts is that of updating values in a table based on some form of already available mapping from the id (or another identifier) to the new value, this is often addressed with an iterative solution in the form: ```python for key, value in mapping.items(): cr.execute( """ UPDATE table SET col = %s WHERE key_col = %s """, [value, key], ) ``` or in a more efficient (only issuing a single query) but hacky way: ```python cr.execute( """ UPDATE table SET col = (%s::jsonb)->>(key_col::text) WHERE key_col = ANY(%s) """, [json.dumps(mapping), list(mapping)], ) ``` With the former being ineffective for big mappings and the latter often requiring some comments at review time to get it right. This commit introduces a util meant to make it easier to efficiently perform such updates.
1 parent 6c9a30c commit 20f04c9

File tree

2 files changed

+208
-1
lines changed

2 files changed

+208
-1
lines changed

‎src/base/tests/test_util.py

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -881,6 +881,115 @@ def test_parallel_execute_retry_on_serialization_failure(self):
881881
cr.execute(util.format_query(cr, "SELECT 1 FROM {}", TEST_TABLE_NAME))
882882
self.assertFalse(cr.rowcount)
883883

884+
def test_update_one_col_from_dict(self):
885+
TEST_TABLE_NAME = "_upgrade_update_one_col_from_dict_test_table"
886+
N_ROWS = 10
887+
888+
cr = self._get_cr()
889+
890+
cr.execute(
891+
util.format_query(
892+
cr,
893+
"""
894+
DROP TABLE IF EXISTS {table};
895+
896+
CREATE TABLE {table} (
897+
id SERIAL PRIMARY KEY,
898+
col1 INTEGER,
899+
col2 INTEGER
900+
);
901+
902+
INSERT INTO {table} (col1, col2) SELECT v, v FROM GENERATE_SERIES(1, %s) as v;
903+
"""
904+
% N_ROWS,
905+
table=TEST_TABLE_NAME,
906+
)
907+
)
908+
mapping = {id: id * 2 for id in range(1, N_ROWS + 1, 2)}
909+
util.update_table_from_dict(cr, TEST_TABLE_NAME, "col1", mapping)
910+
911+
cr.execute(
912+
util.format_query(
913+
cr,
914+
"SELECT id FROM {table} WHERE col2 != id",
915+
table=TEST_TABLE_NAME,
916+
)
917+
)
918+
self.assertFalse(cr.rowcount) # otherwise unintended column is affected
919+
920+
cr.execute(
921+
util.format_query(
922+
cr,
923+
"SELECT id FROM {table} WHERE col1 != id AND MOD(id, 2) = 0",
924+
table=TEST_TABLE_NAME,
925+
)
926+
)
927+
self.assertFalse(cr.rowcount) # otherwise unintended rows are affected
928+
929+
cr.execute(
930+
util.format_query(
931+
cr,
932+
"SELECT id FROM {table} WHERE col1 != 2 * id AND MOD(id, 2) = 1",
933+
table=TEST_TABLE_NAME,
934+
)
935+
)
936+
self.assertFalse(cr.rowcount) # otherwise partial/incorrect updates are performed
937+
938+
def test_update_multiple_cols_from_dict(self):
939+
TEST_TABLE_NAME = "_upgrade_update_multiple_cols_from_dict_test_table"
940+
N_ROWS = 10
941+
942+
cr = self._get_cr()
943+
944+
cr.execute(
945+
util.format_query(
946+
cr,
947+
"""
948+
DROP TABLE IF EXISTS {table};
949+
950+
CREATE TABLE {table} (
951+
id SERIAL PRIMARY KEY,
952+
col1 INTEGER,
953+
col2 INTEGER,
954+
col3 INTEGER
955+
);
956+
957+
INSERT INTO {table} (col1, col2, col3) SELECT v, v, v FROM GENERATE_SERIES(1, %s) as v;
958+
"""
959+
% N_ROWS,
960+
table=TEST_TABLE_NAME,
961+
)
962+
)
963+
mapping = {id: [id * 2, id * 3] for id in range(1, N_ROWS + 1, 2)}
964+
util.update_table_from_dict(cr, TEST_TABLE_NAME, ["col1", "col2"], mapping)
965+
966+
cr.execute(
967+
util.format_query(
968+
cr,
969+
"SELECT id FROM {table} WHERE col3 != id",
970+
table=TEST_TABLE_NAME,
971+
)
972+
)
973+
self.assertFalse(cr.rowcount) # otherwise unintended column is affected
974+
975+
cr.execute(
976+
util.format_query(
977+
cr,
978+
"SELECT id FROM {table} WHERE col1 != id AND MOD(id, 2) = 0",
979+
table=TEST_TABLE_NAME,
980+
)
981+
)
982+
self.assertFalse(cr.rowcount) # otherwise unintended rows are affected
983+
984+
cr.execute(
985+
util.format_query(
986+
cr,
987+
"SELECT id FROM {table} WHERE (col1 != 2 * id OR col2 != 3 * id) AND MOD(id, 2) = 1",
988+
table=TEST_TABLE_NAME,
989+
)
990+
)
991+
self.assertFalse(cr.rowcount) # otherwise partial/incorrect updates are performed
992+
884993
def test_create_column_with_fk(self):
885994
cr = self.env.cr
886995
self.assertFalse(util.column_exists(cr, "res_partner", "_test_lang_id"))

‎src/util/pg.py

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
"""Utility functions for interacting with PostgreSQL."""
33

44
import collections
5+
import json
56
import logging
67
import os
78
import re
@@ -43,7 +44,7 @@
4344

4445
from .exceptions import MigrationError, SleepyDeveloperError
4546
from .helpers import _validate_table, model_of_table
46-
from .misc import Sentinel, log_progress, version_gte
47+
from .misc import Sentinel, chunks, log_progress, version_gte
4748

4849
_logger = logging.getLogger(__name__)
4950

@@ -1621,3 +1622,100 @@ def create_id_sequence(cr, table, set_as_default=True):
16211622
table=table_sql,
16221623
)
16231624
)
1625+
1626+
1627+
def update_table_from_dict(cr, table, columns, mapping, key_col="id", bucket_size=DEFAULT_BUCKET_SIZE):
1628+
"""
1629+
Update table based on mapping.
1630+
1631+
Each `mapping` entry defines the new values for the specified `columns` for the row(s) whose `key_col` value matches the key.
1632+
1633+
.. important::
1634+
1635+
`columns` can be either a string or a list of strings. Crucially the values of the provided `mapping` must have the same
1636+
**type** and, if a list, the same **dimensionality** and **order**. Otherwise:
1637+
1638+
.. code-block:: python
1639+
1640+
columns = ["int_col", "text_col"]
1641+
mapping = {
1642+
1: [111, "foo", True], # third value is ignored
1643+
2: [222], # text_col is set to NULL
1644+
3: ["bar", 333], # will attempt to set `int_col` to "bar" and `text_col` to 789
1645+
4: 444, # text_col is set to NULL
1646+
}
1647+
1648+
.. example::
1649+
1650+
.. code-block:: python
1651+
1652+
# single column update
1653+
util.update_table_from_dict(
1654+
cr,
1655+
"account_move",
1656+
"always_tax_eligible",
1657+
{
1658+
1: True,
1659+
2: False,
1660+
},
1661+
)
1662+
1663+
# multi-column update
1664+
util.update_table_from_dict(
1665+
cr,
1666+
"account_move",
1667+
["closing_return_id", "always_tax_eligible"],
1668+
{
1669+
1: [2, True],
1670+
2: [3, False],
1671+
},
1672+
)
1673+
1674+
.. warning::
1675+
1676+
As a side effect, the cursor may be committed.
1677+
1678+
:param str table: database's table to perform the update of
1679+
:param str | list[str] columns: table's columns to update
1680+
:param dict[any, any | list[any]] mapping: matches values of `key_col` to the new `columns` values
1681+
:param str key_col: column to match against keys of `mapping`
1682+
:param int bucket_size: maximum number of rows to update per single query
1683+
"""
1684+
_validate_table(table)
1685+
if not columns or not mapping:
1686+
return
1687+
1688+
if isinstance(columns, str):
1689+
columns = [columns]
1690+
1691+
query = format_query(
1692+
cr,
1693+
"""
1694+
UPDATE {table} t
1695+
SET ({cols}) = ROW({cols_values})
1696+
FROM JSONB_EACH(%s) m
1697+
WHERE t.{key_col}::varchar = m.key
1698+
""",
1699+
table=table,
1700+
cols=ColumnList.from_unquoted(cr, columns),
1701+
cols_values=SQLStr(
1702+
", ".join(
1703+
"(m.value->>{:d})::{}".format(
1704+
col_idx, sql.Identifier(column_type(cr, table, col_name)).as_string(cr._cnx)
1705+
)
1706+
for col_idx, col_name in enumerate(columns)
1707+
)
1708+
),
1709+
key_col=key_col,
1710+
)
1711+
1712+
if len(mapping) <= 1.1 * bucket_size:
1713+
cr.execute(query, [json.dumps(mapping)])
1714+
else:
1715+
parallel_execute(
1716+
cr,
1717+
[
1718+
cr.mogrify(query, [json.dumps(mapping_chunk)]).decode()
1719+
for mapping_chunk in chunks(mapping.items(), bucket_size, fmt=dict)
1720+
],
1721+
)

0 commit comments

Comments
 (0)