Add insert_only fields to upsert/replace_multiple

Author: bdowning

sql_athame/dataclasses.py

@@ -458,31 +458,46 @@ async def insert_multiple(
 
     @classmethod
     async def upsert_multiple_unnest(
-        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+        cls: Type[T],
+        connection_or_pool: Union[Connection, Pool],
+        rows: Iterable[T],
+        insert_only: FieldNamesSet = (),
     ) -> str:
         return await connection_or_pool.execute(
-            *cls.upsert_sql(cls.insert_multiple_sql(rows))
+            *cls.upsert_sql(cls.insert_multiple_sql(rows), exclude=insert_only)
         )
 
     @classmethod
     async def upsert_multiple_array_safe(
-        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+        cls: Type[T],
+        connection_or_pool: Union[Connection, Pool],
+        rows: Iterable[T],
+        insert_only: FieldNamesSet = (),
     ) -> str:
         last = ""
         for chunk in chunked(rows, 100):
             last = await connection_or_pool.execute(
-                *cls.upsert_sql(cls.insert_multiple_array_safe_sql(chunk))
+                *cls.upsert_sql(
+                    cls.insert_multiple_array_safe_sql(chunk), exclude=insert_only
+                )
             )
         return last
 
     @classmethod
     async def upsert_multiple(
-        cls: Type[T], connection_or_pool: Union[Connection, Pool], rows: Iterable[T]
+        cls: Type[T],
+        connection_or_pool: Union[Connection, Pool],
+        rows: Iterable[T],
+        insert_only: FieldNamesSet = (),
     ) -> str:
         if cls.array_safe_insert:
-            return await cls.upsert_multiple_array_safe(connection_or_pool, rows)
+            return await cls.upsert_multiple_array_safe(
+                connection_or_pool, rows, insert_only=insert_only
+            )
         else:
-            return await cls.upsert_multiple_unnest(connection_or_pool, rows)
+            return await cls.upsert_multiple_unnest(
+                connection_or_pool, rows, insert_only=insert_only
+            )
 
     @classmethod
     def _get_equal_ignoring_fn(
@@ -505,9 +520,11 @@ async def replace_multiple(
         *,
         where: Where,
         ignore: FieldNamesSet = (),
+        insert_only: FieldNamesSet = (),
     ) -> Tuple[List[T], List[T], List[T]]:
+        ignore = sorted(set(ignore) | set(insert_only))
         equal_ignoring = cls._cached(
-            ("equal_ignoring", tuple(sorted(ignore))),
+            ("equal_ignoring", tuple(ignore)),
             lambda: cls._get_equal_ignoring_fn(ignore),
         )
         pending = {row.primary_key(): row for row in map(cls.ensure_model, rows)}
@@ -529,7 +546,9 @@ async def replace_multiple(
         created = list(pending.values())
 
         if created or updated:
-            await cls.upsert_multiple(connection, (*created, *updated))
+            await cls.upsert_multiple(
+                connection, (*created, *updated), insert_only=insert_only
+            )
         if deleted:
             await cls.delete_multiple(connection, deleted)
 
@@ -561,9 +580,11 @@ async def replace_multiple_reporting_differences(
         *,
         where: Where,
         ignore: FieldNamesSet = (),
+        insert_only: FieldNamesSet = (),
     ) -> Tuple[List[T], List[Tuple[T, T, List[str]]], List[T]]:
+        ignore = sorted(set(ignore) | set(insert_only))
         differences_ignoring = cls._cached(
-            ("differences_ignoring", tuple(sorted(ignore))),
+            ("differences_ignoring", tuple(ignore)),
             lambda: cls._get_differences_ignoring_fn(ignore),
         )
 
@@ -588,7 +609,9 @@ async def replace_multiple_reporting_differences(
 
         if created or updated_triples:
             await cls.upsert_multiple(
-                connection, (*created, *(t[1] for t in updated_triples))
+                connection,
+                (*created, *(t[1] for t in updated_triples)),
+                insert_only=insert_only,
             )
         if deleted:
             await cls.delete_multiple(connection, deleted)

tests/test_asyncpg.py

@@ -1,6 +1,8 @@
+import asyncio
 import json
 import os
 from dataclasses import dataclass, field
+from datetime import datetime
 from typing import Optional
 
 import asyncpg
@@ -93,6 +95,37 @@ class Test(ModelBase, table_name="test", primary_key="id"):
     ]
 
 
+async def test_replace_multiple_ignore_insert_only(conn):
+    @dataclass(order=True)
+    class Test(ModelBase, table_name="test", primary_key="id"):
+        id: int
+        a: int
+        created: datetime = field(default_factory=datetime.utcnow)
+        updated: datetime = field(default_factory=datetime.utcnow)
+
+    await conn.execute(*Test.create_table_sql())
+
+    data = [Test(1, 1), Test(2, 1), Test(3, 2)]
+    await Test.insert_multiple(conn, data)
+
+    await asyncio.sleep(0.1)
+    new_data = [Test(1, 1), Test(2, 4), Test(3, 2)]
+    c, u, d = await Test.replace_multiple(
+        conn, new_data, where=[], ignore=["updated"], insert_only=["created"]
+    )
+    assert not c and not d
+    assert len(u) == 1
+
+    db_data = await Test.select(conn, order_by="id")
+    assert data[0] == db_data[0]
+    assert data[2] == db_data[2]
+    orig = data[1]
+    new = new_data[1]
+    db = db_data[1]
+    assert db.created == orig.created
+    assert db.updated == new.updated
+
+
 async def test_replace_multiple_arrays(conn):
     @dataclass(order=True)
     class Test(