django-fast-treenode 2.1.4__py3-none-any.whl → 3.0.0__py3-none-any.whl

treenode/models/mixins/tree.py

@@ -2,17 +2,18 @@
 """
 TreeNode Tree Mixin
 
-Version: 2.1.0
+Version: 3.0.0
 Author: Timur Kady
 Email: timurkady@yandex.com
 """
 
 import json
-from django.db import models, transaction
+import inspect
+from django.db import models, transaction, connection
 from collections import OrderedDict
 from django.core.serializers.json import DjangoJSONEncoder
 
-from ...cache import cached_method
+from ...cache import treenode_cache as cache
 
 
 class TreeNodeTreeMixin(models.Model):
@@ -23,17 +24,41 @@ class TreeNodeTreeMixin(models.Model):
 
         abstract = True
 
-    @classmethod
-    def dump_tree(cls, instance=None):
+    def clone_subtree(self, parent=None):
         """
-        Return an n-dimensional dictionary representing the model tree.
+        Clone self and entire subtree under given parent.
 
-        Introduced for compatibility with other packages.
+        Returns new root of the cloned subtree.
         """
-        return cls.get_tree(cls, instance)
+        model = self._meta.model
+
+        def _clone_node(node, parent):
+            # Copy all regular fields (including ForeignKey via attname)
+            data = {
+                f.attname: getattr(node, f.attname)
+                for f in node._meta.concrete_fields
+                if not f.primary_key and f.name != "_path"
+            }
+            data['parent'] = parent
+
+            # Create a new node
+            new_node = model.objects.create(**data)
+
+            # Copy ManyToMany fields
+            for m2m_field in node._meta.many_to_many:
+                related_ids = getattr(
+                    node, m2m_field.name).values_list('pk', flat=True)
+                getattr(new_node, m2m_field.name).set(related_ids)
+
+            # Clone descendants recursively
+            for child in node.get_children():
+                _clone_node(child, new_node)
+
+            return new_node
+
+        return _clone_node(self, parent)
 
     @classmethod
-    @cached_method
     def get_tree(cls, instance=None):
         """
         Return an n-dimensional dictionary representing the model tree.
@@ -44,39 +69,39 @@ class TreeNodeTreeMixin(models.Model):
         """
         # If instance is passed, we get all its descendants (including itself)
         if instance:
-            queryset = instance.get_descendants_queryset(include_self=True)\
-                .annotate(depth=models.Max("parents_set__depth"))
+            queryset = instance.get_descendants_queryset(include_self=True)
         else:
             # Load all records at once
-            queryset = cls.objects.all()
+            queryset = cls.objects.get_queryset()
 
         # Dictionary for quick access to nodes by id and list for iteration
         nodes_by_id = {}
         nodes_list = []
 
         # Loop through all nodes using an iterator
-        for node in queryset.iterator(chunk_size=1000):
+        for node in queryset.order_by("_path").iterator(chunk_size=1000):
             # Create a dictionary for the node.
-            # In Python 3.7+, the standard dict preserves insertion order.
-            # We'll stick to the order:
-            # id, 'tn_parent', 'tn_priority', 'level', then the rest of
-            # the fields. Сlose the dictionary with the fields 'level', 'path',
-            # 'children'
             node_dict = OrderedDict()
             node_dict['id'] = node.id
-            node_dict['tn_parent'] = node.tn_parent_id
-            node_dict['tn_priority'] = node.tn_priority
-            node_dict['level'] = node.get_depth()
-            node_dict['path'] = node.get_breadcrumbs('tn_priority')
+            node_dict['parent'] = node.parent_id
+            node_dict['priority'] = node.priority
+            node_dict['depth'] = node.get_depth()
+            node_dict['path'] = node.get_order()
+            node_dict['children'] = []  # node.get_children_pks()
 
             # Add the rest of the model fields.
             # Iterate over all the fields obtained via _meta.get_fields()
-            for field in node._meta.get_fields():
+            fields = [
+                f for f in node._meta.get_fields()
+                if f.concrete and not f.auto_created and
+                not f.name.startswith('_')
+            ]
+
+            for field in fields:
                 # Skipping fields that are already added or not required
                 # (e.g. tn_closure or virtual links)
                 if field.name in [
-                        'id', 'tn_parent', 'tn_priority', 'tn_closure',
-                        'children']:
+                        'id', 'parent', 'priority', '_depth', 'children']:
                     continue
 
                 try:
@@ -91,14 +116,6 @@ class TreeNodeTreeMixin(models.Model):
 
                 node_dict[field.name] = value
 
-            # Adding a materialized path
-            node_dict['path'] = None
-            # Adding a nesting level
-            node_dict['level'] = None
-            # We initialize the list of children, which we will then fill
-            # when assembling the tree
-            node_dict['children'] = []
-
             # Save the node both in the list and in the dictionary by id
             # for quick access
             nodes_by_id[node.id] = node_dict
@@ -107,7 +124,7 @@ class TreeNodeTreeMixin(models.Model):
         # Build a tree: assign each node a list of its children
         tree = []
         for node_dict in nodes_list:
-            parent_id = node_dict['tn_parent']
+            parent_id = node_dict['parent']
             # If there is a parent and it is present in nodes_by_id, then
             # add the current node to the list of its children
             if parent_id and parent_id in nodes_by_id:
@@ -122,7 +139,7 @@ class TreeNodeTreeMixin(models.Model):
     @classmethod
     def get_tree_json(cls, instance=None):
         """Represent the tree as a JSON-compatible string."""
-        tree = cls.dump_tree(instance)
+        tree = cls.get_tree(instance)
         return DjangoJSONEncoder().encode(tree)
 
     @classmethod
@@ -130,86 +147,132 @@ class TreeNodeTreeMixin(models.Model):
         """
         Load a tree from a list of dictionaries.
 
-        Loaded nodes are synchronized with the database: existing records
-        are updated, new ones are created.
-        Each dictionary must contain the 'id' key to identify the record.
+        - Ignores visual/path metadata ('path', 'depth', '_path', '_depth')
+        - Accepts 'id', 'name', 'priority', 'parent'
+        - Automatically recalculates _path, _depth, priority after loading
         """
-
-        def flatten_tree(nodes, model_fields):
-            """
-            Recursively traverse the tree and generate lists of nodes.
-
-            Each node in the list is a copy of the dictionary without
-            the service keys 'children', 'level', 'path'.
-            """
+        # Step 1: Flatten nested tree into list of nodes with temporary
+        # parent references
+        def flatten_tree(tree, parent_temp_id=None):
             flat = []
-            for node in nodes:
-                # Create a copy of the dictionary so as not to affect
-                # the original tree
-                node_copy = node.copy()
-                children = node_copy.pop('children', [])
-                # Remove temporary/service fields that are not related to
-                # the model
-                for key in list(node_copy.keys() - model_fields):
-                    del node_copy[key]
-                flat.append(node_copy)
-                # Recursively add all children
-                flat.extend(flatten_tree(children, model_fields))
+            for node in tree:
+                node = node.copy()
+                children = node.pop('children', [])
+                node_id = node.pop('id', None)
+                node.pop('path', None)
+                node.pop('depth', None)
+                node.pop('parent', None)
+                flat_node = {
+                    **node,
+                    '_id': int(node_id) if node_id is not None else None,
+                    '_parent': int(parent_temp_id) if parent_temp_id is not None else None  # noqa: D501
+                }
+                flat.append(flat_node)
+                flat.extend(flatten_tree(children, parent_temp_id=node_id))
             return flat
 
-        # Get a flat list of nodes (from root to leaf)
-        model_fields = [field.name for field in cls._meta.get_fields()]
-        flat_nodes = flatten_tree(tree_data, model_fields)
-
-        # Load all ids for a given model from the database to minimize
-        # the number of database requests
-        existing_ids = set(cls.objects.values_list('id', flat=True))
-
-        # Lists for nodes to update and create
-        nodes_to_update = []
-        nodes_to_create = []
-
-        # Determine which model fields should be updated (excluding 'id')
-        # This assumes that all model fields are used in serialization
-        # (excluding service ones, like children)
-        field_names = model_fields.remove('id')
-
-        # Iterate over each node from the flat list
-        for node_data in flat_nodes:
-            # Collect data for creation/update.
-            # There is already an 'id' in node_data, so we can distinguish
-            # an existing entry from a new one.
-            data = {k: v for k, v in node_data.items()
-                    if k in field_names or k == 'id'}
-
-            # Handle the foreign key tn_parent.
-            # If the value is None, then there is no parent, otherwise
-            # the parent id is expected.
-            # Additional checks can be added here if needed.
-            if 'tn_parent' in data and data['tn_parent'] is None:
-                data['tn_parent'] = None
-
-            # Если id уже есть в базе, то будем обновлять запись,
-            # иначе создаем новую.
-            if data['id'] in existing_ids:
-                nodes_to_update.append(cls(**data))
-            else:
-                nodes_to_create.append(cls(**data))
-
-        # Perform operations in a transaction to ensure data integrity.
-        with transaction.atomic():
-            # bulk_create – creating new nodes
-            if nodes_to_create:
-                cls.objects.bulk_create(nodes_to_create, batch_size=1000)
-            # bulk_update – updating existing nodes.
-            # When bulk_update, we must specify a list of fields to update.
-            if nodes_to_update:
+        raw_data = flatten_tree(tree_data)
+
+        # Step 2: Identify which objects should be updated
+        pks = [r["_id"] for r in raw_data if r["_id"] is not None]
+        existing_ids = set(cls.objects.filter(
+            pk__in=pks).values_list('id', flat=True))
+        to_update_ids = {r["_id"] for r in raw_data if r["_id"] in existing_ids}
+
+        existing_objs = {
+            obj.pk: obj for obj in cls.objects.filter(pk__in=to_update_ids)}
+
+        # Step 3: Determine tree levels for correct creation order
+        levels = {}
+        for record in raw_data:
+            level = 0
+            p = record.get('_parent')
+            while p:
+                level += 1
+                p = next((r['_parent']
+                         for r in raw_data if r['_id'] == p), None)
+            levels.setdefault(level, []).append(record)
+
+        records_by_level = [
+            sorted(levels[level], key=lambda x: x.get('_parent') or -1)
+            for level in sorted(levels)
+        ]
+
+        # Step 4: Map old temporary IDs to actual DB IDs
+        created_objects = []
+        updated_objects = []
+        # ID mapping for existing objects
+        id_map = {pk: pk for pk in to_update_ids}
+
+        model_fields = {
+            f.name for f in cls._meta.get_fields()
+            if f.concrete and not f.auto_created and not f.name.startswith('_')
+        }
+
+        # Step 5: Process records level by level
+        for level_records in records_by_level:
+            objs_to_create = []
+            objs_to_update = []
+            fields_to_update_set = set()
+
+            for record in level_records:
+                temp_id = record['_id']
+                temp_parent = record.get('_parent')
+                data = {k: v for k, v in record.items() if not k.startswith(
+                    '_') and k in model_fields}
+
+                if temp_parent and temp_parent in id_map:
+                    data['parent_id'] = id_map[temp_parent]
+
+                if temp_id in to_update_ids:
+                    obj = existing_objs[temp_id]
+                    for field, value in data.items():
+                        setattr(obj, field, value)
+                    setattr(obj, '_temp_parent', temp_parent)
+                    objs_to_update.append(obj)
+                    fields_to_update_set.update(data.keys())
+                else:
+                    obj = cls(**data)
+                    obj.full_clean()
+                    setattr(obj, '_temp_id', temp_id)
+                    setattr(obj, '_temp_parent', temp_parent)
+                    objs_to_create.append(obj)
+
+            # Step 6: Bulk create and track real IDs
+            if objs_to_create:
+                created = cls.objects.bulk_create(objs_to_create)
+                for obj in created:
+                    temp_id = getattr(obj, '_temp_id')
+                    id_map[temp_id] = obj.id
+                    delattr(obj, '_temp_id')
+                    created_objects.append(obj)
+
+            # Step 7: Bulk update existing objects
+            if objs_to_update:
+                for obj in objs_to_update:
+                    obj.full_clean()
                 cls.objects.bulk_update(
-                    nodes_to_update,
-                    fields=field_names,
-                    batch_size=1000
-                )
-        cls.clear_cache()
+                    objs_to_update, list(fields_to_update_set))
+                updated_objects.extend(objs_to_update)
+
+        all_objects = created_objects + updated_objects
+
+        # Step 8: Finalize parent assignments using ID map
+        for obj in all_objects:
+            temp_parent = getattr(obj, '_temp_parent', None)
+            if temp_parent is not None and temp_parent in id_map:
+                obj.parent_id = id_map[temp_parent]
+            for attr in ['_temp_id', '_temp_parent']:
+                if hasattr(obj, attr):
+                    delattr(obj, attr)
+
+        cls.objects.bulk_update(all_objects, ['parent_id'])
+
+        # Step 9: Rebuild tree structure (path, depth, priority)
+        cls.tasks.run()
+
+        # Step 10: Invalidate tree-related caches
+        cache.invalidate(cls._meta.label)
 
     @classmethod
     def load_tree_json(cls, json_str):
@@ -222,36 +285,9 @@ class TreeNodeTreeMixin(models.Model):
             tree_data = json.loads(json_str)
         except json.JSONDecodeError as e:
             raise ValueError(f"Error decoding JSON: {e}")
-
-        cls.load_tree(tree_data)
-
-    @classmethod
-    @cached_method
-    def get_tree_display(cls, instance=None, symbol="—"):
-        """Get a multiline string representing the model tree."""
-        # If instance is passed, we get all its descendants (including itself)
-        if instance:
-            queryset = instance.get_descendants_queryset(include_self=True)\
-                .prefetch_related("tn_children")\
-                .annotate(depth=models.Max("parents_set__depth"))\
-                .order_by("depth", "tn_parent", "tn_priority")
-        else:
-            queryset = cls.objects.all()\
-                .prefetch_related("tn_children")\
-                .annotate(depth=models.Max("parents_set__depth"))\
-                .order_by("depth", "tn_parent", "tn_priority")
-        # Convert queryset to list for indexed access
-        nodes = list(queryset)
-        sorted_nodes = cls._sort_node_list(nodes)
-        result = []
-        for node in sorted_nodes:
-            # Insert an indent proportional to the depth of the node
-            indent = symbol * node.depth
-            result.append(indent + str(node))
-        return result
+        return tree_data
 
     @classmethod
-    @cached_method
     def get_tree_annotated(cls):
         """
         Get an annotated list from a tree branch.
@@ -285,10 +321,7 @@ class TreeNodeTreeMixin(models.Model):
 
         """
         # Load the tree with the required preloads and depth annotation.
-        queryset = cls.objects.all()\
-            .prefetch_related("tn_children")\
-            .annotate(depth=models.Max("parents_set__depth"))\
-            .order_by("depth", "tn_parent", "tn_priority")
+        queryset = cls.objects.all()
         # Convert queryset to list for indexed access
         nodes = list(queryset)
         total_nodes = len(nodes)
@@ -301,7 +334,7 @@ class TreeNodeTreeMixin(models.Model):
             value = str(node)
             # Determine if there are descendants (use prefetch_related to avoid
             # additional queries)
-            value_open = len(node.tn_children.all()) > 0
+            value_open = len(node.children.all()) > 0
             level = node.depth
 
             # Calculate the "close" field
@@ -326,19 +359,45 @@ class TreeNodeTreeMixin(models.Model):
         return result
 
     @classmethod
-    @transaction.atomic
     def update_tree(cls):
-        """Rebuilds the closure table."""
-        # Clear cache
-        cls.clear_cache()
-        cls.closure_model.delete_all()
-        objs = list(cls.objects.all())
-        cls.closure_model.objects.bulk_create(objs, batch_size=1000)
+        """Rebuld the tree."""
+        cls.tasks.add("update", None)
+        cls.tasks.run()
+
+    def delete_tree(self, include_self=True):
+        """
+        Delete current node and all descendants.
+
+        If include_self=False, only descendants will be deleted.
+        """
+        table = self._meta.db_table
+        path = self._path
+        like_pattern = f"{path}.%"
+
+        if include_self:
+            sql = f"""
+                DELETE FROM {table}
+                WHERE _path = %s OR _path LIKE %s
+            """
+            params = [path, like_pattern]
+        else:
+            sql = f"""
+                DELETE FROM {table}
+                WHERE _path LIKE %s
+            """
+            params = [like_pattern]
+
+        with connection.cursor() as cursor:
+            cursor.execute(sql, params)
+        self.clear_cache()
 
     @classmethod
-    def delete_tree(cls):
+    def delete_forest(cls):
         """Delete the whole tree for the current node class."""
-        cls.clear_cache()
-        cls.objects.all().delete()
+        # cls.objects.all()._raw_delete(cls._base_manager.db)
+        table = cls._meta.db_table
+        with connection.cursor() as cursor:
+            cursor.execute(f"DELETE FROM {table}")
+        cache.invalidate(cls._meta.label)
 
 # The end

treenode/models/mixins/update.py

@@ -0,0 +1,154 @@
+# -*- coding: utf-8 -*-
+"""
+TreeNode Raw SQL Mixin
+
+Version: 3.0.0
+Author: Timur Kady
+Email: timurkady@yandex.com
+"""
+
+from django.db import models, connection
+
+from ...settings import SEGMENT_LENGTH, BASE
+from ...utils.db.sqlcompat import SQLCompat
+
+
+class RawSQLMixin(models.Model):
+    """Raw SQL Mixin."""
+
+    class Meta:
+        """Meta Class."""
+
+        abstract = True
+
+    def refresh(self):
+        """Refresh key fields from DB."""
+        task_query = self._meta.model.tasks
+        if len(task_query.queue) > 0:
+            task_query.run()
+
+        table = self._meta.db_table
+        sql = f"SELECT priority, _path, _depth FROM {table} WHERE id = %s"
+        with connection.cursor() as cursor:
+            cursor.execute(sql, [self.pk])
+            row = cursor.fetchone()
+            self.priority, self._path, self._depth = row
+
+        self._parent_id = self.parent_pk
+        self._priority = self.priority
+
+    def _shift_siblings_forward(self):
+        """
+        Shift the priority of all siblings starting from self.priority.
+
+        Uses direct SQL for maximum speed.
+        """
+        if (self.priority is None) or (self.priority >= BASE - 1):
+            return
+
+        db_table = self._meta.db_table
+
+        if self.parent_id is None:
+            where_clause = "parent_id IS NULL"
+            params = [self.priority]
+        else:
+            where_clause = "parent_id = %s"
+            params = [self.parent_id, self.priority]
+
+        sql = f"""
+            UPDATE {db_table}
+            SET priority = priority + 1
+            WHERE {where_clause} AND priority >= %s
+        """
+
+        with connection.cursor() as cursor:
+            cursor.execute(sql, params)
+
+    def _update_path(self, parent_id):
+        """
+        Rebuild subtree starting from parent_id.
+
+        If parent_id=None, then the whole tree is rebuilt.
+        Only fields are used: parent_id and id. All others (priority, _path,
+        _depth) are recalculated.
+        """
+        db_table = self._meta.db_table
+
+        sorting_field = self.sorting_field
+        sorting_fields = ["priority", "id"] if sorting_field == "priority" else [sorting_field]   # noqa: D501
+        sort_expr = ", ".join([
+            field if "." in field else f"c.{field}"
+            for field in sorting_fields
+        ])
+
+        cte_header = "(id, parent_id, new_priority, new_path, new_depth)"
+
+        row_number_expr = "ROW_NUMBER() OVER (ORDER BY {sort_expr}) - 1"
+        hex_expr = SQLCompat.to_hex(row_number_expr)
+        lpad_expr = SQLCompat.lpad(hex_expr, SEGMENT_LENGTH, "'0'")
+
+        if parent_id is None:
+            new_path_expr = lpad_expr
+            base_sql = f"""
+                SELECT
+                    c.id,
+                    c.parent_id,
+                    {row_number_expr} AS new_priority,
+                    {new_path_expr} AS new_path,
+                    0 AS new_depth
+                FROM {db_table} AS c
+                WHERE c.parent_id IS NULL
+            """
+            params = []
+        else:
+            path_expr = SQLCompat.concat("p._path", "'.'", lpad_expr)
+            base_sql = f"""
+                SELECT
+                    c.id,
+                    c.parent_id,
+                    {row_number_expr} AS new_priority,
+                    {path_expr} AS new_path,
+                    p._depth + 1 AS new_depth
+                FROM {db_table} c
+                JOIN {db_table} p ON c.parent_id = p.id
+                WHERE p.id = %s
+            """
+            params = [parent_id]
+
+        recursive_row_number_expr = "ROW_NUMBER() OVER (PARTITION BY c.parent_id ORDER BY {sort_expr}) - 1"   # noqa: D501
+        recursive_hex_expr = SQLCompat.to_hex(recursive_row_number_expr)
+        recursive_lpad_expr = SQLCompat.lpad(
+            recursive_hex_expr, SEGMENT_LENGTH, "'0'")
+        recursive_path_expr = SQLCompat.concat(
+            "t.new_path", "'.'", recursive_lpad_expr)
+
+        recursive_sql = f"""
+            SELECT
+                c.id,
+                c.parent_id,
+                {recursive_row_number_expr} AS new_priority,
+                {recursive_path_expr} AS new_path,
+                t.new_depth + 1 AS new_depth
+            FROM {db_table} c
+            JOIN tree_cte t ON c.parent_id = t.id
+        """
+
+        final_sql = f"""
+            WITH RECURSIVE tree_cte {cte_header} AS (
+                {base_sql}
+                UNION ALL
+                {recursive_sql}
+            )
+            UPDATE {db_table} AS orig
+            SET
+                priority = t.new_priority,
+                _path = t.new_path,
+                _depth = t.new_depth
+            FROM tree_cte t
+            WHERE orig.id = t.id;
+        """
+
+        self.sqlq.append((final_sql.format(sort_expr=sort_expr), params))
+
+
+# The End