Add insert_only to upsert

Author: bdowning

sql_athame/dataclasses.py

@@ -855,20 +855,31 @@ def upsert_sql(cls, insert_sql: Fragment, exclude: FieldNamesSet = ()) -> Fragme
         return Fragment([insert_sql, cached])
 
     async def upsert(
-        self, connection_or_pool: Union[Connection, Pool], exclude: FieldNamesSet = ()
+        self,
+        connection_or_pool: Union[Connection, Pool],
+        exclude: FieldNamesSet = (),
+        insert_only: FieldNamesSet = (),
     ) -> bool:
         """Insert or update this instance in the database.
 
         Args:
             connection_or_pool: Database connection or pool
             exclude: Field names to exclude from the UPDATE clause
+            insert_only: Field names that should only be set on INSERT, not UPDATE
 
         Returns:
             True if the record was updated, False if it was inserted
+
+        Example:
+            >>> user = User(id=1, name="Alice", created_at=datetime.now())
+            >>> # Only set created_at on INSERT, not UPDATE
+            >>> was_updated = await user.upsert(pool, insert_only={'created_at'})
         """
+        # Combine exclude and insert_only for the UPDATE clause
+        update_exclude = set(exclude) | set(insert_only)
         query = sql(
             "{} RETURNING xmax",
-            self.upsert_sql(self.insert_sql(exclude=exclude), exclude=exclude),
+            self.upsert_sql(self.insert_sql(exclude=exclude), exclude=update_exclude),
         )
         result = await connection_or_pool.fetchrow(*query)
         return result["xmax"] != 0

tests/test_asyncpg.py

@@ -321,3 +321,44 @@ class Test(ModelBase, table_name="table", primary_key="id"):
     await Test.insert_multiple(conn, [])
 
     assert list(await Test.select(conn)) == []
+
+
+async def test_upsert_insert_only(conn):
+    @dataclass
+    class Test(ModelBase, table_name="test_upsert", primary_key="id"):
+        id: int
+        name: str
+        count: int
+        created_at: str
+
+    await conn.execute(*Test.create_table_sql())
+
+    # Initial insert
+    record = Test(1, "Alice", 5, "2023-01-01")
+    was_updated = await record.upsert(conn)
+    assert not was_updated  # Should be False for initial insert
+
+    # Verify record was inserted
+    result = await Test.select(conn, where=sql("id = {}", 1))
+    assert len(result) == 1
+    assert result[0] == record
+
+    # Update without insert_only - should update all fields including created_at
+    updated_record = Test(1, "Alice Updated", 10, "2023-01-02")
+    was_updated = await updated_record.upsert(conn)
+    assert was_updated  # Should be True for update
+
+    result = await Test.select(conn, where=sql("id = {}", 1))
+    assert len(result) == 1
+    assert result[0] == updated_record
+
+    # Update with insert_only - should not update created_at
+    final_record = Test(1, "Alice Final", 15, "2023-01-03")
+    was_updated = await final_record.upsert(conn, insert_only={"created_at"})
+    assert was_updated  # Should be True for update
+
+    result = await Test.select(conn, where=sql("id = {}", 1))
+    assert len(result) == 1
+    # created_at should still be the old value, other fields should be updated
+    expected = Test(1, "Alice Final", 15, "2023-01-02")  # created_at unchanged
+    assert result[0] == expected

tests/test_dataclasses.py

@@ -166,6 +166,37 @@ class Test(ModelBase, table_name="table", primary_key="id"):
     ]
 
 
+def test_upsert_insert_only():
+    @dataclass
+    class Test(ModelBase, table_name="table", primary_key="id"):
+        id: int
+        foo: int
+        bar: str
+        created_at: str
+
+    t = Test(1, 42, "str", "2023-01-01")
+
+    # Test with insert_only parameter - created_at should be excluded from UPDATE
+    assert list(t.upsert_sql(t.insert_sql(), exclude={"created_at"})) == [
+        'INSERT INTO "table" ("id", "foo", "bar", "created_at") VALUES ($1, $2, $3, $4) '
+        'ON CONFLICT ("id") DO UPDATE SET "foo"=EXCLUDED."foo", "bar"=EXCLUDED."bar"',
+        1,
+        42,
+        "str",
+        "2023-01-01",
+    ]
+
+    # Test with both exclude and insert_only-style exclude
+    assert list(t.upsert_sql(t.insert_sql(), exclude={"bar", "created_at"})) == [
+        'INSERT INTO "table" ("id", "foo", "bar", "created_at") VALUES ($1, $2, $3, $4) '
+        'ON CONFLICT ("id") DO UPDATE SET "foo"=EXCLUDED."foo"',
+        1,
+        42,
+        "str",
+        "2023-01-01",
+    ]
+
+
 def test_serial():
     @dataclass
     class Test(ModelBase, table_name="table", primary_key="id"):