fixup! [ADD] util/update_table_from_dict

Author: Pirols

src/base/tests/test_util.py

@@ -881,8 +881,8 @@ 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_table_from_dict(self):
-        TEST_TABLE_NAME = "_upgrade_update_table_from_dict_test_table"
+    def test_update_one_col_from_dict(self):
+        TEST_TABLE_NAME = "_upgrade_update_one_col_from_dict_test_table"
         N_ROWS = 10
 
         cr = self._get_cr()
@@ -905,8 +905,8 @@ def test_update_table_from_dict(self):
                 table=TEST_TABLE_NAME,
             )
         )
-        mapping = {id: {"col1": id * 2} for id in range(1, N_ROWS + 1, 2)}
-        util.update_table_from_dict(cr, TEST_TABLE_NAME, mapping)
+        mapping = {id: id * 2 for id in range(1, N_ROWS + 1, 2)}
+        util.update_table_from_dict(cr, TEST_TABLE_NAME, "col1", mapping)
 
         cr.execute(
             util.format_query(
@@ -933,7 +933,62 @@ def test_update_table_from_dict(self):
                 table=TEST_TABLE_NAME,
             )
         )
-        self.assertFalse(cr.rowcount)  # otherwise not all expected updates are performed
+        self.assertFalse(cr.rowcount)  # otherwise partial/incorrect updates are performed
+
+    def test_update_multiple_cols_from_dict(self):
+        TEST_TABLE_NAME = "_upgrade_update_multiple_cols_from_dict_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;
+                """
+                % N_ROWS,
+                table=TEST_TABLE_NAME,
+            )
+        )
+        mapping = {id: [id * 2, id * 3] for id in range(1, N_ROWS + 1, 2)}
+        util.update_table_from_dict(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)  # otherwise unintended column 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)  # otherwise 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)  # otherwise partial/incorrect updates are performed
 
     def test_create_column_with_fk(self):
         cr = self.env.cr

src/util/pg.py

@@ -1624,79 +1624,101 @@ def create_id_sequence(cr, table, set_as_default=True):
         )
 
 
-def update_table_from_dict(cr, table, mapping, key_col="id", bucket_size=DEFAULT_BUCKET_SIZE):
+def update_table_from_dict(cr, table, columns, mapping, key_col="id", bucket_size=DEFAULT_BUCKET_SIZE):
     """
-    Update table's rows based on mapping.
+    Update table based on mapping.
 
-    Efficiently updates rows in a table by mapping an identifier column (`key_col`) value to the new values for the provided set of columns.
+    Each `mapping` entry defines the new values for the specified `columns` for the row(s) whose `key_col` value matches the key.
+
+    .. important::
+
+       `columns` can be either a string or a list of strings. Crucially the values of the provided `mapping` must have the same type
+       and, if a list, the same **dimensionality** and **order**. Otherwise:
+
+       .. code-block:: python
+
+          columns = ["int_col", "text_col"]
+          mapping = {
+              1: [123, "foo", True],    # third value is ignored
+              2: [456],                 # text_col is set to NULL
+              3: ["bar", 789],          # will attempt to set `int_col` to "bar" and `text_col` to 789
+          }
 
     .. example::
 
        .. code-block:: python
 
+          # single column update
           util.update_table_from_dict(
               cr,
               "account_move",
+              "always_tax_eligible",
               {
-                  1: {"closing_return_id": 2, "always_tax_eligible": True},
-                  2: {"closing_return_id": 3, "always_tax_eligible": False},
+                  1: True,
+                  2: False,
               },
           )
 
-    :param str table: the table to update
-    :param dict[any, dict[str, any]] mapping: mapping of `key_col` identifiers to maps of column names to their new value
-
-                                              .. example::
-
-                                                 .. code-block:: python
-
-                                                    mapping = {
-                                                        1: {"col1": 123, "col2": "foo"},
-                                                        2: {"col1": 456, "col2": "bar"},
-                                                    }
-
-                                              .. warning::
-
-                                                 All maps should have the exact same set of keys (column names). The following
-                                                 example would behave unpredictably:
-
-                                                 .. code-block:: python
+          # multi-column update
+          util.update_table_from_dict(
+              cr,
+              "account_move",
+              ["closing_return_id", "always_tax_eligible"],
+              {
+                  1: [2, True],
+                  2: [3, False],
+              },
+          )
 
-                                                    # WRONG
-                                                    mapping = {
-                                                        1: {"col1": 123, "col2": "foo"},
-                                                        2: {"col1": 456},
-                                                    }
+    .. warning::
 
-                                                 Either resulting in `col2` updates being ignored or setting it to NULL for row 2.
+       As a side effect, the cursor may be committed.
 
-    :param str key_col: The column to match the key against (`id` by default)
+    :param str table: database's table to perform the update of
+    :param str | list[str] columns: table's columns to update
+    :param dict[any, any | list[any]] mapping: matches values of `key_col` to the new `columns` values
+    :param str key_col: column to match against keys of `mapping`
     :param int bucket_size: maximum number of rows to update per single query
     """
-    if not mapping:
-        return
-
     _validate_table(table)
+    if not columns or not mapping:
+        return
 
-    column_names = list(next(iter(mapping.values())).keys())
-    query = cr.mogrify(
-        format_query(
+    if isinstance(columns, str):
+        query = format_query(
             cr,
             """
             UPDATE {table} t
-               SET ({columns_list}) = ROW({values_list})
-              FROM JSONB_EACH(%%s) m
+               SET {col} = m.value::{col_type}
+              FROM JSONB_EACH_TEXT(%s) m
+             WHERE t.{key_col}::text = m.key
+            """,
+            table=table,
+            col=ColumnList.from_unquoted(cr, [columns]),
+            col_type=column_type(cr, table, columns),
+            key_col=key_col,
+        )
+    else:
+        query = format_query(
+            cr,
+            """
+            UPDATE {table} t
+               SET ({cols}) = ROW({cols_values})
+              FROM JSONB_EACH(%s) m
              WHERE t.{key_col}::varchar = m.key
             """,
             table=table,
-            columns_list=ColumnList.from_unquoted(cr, column_names),
-            values_list=sql.SQL(", ").join(
-                sql.SQL("(m.value->>%s)::{}").format(sql.SQL(column_type(cr, table, col))) for col in column_names
+            cols=ColumnList.from_unquoted(cr, columns),
+            cols_values=SQLStr(
+                ", ".join(
+                    "(m.value->>{:d})::{}".format(
+                        col_idx, sql.Identifier(column_type(cr, table, col_name)).as_string(cr._cnx)
+                    )
+                    for col_idx, col_name in enumerate(columns)
+                )
             ),
             key_col=key_col,
-        ),
-        column_names,
-    )
+        )
 
     if len(mapping) <= 1.1 * bucket_size:
         cr.execute(query, [json.dumps(mapping)])