[ADD] util/bulk_update_table

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.

closes #297

Signed-off-by: Christophe Simonis (chs) <[email protected]>

Author: Pirols

src/base/tests/test_util.py

@@ -883,6 +883,115 @@ def test_parallel_execute_retry_on_serialization_failure(self):
         cr.execute(util.format_query(cr, "SELECT 1 FROM {}", TEST_TABLE_NAME))
         self.assertFalse(cr.rowcount)
 
+    def test_update_one_col_from_dict(self):
+        TEST_TABLE_NAME = "_upgrade_bulk_update_one_col_test_table"
+        N_ROWS = 10
+
+        cr = self._get_cr()
+
+        cr.execute(
+            util.format_query(
+                cr,
+                """
+                DROP TABLE IF EXISTS {table};
+
+                CREATE TABLE {table} (
+                    id SERIAL PRIMARY KEY,
+                    col1 INTEGER,
+                    col2 INTEGER
+                );
+
+                INSERT INTO {table} (col1, col2) SELECT v, v FROM GENERATE_SERIES(1, %s) as v;
+                """,
+                table=TEST_TABLE_NAME,
+            ),
+            [N_ROWS],
+        )
+        mapping = {id: id * 2 for id in range(1, N_ROWS + 1, 2)}
+        util.bulk_update_table(cr, TEST_TABLE_NAME, "col1", mapping)
+
+        cr.execute(
+            util.format_query(
+                cr,
+                "SELECT id FROM {table} WHERE col2 != id",
+                table=TEST_TABLE_NAME,
+            )
+        )
+        self.assertFalse(cr.rowcount, "unintended column 'col2' is affected")
+
+        cr.execute(
+            util.format_query(
+                cr,
+                "SELECT id FROM {table} WHERE col1 != id AND MOD(id, 2) = 0",
+                table=TEST_TABLE_NAME,
+            )
+        )
+        self.assertFalse(cr.rowcount, "unintended rows are affected")
+
+        cr.execute(
+            util.format_query(
+                cr,
+                "SELECT id FROM {table} WHERE col1 != 2 * id AND MOD(id, 2) = 1",
+                table=TEST_TABLE_NAME,
+            )
+        )
+        self.assertFalse(cr.rowcount, "partial/incorrect updates are performed")
+
+    def test_update_multiple_cols_from_dict(self):
+        TEST_TABLE_NAME = "_upgrade_bulk_update_multiple_cols_test_table"
+        N_ROWS = 10
+
+        cr = self._get_cr()
+
+        cr.execute(
+            util.format_query(
+                cr,
+                """
+                DROP TABLE IF EXISTS {table};
+
+                CREATE TABLE {table} (
+                    id SERIAL PRIMARY KEY,
+                    col1 INTEGER,
+                    col2 INTEGER,
+                    col3 INTEGER
+                );
+
+                INSERT INTO {table} (col1, col2, col3) SELECT v, v, v FROM GENERATE_SERIES(1, %s) as v;
+                """,
+                table=TEST_TABLE_NAME,
+            ),
+            [N_ROWS],
+        )
+        mapping = {id: [id * 2, id * 3] for id in range(1, N_ROWS + 1, 2)}
+        util.bulk_update_table(cr, TEST_TABLE_NAME, ["col1", "col2"], mapping)
+
+        cr.execute(
+            util.format_query(
+                cr,
+                "SELECT id FROM {table} WHERE col3 != id",
+                table=TEST_TABLE_NAME,
+            )
+        )
+        self.assertFalse(cr.rowcount, "unintended column 'col3' is affected")
+
+        cr.execute(
+            util.format_query(
+                cr,
+                "SELECT id FROM {table} WHERE col1 != id AND MOD(id, 2) = 0",
+                table=TEST_TABLE_NAME,
+            )
+        )
+        self.assertFalse(cr.rowcount, "unintended rows are affected")
+
+        cr.execute(
+            util.format_query(
+                cr,
+                "SELECT id FROM {table} WHERE (col1 != 2 * id OR col2 != 3 * id) AND MOD(id, 2) = 1",
+                table=TEST_TABLE_NAME,
+            )
+        )
+        self.assertFalse(cr.rowcount, "partial/incorrect updates are performed")
+
     def test_create_column_with_fk(self):
         cr = self.env.cr
         self.assertFalse(util.column_exists(cr, "res_partner", "_test_lang_id"))

src/util/pg.py

@@ -32,6 +32,7 @@
 import psycopg2
 from psycopg2 import errorcodes, sql
 from psycopg2.extensions import quote_ident
+from psycopg2.extras import Json
 
 try:
     from odoo.modules import module as odoo_module
@@ -1698,3 +1699,76 @@ def create_id_sequence(cr, table, set_as_default=True):
                 table=table_sql,
             )
         )
+
+
+def bulk_update_table(cr, table, columns, mapping, key_col="id"):
+    """
+    Update table based on mapping.
+
+    Each `mapping` entry defines the new values for the specified `columns` for the row(s)
+    whose `key_col` value matches the key.
+
+    .. example::
+
+       .. code-block:: python
+
+          # single column update
+          util.bulk_update_table(cr, "res_users", "active", {42: False, 27: True})
+
+          # multi-column update
+          util.bulk_update_table(
+              cr,
+              "res_users",
+              ["active", "password"],
+              {
+                  "admin": [True, "1234"],
+                  "demo": [True, "5678"],
+              },
+              key_col="login",
+          )
+
+    :param str table: table to update.
+    :param str | list(str) columns: columns spec for the update. It could be a single
+                                    column name or a list of column names. The `mapping`
+                                    must match the spec.
+    :param dict mapping: values to set, which must match the spec in `columns`,
+                         following the **same** order
+    :param str key_col: column used as key to get the values from `mapping` during the
+                        update.
+
+    .. warning::
+
+       The values in the mapping will be casted to the type of the target column.
+       This function is designed to update scalar values, avoid setting arrays or json
+       data via the mapping.
+    """
+    _validate_table(table)
+    if not columns or not mapping:
+        return
+
+    assert isinstance(mapping, dict)
+    if isinstance(columns, str):
+        columns = [columns]
+    else:
+        n_columns = len(columns)
+        assert all(isinstance(value, (list, tuple)) and len(value) == n_columns for value in mapping.values())
+
+    query = format_query(
+        cr,
+        """
+        UPDATE {table} t
+           SET ({cols}) = ROW({cols_values})
+          FROM JSONB_EACH(%s) m
+         WHERE t.{key_col}::text = m.key
+        """,
+        table=table,
+        cols=ColumnList.from_unquoted(cr, columns),
+        cols_values=SQLStr(
+            ", ".join(
+                "(m.value->>{:d})::{}".format(col_idx, column_type(cr, table, col_name))
+                for col_idx, col_name in enumerate(columns)
+            )
+        ),
+        key_col=key_col,
+    )
+    cr.execute(query, [Json(mapping)])