django-bulk-hooks 0.2.9__py3-none-any.whl → 0.2.93__py3-none-any.whl

django_bulk_hooks/operations/mti_plans.py

@@ -1,87 +1,103 @@
-"""
-MTI operation plans - Data structures for multi-table inheritance operations.
-
-These are pure data structures returned by MTIHandler to be executed by BulkExecutor.
-This separates planning (logic) from execution (database operations).
-"""
-
-from dataclasses import dataclass, field
-from typing import Dict, List, Any
-
-
-@dataclass
-class ParentLevel:
-    """
-    Represents one level in the parent hierarchy for MTI bulk create.
-    
-    Attributes:
-        model_class: The parent model class for this level
-        objects: List of parent instances to create
-        original_object_map: Maps parent instance id() -> original object id()
-        update_conflicts: Whether to enable UPSERT for this level
-        unique_fields: Fields for conflict detection (if update_conflicts=True)
-        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)
-    update_conflicts: bool = False
-    unique_fields: List[str] = field(default_factory=list)
-    update_fields: List[str] = field(default_factory=list)
-
-
-@dataclass
-class MTICreatePlan:
-    """
-    Plan for executing bulk_create on an MTI model.
-    
-    This plan describes WHAT to create, not HOW to create it.
-    The executor is responsible for executing this plan.
-    
-    Attributes:
-        inheritance_chain: List of model classes from root to child
-        parent_levels: List of ParentLevel objects, one per parent model
-        child_objects: List of child instances to create (not yet with parent links)
-        child_model: The child model class
-        original_objects: Original objects provided by user
-        batch_size: Batch size for operations
-    """
-    inheritance_chain: List[Any]
-    parent_levels: List[ParentLevel]
-    child_objects: List[Any]
-    child_model: Any
-    original_objects: List[Any]
-    batch_size: int = None
-
-
-@dataclass
-class ModelFieldGroup:
-    """
-    Represents fields to update for one model in the inheritance chain.
-    
-    Attributes:
-        model_class: The model class
-        fields: List of field names to update on this model
-        filter_field: Field to use for filtering (e.g., 'pk' or parent link attname)
-    """
-    model_class: Any
-    fields: List[str]
-    filter_field: str = "pk"
-
-
-@dataclass
-class MTIUpdatePlan:
-    """
-    Plan for executing bulk_update on an MTI model.
-    
-    Attributes:
-        inheritance_chain: List of model classes from root to child
-        field_groups: List of ModelFieldGroup objects
-        objects: Objects to update
-        batch_size: Batch size for operations
-    """
-    inheritance_chain: List[Any]
-    field_groups: List[ModelFieldGroup]
-    objects: List[Any]
-    batch_size: int = None
-
+"""
+MTI operation plans - Data structures for multi-table inheritance operations.
+
+These are pure data structures returned by MTIHandler to be executed by BulkExecutor.
+This separates planning (logic) from execution (database operations).
+"""
+
+from dataclasses import dataclass
+from dataclasses import field
+from typing import Any
+
+
+@dataclass
+class ParentLevel:
+    """
+    Represents one level in the parent hierarchy for MTI bulk create.
+
+    Attributes:
+        model_class: The parent model class for this level
+        objects: List of parent instances to create
+        original_object_map: Maps parent instance id() -> original object id()
+        update_conflicts: Whether to enable UPSERT for this level
+        unique_fields: Fields for conflict detection (if update_conflicts=True)
+        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)
+    update_conflicts: bool = False
+    unique_fields: list[str] = field(default_factory=list)
+    update_fields: list[str] = field(default_factory=list)
+
+
+@dataclass
+class MTICreatePlan:
+    """
+    Plan for executing bulk_create on an MTI model.
+
+    This plan describes WHAT to create, not HOW to create it.
+    The executor is responsible for executing this plan.
+
+    Attributes:
+        inheritance_chain: List of model classes from root to child
+        parent_levels: List of ParentLevel objects, one per parent model
+        child_objects: List of child instances to create (not yet with parent links)
+        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 (original, unfiltered)
+        update_fields: Fields to update on conflict (original, unfiltered)
+        child_unique_fields: Pre-filtered field objects for child table conflict detection
+        child_update_fields: Pre-filtered field objects for child table updates
+    """
+
+    inheritance_chain: list[Any]
+    parent_levels: list[ParentLevel]
+    child_objects: list[Any]
+    child_model: 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)
+    child_unique_fields: list = field(default_factory=list)  # Field objects for child table
+    child_update_fields: list = field(default_factory=list)  # Field objects for child table
+
+
+@dataclass
+class ModelFieldGroup:
+    """
+    Represents fields to update for one model in the inheritance chain.
+
+    Attributes:
+        model_class: The model class
+        fields: List of field names to update on this model
+        filter_field: Field to use for filtering (e.g., 'pk' or parent link attname)
+    """
+
+    model_class: Any
+    fields: list[str]
+    filter_field: str = "pk"
+
+
+@dataclass
+class MTIUpdatePlan:
+    """
+    Plan for executing bulk_update on an MTI model.
+
+    Attributes:
+        inheritance_chain: List of model classes from root to child
+        field_groups: List of ModelFieldGroup objects
+        objects: Objects to update
+        batch_size: Batch size for operations
+    """
+
+    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,196 @@
+"""
+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
+
+from django_bulk_hooks.operations.field_utils import get_field_value_for_db
+
+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, query_model=None):
+        """
+        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
+            query_model: Optional model class to query (for MTI, may be different from self.model_cls)
+
+        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(), {}
+
+        # Use query_model if provided (for MTI scenarios), otherwise use self.model_cls
+        query_model = query_model or self.model_cls
+
+        # 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 = {}
+            normalized_values = []
+
+            for field_name in unique_fields:
+                # Use centralized field value extraction for consistent FK handling
+                value = get_field_value_for_db(obj, field_name, query_model)
+                if value is None:
+                    # Can't match on None values
+                    break
+                lookup[field_name] = value
+                normalized_values.append(value)
+            else:
+                # All unique fields have values, add to query
+                if lookup:
+                    queries.append(Q(**lookup))
+                    # Store normalized values for comparison with database results
+                    obj_to_unique_values[id(obj)] = tuple(normalized_values)
+
+        if not queries:
+            return set(), {}
+
+        # Query for existing records
+        combined_query = queries[0]
+        for q in queries[1:]:
+            combined_query |= q
+
+        logger.info(f"Classifying for upsert: model={query_model.__name__}, query={combined_query}, unique_fields={unique_fields}")
+        queryset = query_model.objects.filter(combined_query)
+        logger.info(f"Queryset SQL: {queryset.query}")
+        logger.info(f"All records in table: {query_model.objects.all().count()}")
+        existing_records = list(queryset.values("pk", *unique_fields))
+        logger.info(f"Found {len(existing_records)} existing records: {existing_records}")
+
+        # 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 {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 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