Add insert_only to ColumnInfo

Author: bdowning

sql_athame/dataclasses.py

@@ -45,6 +45,7 @@ class ColumnInfo:
         constraints: Additional SQL constraints (e.g., 'UNIQUE', 'CHECK (value > 0)')
         serialize: Function to transform Python values before database storage
         deserialize: Function to transform database values back to Python objects
+        insert_only: Whether this field should only be set on INSERT, not UPDATE in upsert operations
 
     Example:
         >>> from dataclasses import dataclass
@@ -58,6 +59,7 @@ class ColumnInfo:
         ...     name: str
         ...     price: Annotated[float, ColumnInfo(constraints="CHECK (price > 0)")]
         ...     tags: Annotated[list, ColumnInfo(type="JSONB", serialize=json.dumps, deserialize=json.loads)]
+        ...     created_at: Annotated[datetime, ColumnInfo(insert_only=True)]
     """
 
     type: Optional[str] = None
@@ -69,6 +71,7 @@ class ColumnInfo:
 
     serialize: Optional[Callable[[Any], Any]] = None
     deserialize: Optional[Callable[[Any], Any]] = None
+    insert_only: Optional[bool] = None
 
     def __post_init__(self, constraints: Union[str, Iterable[str], None]) -> None:
         if constraints is not None:
@@ -94,6 +97,7 @@ def merge(a: "ColumnInfo", b: "ColumnInfo") -> "ColumnInfo":
             _constraints=(*a._constraints, *b._constraints),
             serialize=b.serialize if b.serialize is not None else a.serialize,
             deserialize=b.deserialize if b.deserialize is not None else a.deserialize,
+            insert_only=b.insert_only if b.insert_only is not None else a.insert_only,
         )
 
 
@@ -113,6 +117,7 @@ class ConcreteColumnInfo:
         constraints: Tuple of SQL constraint strings
         serialize: Optional serialization function
         deserialize: Optional deserialization function
+        insert_only: Whether this field should only be set on INSERT, not UPDATE
     """
 
     field: Field
@@ -123,6 +128,7 @@ class ConcreteColumnInfo:
     constraints: tuple[str, ...]
     serialize: Optional[Callable[[Any], Any]] = None
     deserialize: Optional[Callable[[Any], Any]] = None
+    insert_only: bool = False
 
     @staticmethod
     def from_column_info(
@@ -156,6 +162,7 @@ def from_column_info(
             constraints=info._constraints,
             serialize=info.serialize,
             deserialize=info.deserialize,
+            insert_only=bool(info.insert_only),
         )
 
     def create_table_string(self) -> str:
@@ -365,6 +372,20 @@ def field_names(cls, *, exclude: FieldNamesSet = ()) -> list[str]:
             if ci.field.name not in exclude
         ]
 
+    @classmethod
+    def insert_only_field_names(cls) -> set[str]:
+        """Get set of field names marked as insert_only in ColumnInfo.
+
+        Returns:
+            Set of field names that should only be set on INSERT, not UPDATE
+        """
+        return cls._cached(
+            ("insert_only_field_names",),
+            lambda: {
+                ci.field.name for ci in cls.column_info().values() if ci.insert_only
+            },
+        )
+
     @classmethod
     def field_names_sql(
         cls,
@@ -874,9 +895,15 @@ async def upsert(
             >>> 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'})
+
+        Note:
+            Fields marked with ColumnInfo(insert_only=True) are automatically
+            treated as insert-only and combined with the insert_only parameter.
         """
+        # Combine auto-detected insert_only fields with manual ones
+        all_insert_only = self.insert_only_field_names() | set(insert_only)
         # Combine exclude and insert_only for the UPDATE clause
-        update_exclude = set(exclude) | set(insert_only)
+        update_exclude = set(exclude) | all_insert_only
         query = sql(
             "{} RETURNING xmax",
             self.upsert_sql(self.insert_sql(exclude=exclude), exclude=update_exclude),
@@ -1193,19 +1220,26 @@ async def upsert_multiple(
 
         Example:
             >>> await User.upsert_multiple(pool, users, insert_only={'created_at'})
+
+        Note:
+            Fields marked with ColumnInfo(insert_only=True) are automatically
+            treated as insert-only and combined with the insert_only parameter.
         """
+        # Combine auto-detected insert_only fields with manual ones
+        all_insert_only = cls.insert_only_field_names() | set(insert_only)
+
         if cls.insert_multiple_mode == "executemany":
             await cls.upsert_multiple_executemany(
-                connection_or_pool, rows, insert_only=insert_only
+                connection_or_pool, rows, insert_only=all_insert_only
             )
             return "INSERT"
         elif cls.insert_multiple_mode == "array_safe":
             return await cls.upsert_multiple_array_safe(
-                connection_or_pool, rows, insert_only=insert_only
+                connection_or_pool, rows, insert_only=all_insert_only
             )
         else:
             return await cls.upsert_multiple_unnest(
-                connection_or_pool, rows, insert_only=insert_only
+                connection_or_pool, rows, insert_only=all_insert_only
             )
 
     @classmethod
@@ -1259,8 +1293,14 @@ async def plan_replace_multiple(
             ...     conn, new_users, where=sql("department_id = {}", dept_id)
             ... )
             >>> print(f"Will create {len(plan.created)}, update {len(plan.updated)}, delete {len(plan.deleted)}")
+
+        Note:
+            Fields marked with ColumnInfo(insert_only=True) are automatically
+            treated as insert-only and combined with the insert_only parameter.
         """
-        ignore = sorted(set(ignore) | set(insert_only))
+        # Combine auto-detected insert_only fields with manual ones
+        all_insert_only = cls.insert_only_field_names() | set(insert_only)
+        ignore = sorted(set(ignore) | all_insert_only)
         equal_ignoring = cls._cached(
             ("equal_ignoring", tuple(ignore)),
             lambda: cls._get_equal_ignoring_fn(ignore),
@@ -1283,7 +1323,7 @@ async def plan_replace_multiple(
 
         created = list(pending.values())
 
-        return ReplaceMultiplePlan(cls, insert_only, created, updated, deleted)
+        return ReplaceMultiplePlan(cls, all_insert_only, created, updated, deleted)
 
     @classmethod
     async def replace_multiple(
@@ -1315,6 +1355,10 @@ async def replace_multiple(
             >>> created, updated, deleted = await User.replace_multiple(
             ...     conn, new_users, where=sql("department_id = {}", dept_id)
             ... )
+
+        Note:
+            Fields marked with ColumnInfo(insert_only=True) are automatically
+            treated as insert-only and combined with the insert_only parameter.
         """
         plan = await cls.plan_replace_multiple(
             connection, rows, where=where, ignore=ignore, insert_only=insert_only
@@ -1380,8 +1424,14 @@ async def replace_multiple_reporting_differences(
             ... )
             >>> for old, new, fields in updates:
             ...     print(f"Updated {old.name}: changed {', '.join(fields)}")
+
+        Note:
+            Fields marked with ColumnInfo(insert_only=True) are automatically
+            treated as insert-only and combined with the insert_only parameter.
         """
-        ignore = sorted(set(ignore) | set(insert_only))
+        # Combine auto-detected insert_only fields with manual ones
+        all_insert_only = cls.insert_only_field_names() | set(insert_only)
+        ignore = sorted(set(ignore) | all_insert_only)
         differences_ignoring = cls._cached(
             ("differences_ignoring", tuple(ignore)),
             lambda: cls._get_differences_ignoring_fn(ignore),
@@ -1410,7 +1460,7 @@ async def replace_multiple_reporting_differences(
             await cls.upsert_multiple(
                 connection,
                 (*created, *(t[1] for t in updated_triples)),
-                insert_only=insert_only,
+                insert_only=all_insert_only,
             )
         if deleted:
             await cls.delete_multiple(connection, deleted)

tests/test_dataclasses.py

@@ -245,3 +245,120 @@ class Test(ModelBase, table_name="table"):
     assert Test.from_prepended_mapping(
         {"p_foo": "FOO", "p_bar": "BAR", "foo": "not foo", "other": "other"}, "p_"
     ) == Test("foo", "BAR")
+
+
+def test_insert_only_column_info():
+    """Test that ColumnInfo(insert_only=True) works correctly."""
+
+    @dataclass
+    class Test(ModelBase, table_name="table", primary_key="id"):
+        id: int
+        name: str
+        created_at: Annotated[str, ColumnInfo(insert_only=True)]
+        updated_at: str
+
+    # Test that insert_only fields are detected
+    insert_only_fields = Test.insert_only_field_names()
+    assert insert_only_fields == {"created_at"}
+
+    # Test that upsert SQL excludes insert_only fields from UPDATE when explicitly passed
+    test_instance = Test(1, "Alice", "2023-01-01", "2023-01-02")
+    insert_sql = test_instance.insert_sql()
+
+    # Pass the insert_only fields to upsert_sql to exclude them from UPDATE
+    upsert_sql = Test.upsert_sql(insert_sql, exclude=Test.insert_only_field_names())
+
+    # The upsert should include created_at in INSERT but exclude it from UPDATE
+    upsert_query = upsert_sql.query()[0]
+    assert "created_at" in upsert_query  # Should be in INSERT part
+    # The UPDATE SET clause should only include name and updated_at
+    assert '"name"=EXCLUDED."name"' in upsert_query
+    assert '"updated_at"=EXCLUDED."updated_at"' in upsert_query
+    assert '"created_at"=EXCLUDED."created_at"' not in upsert_query
+
+
+def test_insert_only_automatic_handling():
+    """Test that the upsert() method automatically handles insert_only fields."""
+
+    @dataclass
+    class Test(ModelBase, table_name="table", primary_key="id"):
+        id: int
+        name: str
+        created_at: Annotated[str, ColumnInfo(insert_only=True)]
+        updated_at: str
+
+    # Test automatic handling by checking the generated SQL from the upsert method
+    test_instance = Test(1, "Alice", "2023-01-01", "2023-01-02")
+
+    # Get the SQL that would be generated for upsert operation
+    # We simulate what happens inside the upsert method
+    all_insert_only = test_instance.insert_only_field_names()
+    insert_sql = test_instance.insert_sql()
+    upsert_sql = Test.upsert_sql(insert_sql, exclude=all_insert_only)
+
+    upsert_query = upsert_sql.query()[0]
+
+    # created_at should be excluded from UPDATE clause automatically
+    assert '"name"=EXCLUDED."name"' in upsert_query
+    assert '"updated_at"=EXCLUDED."updated_at"' in upsert_query
+    assert '"created_at"=EXCLUDED."created_at"' not in upsert_query
+
+
+def test_insert_only_merge_with_manual():
+    """Test that ColumnInfo insert_only merges with manual insert_only parameter."""
+
+    @dataclass
+    class Test(ModelBase, table_name="table", primary_key="id"):
+        id: int
+        name: str
+        created_at: Annotated[str, ColumnInfo(insert_only=True)]  # Auto insert-only
+        updated_at: str
+        version: int
+
+    # Verify auto-detected fields
+    assert Test.insert_only_field_names() == {"created_at"}
+
+    # Test combining auto-detected with manual insert_only
+    test_instance = Test(1, "Alice", "2023-01-01", "2023-01-02", 1)
+
+    # Simulate what happens in upsert methods - combine auto and manual
+    auto_insert_only = Test.insert_only_field_names()
+    manual_insert_only = {"version"}
+    all_insert_only = auto_insert_only | manual_insert_only
+
+    insert_sql = test_instance.insert_sql()
+    upsert_sql = Test.upsert_sql(insert_sql, exclude=all_insert_only)
+    upsert_query = upsert_sql.query()[0]
+
+    # Both created_at (auto) and version (manual) should be excluded from UPDATE
+    assert '"name"=EXCLUDED."name"' in upsert_query
+    assert '"updated_at"=EXCLUDED."updated_at"' in upsert_query
+    assert '"created_at"=EXCLUDED."created_at"' not in upsert_query
+    assert '"version"=EXCLUDED."version"' not in upsert_query
+
+
+def test_column_info_merge_insert_only():
+    """Test that ColumnInfo.merge handles insert_only properly."""
+
+    base_info = ColumnInfo(type="TEXT")
+    insert_only_info = ColumnInfo(insert_only=True)
+
+    # Test merging - insert_only should be preserved
+    merged = ColumnInfo.merge(base_info, insert_only_info)
+    assert merged.insert_only is True
+    assert merged.type == "TEXT"
+
+    # Test merging the other way
+    merged2 = ColumnInfo.merge(insert_only_info, base_info)
+    assert merged2.insert_only is True  # Should remain True
+    assert merged2.type == "TEXT"
+
+    # Test with both having insert_only set
+    both_false = ColumnInfo(type="INTEGER", insert_only=False)
+    merged3 = ColumnInfo.merge(both_false, insert_only_info)
+    assert merged3.insert_only is True  # True should take precedence
+
+    # Test with None (default)
+    none_info = ColumnInfo(type="BIGINT")
+    merged4 = ColumnInfo.merge(none_info, insert_only_info)
+    assert merged4.insert_only is True