Move TreeNode next-value logic into PL/pgSQL

The purpose of this was to solve a bug when more than 10 entries at
any given level (including the root level) were inserted. The previous
next-value logic sorted lexicographically but didn't pad values when
converted back from integers, so 1000 got sorted between 100 and 200.
This led to constraint violations when more than 10 rows were
inserted.

Because the previous Python code that generated a SQL expression to
calculate the next value was a nightmare to read, and relied on a dirty
hack overriding parts of the Django ORM's internals, I took this
opportunity to rewrite it as a PL/pgSQL function; this means it can
actually be written in an imperative manner where appropriate, without
having to do multiple round-trips to the database.

Author: daisylb

django_pgtree/migrations/0001_initial.py

@@ -0,0 +1,58 @@
+# Generated by Django 2.1.2 on 2018-10-25 04:58
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = []
+
+    operations = [
+        migrations.RunSQL(
+            """
+            CREATE OR REPLACE FUNCTION djpgtree_next(
+                tbl regclass,
+                prefix ltree,
+                gap bigint,
+                pad_length int
+            ) RETURNS ltree AS $function$
+                DECLARE
+                    sibling_query lquery;
+                    previous_highest ltree;
+                    previous_rightmost_label text;
+                    next_rightmost_segment text;
+                BEGIN
+                    -- Generate a lquery that matches all would-be siblings of
+                    -- the new row
+                    IF prefix = ''::ltree THEN
+                        sibling_query = '*{1}';
+                    ELSE
+                        sibling_query = prefix::text || '.*{1}';
+                    END IF;
+
+                    -- Find the existing sibling with the highest tree_path
+                    EXECUTE format($$
+                        SELECT tree_path
+                        FROM %s
+                        WHERE tree_path ~ %L
+                        ORDER BY tree_path DESC LIMIT 1
+                    $$, tbl, sibling_query) INTO previous_highest;
+
+                    IF previous_highest IS NULL THEN
+                        -- If there is no such row, start at the gap we were passed in,
+                        -- to allow room for other rows to be moved above us
+                        RETURN prefix || LPAD(gap::text, pad_length, '0');
+                    ELSE
+                        -- Otherwise, parse the rightmost label as a number, adding
+                        -- the gap to it, and reattach to the prefix
+                        previous_rightmost_label = subpath(previous_highest, -1);
+                        next_rightmost_segment = previous_rightmost_label::bigint + gap;
+                        RETURN prefix || LPAD(next_rightmost_segment, pad_length, '0');
+                    END IF;
+                END
+            $function$ LANGUAGE plpgsql;
+
+        """,
+            "DROP FUNCTION djpgtree_next(regclass, ltree, bigint, int)",
+        )
+    ]

django_pgtree/migrations/__init__.py

@@ -1,11 +1,11 @@
 from django.contrib.postgres.indexes import GistIndex
 from django.db import models
-from django.db.models import functions as f
 from django.db.transaction import atomic
 
 from .fields import LtreeField
 
-GAP = 1000000000
+GAP = 10 ** 9
+PAD_LENGTH = 18
 
 
 class LtreeConcat(models.Func):
@@ -21,6 +21,10 @@ class Text2Ltree(models.Func):
     function = "text2ltree"
 
 
+class DjPgTreeNext(models.Func):
+    function = "djpgtree_next"
+
+
 class TreeQuerySet(models.QuerySet):
     def roots(self):
         return self.filter(tree_path__matches_lquery=["*{1}"])
@@ -56,47 +60,14 @@ def parent(self, new_parent):
         # Replace our tree_path with a new one that has our new parent's
         self.__new_parent = new_parent
 
-    def __next_tree_path_qx(self, prefix=None):
-        if prefix is None:
-            prefix = []
-
-        # These are all the siblings of the target position, in reverse tree order.
-        # If we don't have a prefix, this will be all root nodes.
-        sibling_queryset = self.__class__.objects.filter(
-            tree_path__matches_lquery=[*prefix, "*{1}"]
-        ).order_by("-tree_path")
-        # This query expression is the full ltree of the last sibling by tree order.
-        last_sibling_tree_path = models.Subquery(
-            sibling_queryset.values("tree_path")[:1]
-        )
-
-        # Django doesn't allow the use of column references in an INSERT statement,
-        # because it makes the assumption that they refer to columns in the
-        # to-be-inserted row, the values for which aren't yet known.
-        # Unfortunately, this means we can't use a subquery that refers to column
-        # values anywhere internally, even though the columns it refers to are subquery
-        # result columns. To get around this, we override the contains_column_references
-        # property on the subquery with a static False, so that Django's check doesn't
-        # cross the subquery boundary.
-        last_sibling_tree_path.contains_column_references = False
-
-        # This query expression is the rightmost component of that ltree. The double
-        # cast is because PostgreSQL doesn't let you cast directly from ltree to bigint.
-        last_sibling_last_value = f.Cast(
-            f.Cast(Subpath(last_sibling_tree_path, -1), models.CharField()),
-            models.BigIntegerField(),
-        )
-        # This query expression is an ltree containing that value, plus GAP, or just
-        # GAP if there is no existing siblings. Again, we need to double cast.
-        new_last_value = Text2Ltree(
-            f.Cast(f.Coalesce(last_sibling_last_value, 0) + (GAP), models.CharField())
+    def __next_tree_path_qx(self, prefix=()):
+        return DjPgTreeNext(
+            models.Value(self._meta.db_table),
+            models.Value(".".join(prefix)),
+            GAP,
+            PAD_LENGTH,
         )
 
-        # If we have a prefix, we prepend that to the resulting ltree.
-        if not prefix:
-            return new_last_value
-        return LtreeConcat(models.Value(".".join(prefix)), new_last_value)
-
     def relocate(self, *, after=None, before=None):
         if after is None and before is None:
             raise ValueError("You must supply at least one of before or after")
@@ -133,11 +104,15 @@ def relocate(self, *, after=None, before=None):
 
         next_v = int(new_next_child.tree_path[-1])
         if new_prev_child is None:
-            self.tree_path = new_next_child.tree_path[:-1] + [str(next_v // 2)]
+            self.tree_path = new_next_child.tree_path[:-1] + [
+                str(next_v // 2).zfill(PAD_LENGTH)
+            ]
         else:
             prev_v = int(new_prev_child.tree_path[-1])
             this_v = prev_v + (next_v - prev_v) // 2
-            self.tree_path = new_prev_child.tree_path[:-1] + [str(this_v)]
+            self.tree_path = new_prev_child.tree_path[:-1] + [
+                str(this_v).zfill(PAD_LENGTH)
+            ]
 
     def save(self, *args, **kwargs):  # pylint: disable=arguments-differ
         tree_path_needs_refresh = False

django_pgtree/models.py

@@ -93,3 +93,9 @@ def test_relocate_in_between(animal):
     seal.relocate(after=cat)
     seal.save()
     assert [x.name for x in cat.parent.children] == ["Cat", "Seal", "Dog", "Bear"]
+
+
+def test_ordering_past_10():
+    for i in range(1, 12):
+        T.objects.create(name=str(i))
+    assert [x.name for x in T.objects.all()] == [str(x) for x in range(1, 12)]

django_pgtree/tests.py

@@ -1,7 +1,7 @@
 Installing django-pgtree
 ========================
 
-First, install the ``django-pgtree`` package from PyPI using your package manager of choice:
+Install the ``django-pgtree`` package from PyPI using your package manager of choice:
 
 .. code-block:: sh
 
@@ -11,11 +11,10 @@ First, install the ``django-pgtree`` package from PyPI using your package manage
     # or
     $ pipenv install django-pgtree
 
-.. note::
 
-    Currently, you don't have to add anything to ``INSTALLED_APPS``.
+Then, add ``django_pgtree`` to ``INSTALLED_APPS``.
 
-Second, run the following SQL command on any database that's going to use django-pgtree fields or models:
+On your development database and when you first deploy, you'll also need to run the following SQL command on any database that's going to use django-pgtree fields or models:
 
 .. code-block:: sql
 

docs/install.rst

@@ -37,6 +37,7 @@
     "django.contrib.sessions",
     "django.contrib.messages",
     "django.contrib.staticfiles",
+    "django_pgtree",
     "testproject.testapp",
 ]