django-bulk-hooks 0.2.15__py3-none-any.whl → 0.2.17__py3-none-any.whl

django_bulk_hooks/operations/mti_handler.py

@@ -8,6 +8,7 @@ It returns plans (data structures) that the BulkExecutor executes.
 """
 
 import logging
+
 from django.db.models import AutoField
 
 logger = logging.getLogger(__name__)
@@ -121,6 +122,8 @@ class MTIHandler:
         update_conflicts=False,
         unique_fields=None,
         update_fields=None,
+        existing_record_ids=None,
+        existing_pks_map=None,
     ):
         """
         Build an execution plan for bulk creating MTI model instances.
@@ -134,21 +137,36 @@ class MTIHandler:
             update_conflicts: Enable UPSERT on conflict
             unique_fields: Fields for conflict detection
             update_fields: Fields to update on conflict
+            existing_record_ids: Set of id() for objects that exist in DB (from RecordClassifier)
+            existing_pks_map: Dict mapping id(obj) -> pk for existing records (from RecordClassifier)
             
         Returns:
             MTICreatePlan object
         """
-        from django_bulk_hooks.operations.mti_plans import MTICreatePlan, ParentLevel
-        
+        from django_bulk_hooks.operations.mti_plans import MTICreatePlan
+
         if not objs:
             return None
-        
+
         inheritance_chain = self.get_inheritance_chain()
         if len(inheritance_chain) <= 1:
             raise ValueError("build_create_plan called on non-MTI model")
-        
+
         batch_size = batch_size or len(objs)
-        
+
+        # Use provided classification (no more DB query here!)
+        if existing_record_ids is None:
+            existing_record_ids = set()
+        if existing_pks_map is None:
+            existing_pks_map = {}
+
+        # Set PKs on existing objects so they can be updated
+        if existing_pks_map:
+            for obj in objs:
+                if id(obj) in existing_pks_map:
+                    obj.pk = existing_pks_map[id(obj)]
+                    obj.id = existing_pks_map[id(obj)]
+
         # Build parent levels
         parent_levels = self._build_parent_levels(
             objs,
@@ -157,13 +175,13 @@ class MTIHandler:
             unique_fields=unique_fields,
             update_fields=update_fields,
         )
-        
+
         # Build child object templates (without parent links - executor adds them)
         child_objects = []
         for obj in objs:
             child_obj = self._create_child_instance_template(obj, inheritance_chain[-1])
             child_objects.append(child_obj)
-        
+
         return MTICreatePlan(
             inheritance_chain=inheritance_chain,
             parent_levels=parent_levels,
@@ -171,6 +189,10 @@ class MTIHandler:
             child_model=inheritance_chain[-1],
             original_objects=objs,
             batch_size=batch_size,
+            existing_record_ids=existing_record_ids,
+            update_conflicts=update_conflicts,
+            unique_fields=unique_fields or [],
+            update_fields=update_fields or [],
         )
 
     def _build_parent_levels(
@@ -190,38 +212,38 @@ class MTIHandler:
             List of ParentLevel objects
         """
         from django_bulk_hooks.operations.mti_plans import ParentLevel
-        
+
         parent_levels = []
         parent_instances_map = {}  # Maps obj id() -> {model_class: parent_instance}
-        
+
         for level_idx, model_class in enumerate(inheritance_chain[:-1]):
             parent_objs_for_level = []
-            
+
             for obj in objs:
                 # Get current parent from previous level
                 current_parent = None
                 if level_idx > 0:
                     prev_parents = parent_instances_map.get(id(obj), {})
                     current_parent = prev_parents.get(inheritance_chain[level_idx - 1])
-                
+
                 # Create parent instance
                 parent_obj = self._create_parent_instance(obj, model_class, current_parent)
                 parent_objs_for_level.append(parent_obj)
-                
+
                 # Store in map
                 if id(obj) not in parent_instances_map:
                     parent_instances_map[id(obj)] = {}
                 parent_instances_map[id(obj)][model_class] = parent_obj
-            
+
             # Determine upsert parameters for this level
             level_update_conflicts = False
             level_unique_fields = []
             level_update_fields = []
-            
+
             if update_conflicts and unique_fields:
                 # Filter unique_fields and update_fields to only those in this model
                 model_fields_by_name = {f.name: f for f in model_class._meta.local_fields}
-                
+
                 # Normalize unique fields
                 normalized_unique = []
                 for uf in unique_fields or []:
@@ -229,19 +251,19 @@ class MTIHandler:
                         normalized_unique.append(uf)
                     elif uf.endswith("_id") and uf[:-3] in model_fields_by_name:
                         normalized_unique.append(uf[:-3])
-                
+
                 # Check if this model has a matching constraint
                 if normalized_unique and self._has_matching_constraint(model_class, normalized_unique):
                     # Filter update fields
                     filtered_updates = [
                         uf for uf in (update_fields or []) if uf in model_fields_by_name
                     ]
-                    
+
                     if filtered_updates:
                         level_update_conflicts = True
                         level_unique_fields = normalized_unique
                         level_update_fields = filtered_updates
-            
+
             # Create parent level
             parent_level = ParentLevel(
                 model_class=model_class,
@@ -252,7 +274,7 @@ class MTIHandler:
                 update_fields=level_update_fields,
             )
             parent_levels.append(parent_level)
-        
+
         return parent_levels
 
     def _has_matching_constraint(self, model_class, normalized_unique):
@@ -260,21 +282,29 @@ class MTIHandler:
         try:
             from django.db.models import UniqueConstraint
             constraint_field_sets = [
-                tuple(c.fields) for c in model_class._meta.constraints 
+                tuple(c.fields) for c in model_class._meta.constraints
                 if isinstance(c, UniqueConstraint)
             ]
         except Exception:
             constraint_field_sets = []
-        
+
         # Check unique_together
         ut = getattr(model_class._meta, "unique_together", ()) or ()
         if isinstance(ut, tuple) and ut and not isinstance(ut[0], (list, tuple)):
             ut = (ut,)
         ut_field_sets = [tuple(group) for group in ut]
-        
+
+        # Check individual field uniqueness
+        unique_field_sets = []
+        for field in model_class._meta.local_fields:
+            if field.unique and not field.primary_key:
+                unique_field_sets.append((field.name,))
+
         # Compare as sets
         provided_set = set(normalized_unique)
-        for group in constraint_field_sets + ut_field_sets:
+        all_constraint_sets = constraint_field_sets + ut_field_sets + unique_field_sets
+
+        for group in all_constraint_sets:
             if provided_set == set(group):
                 return True
         return False
@@ -292,13 +322,13 @@ class MTIHandler:
             Parent model instance (not saved)
         """
         parent_obj = parent_model()
-        
+
         # Copy field values from source
         for field in parent_model._meta.local_fields:
             if hasattr(source_obj, field.name):
                 value = getattr(source_obj, field.name, None)
                 if value is not None:
-                    if (field.is_relation and not field.many_to_many and 
+                    if (field.is_relation and not field.many_to_many and
                         not field.one_to_many):
                         # Handle FK fields
                         if hasattr(value, "pk") and value.pk is not None:
@@ -307,7 +337,7 @@ class MTIHandler:
                             setattr(parent_obj, field.attname, value)
                     else:
                         setattr(parent_obj, field.name, value)
-        
+
         # Link to parent if exists
         if current_parent is not None:
             for field in parent_model._meta.local_fields:
@@ -315,22 +345,22 @@ class MTIHandler:
                     field.remote_field.model == current_parent.__class__):
                     setattr(parent_obj, field.name, current_parent)
                     break
-        
+
         # Copy object state
-        if hasattr(source_obj, '_state') and hasattr(parent_obj, '_state'):
+        if hasattr(source_obj, "_state") and hasattr(parent_obj, "_state"):
             parent_obj._state.adding = source_obj._state.adding
-            if hasattr(source_obj._state, 'db'):
+            if hasattr(source_obj._state, "db"):
                 parent_obj._state.db = source_obj._state.db
-        
+
         # Handle auto_now_add and auto_now fields
         for field in parent_model._meta.local_fields:
-            if getattr(field, 'auto_now_add', False):
+            if getattr(field, "auto_now_add", False):
                 if getattr(parent_obj, field.name) is None:
                     field.pre_save(parent_obj, add=True)
                     setattr(parent_obj, field.attname, field.value_from_object(parent_obj))
-            elif getattr(field, 'auto_now', False):
+            elif getattr(field, "auto_now", False):
                 field.pre_save(parent_obj, add=True)
-        
+
         return parent_obj
 
     def _create_child_instance_template(self, source_obj, child_model):
@@ -347,22 +377,22 @@ class MTIHandler:
             Child model instance (not saved, no parent links)
         """
         child_obj = child_model()
-        
+
         # Copy field values (excluding AutoField and parent links)
         for field in child_model._meta.local_fields:
             if isinstance(field, AutoField):
                 continue
-            
+
             # Skip parent link fields - executor will set these
-            if field.is_relation and hasattr(field, 'related_model'):
+            if field.is_relation and hasattr(field, "related_model"):
                 # Check if this field is a parent link
                 if child_model._meta.get_ancestor_link(field.related_model) == field:
                     continue
-            
+
             if hasattr(source_obj, field.name):
                 value = getattr(source_obj, field.name, None)
                 if value is not None:
-                    if (field.is_relation and not field.many_to_many and 
+                    if (field.is_relation and not field.many_to_many and
                         not field.one_to_many):
                         if hasattr(value, "pk") and value.pk is not None:
                             setattr(child_obj, field.attname, value.pk)
@@ -370,22 +400,22 @@ class MTIHandler:
                             setattr(child_obj, field.attname, value)
                     else:
                         setattr(child_obj, field.name, value)
-        
+
         # Copy object state
-        if hasattr(source_obj, '_state') and hasattr(child_obj, '_state'):
+        if hasattr(source_obj, "_state") and hasattr(child_obj, "_state"):
             child_obj._state.adding = source_obj._state.adding
-            if hasattr(source_obj._state, 'db'):
+            if hasattr(source_obj._state, "db"):
                 child_obj._state.db = source_obj._state.db
-        
+
         # Handle auto_now_add and auto_now fields
         for field in child_model._meta.local_fields:
-            if getattr(field, 'auto_now_add', False):
+            if getattr(field, "auto_now_add", False):
                 if getattr(child_obj, field.name) is None:
                     field.pre_save(child_obj, add=True)
                     setattr(child_obj, field.attname, field.value_from_object(child_obj))
-            elif getattr(field, 'auto_now', False):
+            elif getattr(field, "auto_now", False):
                 field.pre_save(child_obj, add=True)
-        
+
         return child_obj
 
     # ==================== MTI BULK UPDATE PLANNING ====================
@@ -404,48 +434,49 @@ class MTIHandler:
         Returns:
             MTIUpdatePlan object
         """
-        from django_bulk_hooks.operations.mti_plans import MTIUpdatePlan, ModelFieldGroup
-        
+        from django_bulk_hooks.operations.mti_plans import ModelFieldGroup
+        from django_bulk_hooks.operations.mti_plans import MTIUpdatePlan
+
         if not objs:
             return None
-        
+
         inheritance_chain = self.get_inheritance_chain()
         if len(inheritance_chain) <= 1:
             raise ValueError("build_update_plan called on non-MTI model")
-        
+
         batch_size = batch_size or len(objs)
-        
+
         # Handle auto_now fields
         for obj in objs:
             for model in inheritance_chain:
                 for field in model._meta.local_fields:
-                    if getattr(field, 'auto_now', False):
+                    if getattr(field, "auto_now", False):
                         field.pre_save(obj, add=False)
-        
+
         # Add auto_now fields to update list
         auto_now_fields = set()
         for model in inheritance_chain:
             for field in model._meta.local_fields:
-                if getattr(field, 'auto_now', False):
+                if getattr(field, "auto_now", False):
                     auto_now_fields.add(field.name)
-        
+
         all_fields = list(fields) + list(auto_now_fields)
-        
+
         # Group fields by model
         field_groups = []
         for model_idx, model in enumerate(inheritance_chain):
             model_fields = []
-            
+
             for field_name in all_fields:
                 try:
                     field = self.model_cls._meta.get_field(field_name)
                     if field in model._meta.local_fields:
                         # Skip auto_now_add fields for updates
-                        if not getattr(field, 'auto_now_add', False):
+                        if not getattr(field, "auto_now_add", False):
                             model_fields.append(field_name)
                 except Exception:
                     continue
-            
+
             if model_fields:
                 # Determine filter field
                 if model_idx == 0:
@@ -458,13 +489,13 @@ class MTIHandler:
                             parent_link = model._meta.parents[parent_model]
                             break
                     filter_field = parent_link.attname if parent_link else "pk"
-                
+
                 field_groups.append(ModelFieldGroup(
                     model_class=model,
                     fields=model_fields,
                     filter_field=filter_field,
                 ))
-        
+
         return MTIUpdatePlan(
             inheritance_chain=inheritance_chain,
             field_groups=field_groups,

django_bulk_hooks/operations/mti_plans.py

@@ -5,8 +5,9 @@ These are pure data structures returned by MTIHandler to be executed by BulkExec
 This separates planning (logic) from execution (database operations).
 """
 
-from dataclasses import dataclass, field
-from typing import Dict, List, Any
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
 
 
 @dataclass
@@ -23,11 +24,11 @@ class ParentLevel:
         update_fields: Fields to update on conflict (if update_conflicts=True)
     """
     model_class: Any
-    objects: List[Any]
-    original_object_map: Dict[int, int] = field(default_factory=dict)
+    objects: list[Any]
+    original_object_map: dict[int, int] = field(default_factory=dict)
     update_conflicts: bool = False
-    unique_fields: List[str] = field(default_factory=list)
-    update_fields: List[str] = field(default_factory=list)
+    unique_fields: list[str] = field(default_factory=list)
+    update_fields: list[str] = field(default_factory=list)
 
 
 @dataclass
@@ -45,13 +46,21 @@ class MTICreatePlan:
         child_model: The child model class
         original_objects: Original objects provided by user
         batch_size: Batch size for operations
+        existing_record_ids: Set of id() of original objects that represent existing DB records
+        update_conflicts: Whether this is an upsert operation
+        unique_fields: Fields used for conflict detection
+        update_fields: Fields to update on conflict
     """
-    inheritance_chain: List[Any]
-    parent_levels: List[ParentLevel]
-    child_objects: List[Any]
+    inheritance_chain: list[Any]
+    parent_levels: list[ParentLevel]
+    child_objects: list[Any]
     child_model: Any
-    original_objects: List[Any]
+    original_objects: list[Any]
     batch_size: int = None
+    existing_record_ids: set = field(default_factory=set)
+    update_conflicts: bool = False
+    unique_fields: list[str] = field(default_factory=list)
+    update_fields: list[str] = field(default_factory=list)
 
 
 @dataclass
@@ -65,7 +74,7 @@ class ModelFieldGroup:
         filter_field: Field to use for filtering (e.g., 'pk' or parent link attname)
     """
     model_class: Any
-    fields: List[str]
+    fields: list[str]
     filter_field: str = "pk"
 
 
@@ -80,8 +89,8 @@ class MTIUpdatePlan:
         objects: Objects to update
         batch_size: Batch size for operations
     """
-    inheritance_chain: List[Any]
-    field_groups: List[ModelFieldGroup]
-    objects: List[Any]
+    inheritance_chain: list[Any]
+    field_groups: list[ModelFieldGroup]
+    objects: list[Any]
     batch_size: int = None
 

django_bulk_hooks/operations/record_classifier.py

@@ -0,0 +1,184 @@
+"""
+Record classification service for database queries.
+
+This service handles all database queries related to classifying and fetching
+records based on various criteria (PKs, unique fields, etc.).
+
+Separates data access concerns from business logic.
+"""
+
+import logging
+
+from django.db.models import Q
+
+logger = logging.getLogger(__name__)
+
+
+class RecordClassifier:
+    """
+    Service for classifying and fetching records via database queries.
+    
+    This is the SINGLE point of truth for record classification queries.
+    Keeps database access logic separate from business/planning logic.
+    """
+
+    def __init__(self, model_cls):
+        """
+        Initialize classifier for a specific model.
+
+        Args:
+            model_cls: The Django model class
+        """
+        self.model_cls = model_cls
+
+    def classify_for_upsert(self, objs, unique_fields):
+        """
+        Classify records as new or existing based on unique_fields.
+        
+        Queries the database to check which records already exist based on the
+        unique_fields constraint.
+        
+        Args:
+            objs: List of model instances
+            unique_fields: List of field names that form the unique constraint
+            
+        Returns:
+            Tuple of (existing_record_ids, existing_pks_map)
+            - existing_record_ids: Set of id() for objects that exist in DB
+            - existing_pks_map: Dict mapping id(obj) -> pk for existing records
+        """
+        if not unique_fields or not objs:
+            return set(), {}
+
+        # Build a query to find existing records
+        queries = []
+        obj_to_unique_values = {}
+
+        for obj in objs:
+            # Build lookup dict for this object's unique fields
+            lookup = {}
+            for field_name in unique_fields:
+                value = getattr(obj, field_name, None)
+                if value is None:
+                    # Can't match on None values
+                    break
+                lookup[field_name] = value
+            else:
+                # All unique fields have values, add to query
+                if lookup:
+                    queries.append(Q(**lookup))
+                    obj_to_unique_values[id(obj)] = tuple(lookup.values())
+
+        if not queries:
+            return set(), {}
+
+        # Query for existing records
+        combined_query = queries[0]
+        for q in queries[1:]:
+            combined_query |= q
+
+        existing_records = list(
+            self.model_cls.objects.filter(combined_query).values("pk", *unique_fields),
+        )
+
+        # Map existing records back to original objects
+        existing_record_ids = set()
+        existing_pks_map = {}
+
+        for record in existing_records:
+            record_values = tuple(record[field] for field in unique_fields)
+            # Find which object(s) match these values
+            for obj_id, obj_values in obj_to_unique_values.items():
+                if obj_values == record_values:
+                    existing_record_ids.add(obj_id)
+                    existing_pks_map[obj_id] = record["pk"]
+
+        logger.info(
+            f"Classified {len(existing_record_ids)} existing and "
+            f"{len(objs) - len(existing_record_ids)} new records for upsert",
+        )
+
+        return existing_record_ids, existing_pks_map
+
+    def fetch_by_pks(self, pks, select_related=None, prefetch_related=None):
+        """
+        Fetch records by primary keys with optional relationship loading.
+        
+        Args:
+            pks: List of primary key values
+            select_related: Optional list of fields to select_related
+            prefetch_related: Optional list of fields to prefetch_related
+            
+        Returns:
+            Dict[pk, instance] for O(1) lookups
+        """
+        if not pks:
+            return {}
+
+        queryset = self.model_cls._base_manager.filter(pk__in=pks)
+
+        if select_related:
+            queryset = queryset.select_related(*select_related)
+
+        if prefetch_related:
+            queryset = queryset.prefetch_related(*prefetch_related)
+
+        return {obj.pk: obj for obj in queryset}
+
+    def fetch_by_unique_constraint(self, field_values_map):
+        """
+        Fetch records matching a unique constraint.
+        
+        Args:
+            field_values_map: Dict of {field_name: value} for unique constraint
+            
+        Returns:
+            Model instance if found, None otherwise
+        """
+        try:
+            return self.model_cls.objects.get(**field_values_map)
+        except self.model_cls.DoesNotExist:
+            return None
+        except self.model_cls.MultipleObjectsReturned:
+            logger.warning(
+                f"Multiple {self.model_cls.__name__} records found for "
+                f"unique constraint {field_values_map}",
+            )
+            return self.model_cls.objects.filter(**field_values_map).first()
+
+    def exists_by_pks(self, pks):
+        """
+        Check if records exist by primary keys without fetching them.
+        
+        Args:
+            pks: List of primary key values
+            
+        Returns:
+            Set of PKs that exist in the database
+        """
+        if not pks:
+            return set()
+
+        existing_pks = self.model_cls.objects.filter(
+            pk__in=pks,
+        ).values_list("pk", flat=True)
+
+        return set(existing_pks)
+
+    def count_by_unique_fields(self, objs, unique_fields):
+        """
+        Count how many objects already exist based on unique fields.
+        
+        Useful for validation or reporting before upsert operations.
+        
+        Args:
+            objs: List of model instances
+            unique_fields: List of field names that form the unique constraint
+            
+        Returns:
+            Tuple of (existing_count, new_count)
+        """
+        existing_ids, _ = self.classify_for_upsert(objs, unique_fields)
+        existing_count = len(existing_ids)
+        new_count = len(objs) - existing_count
+        return existing_count, new_count

django_bulk_hooks/queryset.py

@@ -7,7 +7,9 @@ complex coordination required for bulk operations with hooks.
 """
 
 import logging
-from django.db import models, transaction
+
+from django.db import models
+from django.db import transaction
 
 logger = logging.getLogger(__name__)
 
@@ -98,7 +100,7 @@ class HookQuerySet(models.QuerySet):
             fields = self.coordinator.analyzer.detect_changed_fields(objs)
             if not fields:
                 logger.debug(
-                    f"bulk_update: No fields changed for {len(objs)} {self.model.__name__} objects"
+                    f"bulk_update: No fields changed for {len(objs)} {self.model.__name__} objects",
                 )
                 return 0
 
@@ -133,7 +135,7 @@ class HookQuerySet(models.QuerySet):
 
     @transaction.atomic
     def bulk_delete(
-        self, objs, bypass_hooks=False, bypass_validation=False, **kwargs
+        self, objs, bypass_hooks=False, bypass_validation=False, **kwargs,
     ):
         """
         Delete multiple objects with hook support.