django-bulk-hooks 0.2.60__py3-none-any.whl → 0.2.62__py3-none-any.whl

django_bulk_hooks/operations/bulk_executor.py

@@ -11,7 +11,12 @@ from django.db.models import AutoField, ForeignKey, Case, When, Value
 from django.db.models.constants import OnConflict
 from django.db.models.functions import Cast
 
-from django_bulk_hooks.operations.field_utils import get_field_value_for_db
+from django_bulk_hooks.operations.field_utils import (
+    get_field_value_for_db,
+    collect_auto_now_fields_for_inheritance_chain,
+    pre_save_auto_now_fields,
+)
+from django_bulk_hooks.helpers import tag_upsert_metadata
 
 logger = logging.getLogger(__name__)
 
@@ -42,6 +47,52 @@ class BulkExecutor:
         self.record_classifier = record_classifier
         self.model_cls = queryset.model
 
+    def _handle_upsert_metadata_tagging(
+        self, result_objects, objs, update_conflicts, unique_fields, existing_record_ids=None, existing_pks_map=None
+    ):
+        """
+
+        Handle classification and metadata tagging for upsert operations.
+
+
+
+        This centralizes the logic that was duplicated between MTI and non-MTI paths.
+
+
+
+        Args:
+
+            result_objects: List of objects returned from the bulk operation
+
+            objs: Original list of objects passed to bulk_create
+
+            update_conflicts: Whether this was an upsert operation
+
+            unique_fields: Fields used for conflict detection
+
+            existing_record_ids: Pre-classified existing record IDs (optional)
+
+            existing_pks_map: Pre-classified existing PK mapping (optional)
+
+
+
+        Returns:
+
+            None - modifies result_objects in place with metadata
+
+        """
+
+        if not (update_conflicts and unique_fields):
+            return
+
+        # Classify records if not already done
+
+        if existing_record_ids is None or existing_pks_map is None:
+            existing_record_ids, existing_pks_map = self.record_classifier.classify_for_upsert(objs, unique_fields)
+
+        # Tag the metadata
+        tag_upsert_metadata(result_objects, existing_record_ids, existing_pks_map)
+
     def bulk_create(
         self,
         objs,
@@ -77,7 +128,6 @@ class BulkExecutor:
 
         # Check if this is an MTI model and route accordingly
         if self.mti_handler.is_mti_model():
-
             # Use pre-classified records if provided, otherwise classify now
             if existing_record_ids is None or existing_pks_map is None:
                 existing_record_ids = set()
@@ -87,11 +137,13 @@ class BulkExecutor:
                     # This handles the schema migration case where parent exists but child doesn't
                     query_model = self.mti_handler.find_model_with_unique_fields(unique_fields)
                     logger.info(f"MTI upsert: querying {query_model.__name__} for unique fields {unique_fields}")
-                    
+
                     existing_record_ids, existing_pks_map = self.record_classifier.classify_for_upsert(
                         objs, unique_fields, query_model=query_model
                     )
-                    logger.info(f"MTI Upsert classification: {len(existing_record_ids)} existing, {len(objs) - len(existing_record_ids)} new")
+                    logger.info(
+                        f"MTI Upsert classification: {len(existing_record_ids)} existing, {len(objs) - len(existing_record_ids)} new"
+                    )
                     logger.info(f"existing_record_ids: {existing_record_ids}")
                     logger.info(f"existing_pks_map: {existing_pks_map}")
 
@@ -108,29 +160,20 @@ class BulkExecutor:
             # Execute the plan
             result = self._execute_mti_create_plan(plan)
 
-            # Tag objects with upsert metadata for hook dispatching
-            if update_conflicts and unique_fields:
-                self._tag_upsert_metadata(result, existing_record_ids, existing_pks_map)
-
-            return result
-
-        # Non-MTI model - use Django's native bulk_create
-        result = self._execute_bulk_create(
-            objs,
-            batch_size,
-            ignore_conflicts,
-            update_conflicts,
-            update_fields,
-            unique_fields,
-            **kwargs,
-        )
+        else:
+            # Non-MTI model - use Django's native bulk_create
+            result = self._execute_bulk_create(
+                objs,
+                batch_size,
+                ignore_conflicts,
+                update_conflicts,
+                update_fields,
+                unique_fields,
+                **kwargs,
+            )
 
-        # Tag objects with upsert metadata for hook dispatching
-        if update_conflicts and unique_fields:
-            # Use pre-classified results if available, otherwise classify now
-            if existing_record_ids is None:
-                existing_record_ids, existing_pks_map = self.record_classifier.classify_for_upsert(objs, unique_fields)
-            self._tag_upsert_metadata(result, existing_record_ids, existing_pks_map)
+        # Unified upsert metadata handling for both paths
+        self._handle_upsert_metadata_tagging(result, objs, update_conflicts, unique_fields, existing_record_ids, existing_pks_map)
 
         return result
 
@@ -185,23 +228,18 @@ class BulkExecutor:
         # Ensure auto_now fields are included and pre-saved for all models
         # This handles both MTI and non-MTI models uniformly (SOC & DRY)
         fields = list(fields)  # Make a copy so we can modify it
-        
+
         # Get models to check - for MTI, check entire inheritance chain
         if self.mti_handler.is_mti_model():
             models_to_check = self.mti_handler.get_inheritance_chain()
         else:
             models_to_check = [self.model_cls]
-        
-        # Collect all auto_now fields and pre-save them
-        auto_now_fields = set()
-        for model in models_to_check:
-            for field in model._meta.local_fields:
-                if getattr(field, "auto_now", False) and not getattr(field, "auto_now_add", False):
-                    auto_now_fields.add(field.name)
-                    # Pre-save the field to set the value on instances
-                    for obj in objs:
-                        field.pre_save(obj, add=False)
-        
+
+        # Use unified auto-now field handling
+        from django_bulk_hooks.operations.field_utils import handle_auto_now_fields_for_inheritance_chain
+
+        auto_now_fields = handle_auto_now_fields_for_inheritance_chain(models_to_check, objs, for_update=True)
+
         # Add auto_now fields to the update list if not already present
         for auto_now_field in auto_now_fields:
             if auto_now_field not in fields:
@@ -242,7 +280,6 @@ class BulkExecutor:
         if not plan:
             return []
 
-
         with transaction.atomic(using=self.queryset.db, savepoint=False):
             # Step 1: Upsert all parent objects level by level using Django's native upsert
             parent_instances_map = {}  # Maps original obj id() -> {model: parent_instance}
@@ -250,15 +287,15 @@ class BulkExecutor:
             for parent_level in plan.parent_levels:
                 # Use base QuerySet to avoid recursion
                 base_qs = BaseQuerySet(model=parent_level.model_class, using=self.queryset.db)
-                
+
                 # Build bulk_create kwargs
                 bulk_kwargs = {"batch_size": len(parent_level.objects)}
-                
+
                 if parent_level.update_conflicts:
                     # Let Django handle the upsert - it will INSERT or UPDATE as needed
                     bulk_kwargs["update_conflicts"] = True
                     bulk_kwargs["unique_fields"] = parent_level.unique_fields
-                    
+
                     # Filter update fields to only those that exist in this parent model
                     parent_model_fields = {field.name for field in parent_level.model_class._meta.local_fields}
                     filtered_update_fields = [field for field in parent_level.update_fields if field in parent_model_fields]
@@ -323,9 +360,7 @@ class BulkExecutor:
 
                 existing_child_pks = set()
                 if parent_pks_to_check:
-                    existing_child_pks = set(
-                        base_qs.filter(pk__in=parent_pks_to_check).values_list('pk', flat=True)
-                    )
+                    existing_child_pks = set(base_qs.filter(pk__in=parent_pks_to_check).values_list("pk", flat=True))
 
                 # Split based on whether child record exists
                 for child_obj in plan.child_objects:
@@ -364,59 +399,43 @@ class BulkExecutor:
                 # For MTI, we need to include the parent link (which is the PK)
                 filtered_fields = [f for f in opts.local_fields if not f.generated]
 
-                # FIX: Pass conflict resolution parameters to _batched_insert for MTI child tables
-                # Previously, _batched_insert was called without on_conflict/unique_fields/update_fields,
-                # causing IntegrityError when child tables have unique constraints during upsert operations.
-                # See: https://github.com/user/repo/issues/XXX
-                # Prepare conflict resolution parameters for upsert
+                # Prepare conflict resolution parameters for upsert using pre-computed fields
                 on_conflict = None
                 batched_unique_fields = None
                 batched_update_fields = None
 
-                if plan.update_conflicts:
-                    # Filter unique_fields and update_fields to only those on child model
-                    # Django's _batched_insert expects field objects, not field names
-                    child_model_fields_dict = {field.name: field for field in plan.child_model._meta.local_fields}
-
-                    # Unique fields may be on parent or child - filter to child only for child table insert
-                    # Convert field names to field objects
-                    if plan.unique_fields:
-                        batched_unique_fields = [
-                            child_model_fields_dict[fname]
-                            for fname in plan.unique_fields
-                            if fname in child_model_fields_dict
-                        ]
-
-                    # Update fields - filter to child only
-                    # Keep as strings - Django's _batched_insert accepts field name strings for update_fields
-                    if plan.update_fields:
-                        batched_update_fields = [
-                            fname
-                            for fname in plan.update_fields
-                            if fname in child_model_fields_dict
-                        ]
-
-                    # Only set on_conflict if we have unique fields for this table
-                    # Note: If unique_fields are all on parent, batched_unique_fields will be empty,
-                    # meaning no conflict resolution needed for child table
-                    if batched_unique_fields:
-                        if batched_update_fields:
-                            # We have both unique fields and update fields on child - use UPDATE
-                            on_conflict = OnConflict.UPDATE
-                        else:
-                            # We have unique fields on child but no update fields - use IGNORE
-                            # This handles the case where all update fields are on parent tables
-                            on_conflict = OnConflict.IGNORE
-                            # Clear batched_update_fields to avoid issues
-                            batched_update_fields = None
+                # Only set up upsert logic if we have child-specific unique fields
+                if plan.update_conflicts and plan.child_unique_fields:
+                    batched_unique_fields = plan.child_unique_fields
+                    batched_update_fields = plan.child_update_fields
+
+                    if batched_update_fields:
+                        # We have both unique fields and update fields on child - use UPDATE
+                        on_conflict = OnConflict.UPDATE
+                    else:
+                        # We have unique fields on child but no update fields - use IGNORE
+                        # This handles the case where all update fields are on parent tables
+                        on_conflict = OnConflict.IGNORE
+                        batched_update_fields = None
+
+                # Build kwargs for _batched_insert call
+                kwargs = {
+                    "batch_size": len(objs_without_pk),
+                }
+                # Only pass conflict resolution parameters if we have unique fields for this table
+                if batched_unique_fields:
+                    kwargs.update(
+                        {
+                            "on_conflict": on_conflict,
+                            "update_fields": batched_update_fields,
+                            "unique_fields": batched_unique_fields,
+                        }
+                    )
 
                 returned_columns = base_qs._batched_insert(
                     objs_without_pk,
                     filtered_fields,
-                    batch_size=len(objs_without_pk),
-                    on_conflict=on_conflict,
-                    update_fields=batched_update_fields,
-                    unique_fields=batched_unique_fields,
+                    **kwargs,
                 )
                 if returned_columns:
                     for obj, results in zip(objs_without_pk, returned_columns):
@@ -596,39 +615,3 @@ class BulkExecutor:
         from django.db.models import QuerySet
 
         return QuerySet.delete(self.queryset)
-
-    def _tag_upsert_metadata(self, result_objects, existing_record_ids, existing_pks_map):
-        """
-        Tag objects with metadata indicating whether they were created or updated.
-
-        This metadata is used by the coordinator to determine which hooks to fire.
-        The metadata is temporary and will be cleaned up after hook execution.
-
-        Args:
-            result_objects: List of objects returned from bulk operation
-            existing_record_ids: Set of id() for objects that existed before the operation
-            existing_pks_map: Dict mapping id(obj) -> pk for existing records
-        """
-        created_count = 0
-        updated_count = 0
-
-        # Create a set of PKs that existed before the operation
-        existing_pks = set(existing_pks_map.values())
-
-        for obj in result_objects:
-            # Use PK to determine if this record was created or updated
-            # If the PK was in the existing_pks_map, it was updated; otherwise created
-            was_created = obj.pk not in existing_pks
-            obj._bulk_hooks_was_created = was_created
-            obj._bulk_hooks_upsert_metadata = True
-
-            if was_created:
-                created_count += 1
-            else:
-                updated_count += 1
-
-        logger.info(
-            f"Tagged upsert metadata: {created_count} created, {updated_count} updated "
-            f"(total={len(result_objects)}, existing_pks={len(existing_pks)})"
-        )
-