diff --git a/tests/testapp/models.py b/tests/testapp/models.py index 473003e..03ded71 100644 --- a/tests/testapp/models.py +++ b/tests/testapp/models.py @@ -122,3 +122,17 @@ class Meta: class InheritConcreteGrandChildModel(InheritAbstractChildModel): pass + + +class RelatedOrderModel(TreeNode): + name = models.CharField(max_length=100) + + +class OneToOneRelatedOrder(models.Model): + relatedmodel = models.OneToOneField( + RelatedOrderModel, + on_delete=models.CASCADE, + primary_key=True, + related_name="related", + ) + order = models.PositiveIntegerField(default=0) diff --git a/tests/testapp/test_queries.py b/tests/testapp/test_queries.py index b4aeca0..30512ef 100644 --- a/tests/testapp/test_queries.py +++ b/tests/testapp/test_queries.py @@ -19,6 +19,8 @@ TreeNodeIsOptional, UnorderedModel, UUIDModel, + RelatedOrderModel, + OneToOneRelatedOrder, ) from tree_queries.compiler import SEPARATOR, TreeQuery from tree_queries.query import pk @@ -55,9 +57,12 @@ def test_no_attributes(self): def test_attributes(self): tree = self.create_tree() - child2_2 = Model.objects.with_tree_fields().get(pk=tree.child2_2.pk) + # Ordering should be deterministic + child2_2 = Model.objects.with_tree_fields().order_siblings_by("order", "pk").get(pk=tree.child2_2.pk) self.assertEqual(child2_2.tree_depth, 2) - self.assertEqual(child2_2.tree_ordering, [0, 1, 1]) + # Tree ordering is an array of the ranks assigned to a comment's + # ancestors when they are ordered without respect for tree relations. + self.assertEqual(child2_2.tree_ordering, [1, 5, 6]) self.assertEqual( child2_2.tree_path, [tree.root.pk, tree.child2.pk, tree.child2_2.pk] ) @@ -649,3 +654,90 @@ def test_polymorphic_queries(self): ("child2_2", [1, 3, 6]), ], ) + + def test_descending_order(self): + tree = self.create_tree() + + nodes = Model.objects.order_siblings_by("-order") + self.assertEqual( + list(nodes), + [ + tree.root, + tree.child2, + tree.child2_2, + tree.child2_1, + tree.child1, + tree.child1_1, + ], + ) + + def test_multi_field_order(self): + tree = type("Namespace", (), {})() # SimpleNamespace for PY2... + + tree.root = MultiOrderedModel.objects.create(name="root") + tree.child1 = MultiOrderedModel.objects.create( + parent=tree.root, first_position=0, second_position=1, name="1" + ) + tree.child2 = MultiOrderedModel.objects.create( + parent=tree.root, first_position=0, second_position=0, name="2" + ) + tree.child1_1 = MultiOrderedModel.objects.create( + parent=tree.child1, first_position=1, second_position=1, name="1-1" + ) + tree.child2_1 = MultiOrderedModel.objects.create( + parent=tree.child2, first_position=0, second_position=1, name="2-1" + ) + tree.child2_2 = MultiOrderedModel.objects.create( + parent=tree.child2, first_position=1, second_position=0, name="2-2" + ) + + nodes = MultiOrderedModel.objects.order_siblings_by("first_position", "-second_position") + self.assertEqual( + list(nodes), + [ + tree.root, + tree.child1, + tree.child1_1, + tree.child2, + tree.child2_1, + tree.child2_2, + ] + ) + + def test_order_by_related(self): + tree = type("Namespace", (), {})() # SimpleNamespace for PY2... + + tree.root = RelatedOrderModel.objects.create(name="root") + tree.child1 = RelatedOrderModel.objects.create(parent=tree.root, name="1") + tree.child1_related = OneToOneRelatedOrder.objects.create( + relatedmodel=tree.child1, order=0 + ) + tree.child2 = RelatedOrderModel.objects.create(parent=tree.root, name="2") + tree.child2_related = OneToOneRelatedOrder.objects.create( + relatedmodel=tree.child2, order=1 + ) + tree.child1_1 = RelatedOrderModel.objects.create(parent=tree.child1, name="1-1") + tree.child1_1_related = OneToOneRelatedOrder.objects.create( + relatedmodel=tree.child1_1, order=0 + ) + tree.child2_1 = RelatedOrderModel.objects.create(parent=tree.child2, name="2-1") + tree.child2_1_related = OneToOneRelatedOrder.objects.create( + relatedmodel=tree.child2_1, order=0 + ) + tree.child2_2 = RelatedOrderModel.objects.create(parent=tree.child2, name="2-2") + tree.child2_2_related = OneToOneRelatedOrder.objects.create( + relatedmodel=tree.child2_2, order=1 + ) + + nodes = RelatedOrderModel.objects.order_siblings_by("related__order") + self.assertEqual( + list(nodes), + [ + tree.root, + tree.child1, + tree.child1_1, + tree.child2, + tree.child2_1, + tree.child2_2, + ] + ) diff --git a/tree_queries/compiler.py b/tree_queries/compiler.py index 8d0579c..800dd61 100644 --- a/tree_queries/compiler.py +++ b/tree_queries/compiler.py @@ -1,7 +1,4 @@ -import warnings - -from django.core.exceptions import FieldDoesNotExist -from django.db import connections, models +from django.db import connections from django.db.models.sql.compiler import SQLCompiler from django.db.models.sql.query import Query @@ -34,49 +31,24 @@ def get_sibling_order(self): return self.sibling_order opts = _find_tree_model(self.model)._meta if opts.ordering: - if len(opts.ordering) > 1: - warnings.warn( - f"Model '{opts.app_label}.{opts.model_name}' is used" - " with django-tree-queries and has more than one field in" - " Meta.ordering. django-tree-queries only uses the first" - " field in the list.", - RuntimeWarning, - stacklevel=1, - ) - return opts.ordering[0] + return opts.ordering return opts.pk.attname class TreeCompiler(SQLCompiler): - CTE_POSTGRESQL_WITH_TEXT_ORDERING = """ - WITH RECURSIVE __tree ( - "tree_depth", - "tree_path", - "tree_ordering", - "tree_pk" + CTE_POSTGRESQL_WITH_INTEGER_ORDERING = """ + WITH RECURSIVE __rank_table( + "{pk}", + "{parent}", + "rank_order" ) AS ( SELECT - 0 AS tree_depth, - array[T.{pk}] AS tree_path, - array[{order_by}]::text[] AS tree_ordering, - T."{pk}" - FROM {db_table} T - WHERE T."{parent}" IS NULL - - UNION ALL - - SELECT - __tree.tree_depth + 1 AS tree_depth, - __tree.tree_path || T.{pk}, - __tree.tree_ordering || {order_by}::text, - T."{pk}" - FROM {db_table} T - JOIN __tree ON T."{parent}" = __tree.tree_pk - ) - """ - - CTE_POSTGRESQL_WITH_INTEGER_ORDERING = """ - WITH RECURSIVE __tree ( + {rank_pk}, + {rank_parent}, + ROW_NUMBER() OVER (ORDER BY {rank_order_by}) + FROM {rank_from} + ), + __tree ( "tree_depth", "tree_path", "tree_ordering", @@ -85,9 +57,9 @@ class TreeCompiler(SQLCompiler): SELECT 0 AS tree_depth, array[T.{pk}] AS tree_path, - array[{order_by}] AS tree_ordering, + array[T.rank_order] AS tree_ordering, T."{pk}" - FROM {db_table} T + FROM __rank_table T WHERE T."{parent}" IS NULL UNION ALL @@ -95,47 +67,30 @@ class TreeCompiler(SQLCompiler): SELECT __tree.tree_depth + 1 AS tree_depth, __tree.tree_path || T.{pk}, - __tree.tree_ordering || {order_by}, + __tree.tree_ordering || T.rank_order, T."{pk}" - FROM {db_table} T + FROM __rank_table T JOIN __tree ON T."{parent}" = __tree.tree_pk ) """ CTE_MYSQL_WITH_INTEGER_ORDERING = """ - WITH RECURSIVE __tree(tree_depth, tree_path, tree_ordering, tree_pk) AS ( - SELECT - 0, - -- Limit to max. 50 levels... - CAST(CONCAT("{sep}", {pk}, "{sep}") AS char(1000)), - CAST(CONCAT("{sep}", LPAD(CONCAT({order_by}, "{sep}"), 20, "0")) - AS char(1000)), - T.{pk} - FROM {db_table} T - WHERE T.{parent} IS NULL - - UNION ALL - + WITH RECURSIVE __rank_table({pk}, {parent}, rank_order) AS ( SELECT - __tree.tree_depth + 1, - CONCAT(__tree.tree_path, T2.{pk}, "{sep}"), - CONCAT(__tree.tree_ordering, LPAD(CONCAT(T2.{order_by}, "{sep}"), 20, "0")), - T2.{pk} - FROM __tree, {db_table} T2 - WHERE __tree.tree_pk = T2.{parent} - ) - """ - - CTE_MYSQL_WITH_TEXT_ORDERING = """ - WITH RECURSIVE __tree(tree_depth, tree_path, tree_ordering, tree_pk) AS ( + {rank_pk}, + {rank_parent}, + ROW_NUMBER() OVER (ORDER BY {rank_order_by}) + FROM {rank_from} + ), + __tree(tree_depth, tree_path, tree_ordering, tree_pk) AS ( SELECT 0, -- Limit to max. 50 levels... CAST(CONCAT("{sep}", {pk}, "{sep}") AS char(1000)), - CAST(CONCAT("{sep}", CONCAT({order_by}, "{sep}")) + CAST(CONCAT("{sep}", LPAD(CONCAT(T.rank_order, "{sep}"), 20, "0")) AS char(1000)), T.{pk} - FROM {db_table} T + FROM __rank_table T WHERE T.{parent} IS NULL UNION ALL @@ -143,21 +98,28 @@ class TreeCompiler(SQLCompiler): SELECT __tree.tree_depth + 1, CONCAT(__tree.tree_path, T2.{pk}, "{sep}"), - CONCAT(__tree.tree_ordering, CONCAT(T2.{order_by}, "{sep}")), + CONCAT(__tree.tree_ordering, LPAD(CONCAT(T2.rank_order, "{sep}"), 20, "0")), T2.{pk} - FROM __tree, {db_table} T2 + FROM __tree, __rank_table T2 WHERE __tree.tree_pk = T2.{parent} ) """ CTE_SQLITE3_WITH_INTEGER_ORDERING = """ - WITH RECURSIVE __tree(tree_depth, tree_path, tree_ordering, tree_pk) AS ( + WITH RECURSIVE __rank_table({pk}, {parent}, rank_order) AS ( + SELECT + {rank_pk}, + {rank_parent}, + row_number() OVER (ORDER BY {rank_order_by}) + FROM {rank_from} + ), + __tree(tree_depth, tree_path, tree_ordering, tree_pk) AS ( SELECT 0 tree_depth, printf("{sep}%%s{sep}", {pk}) tree_path, - printf("{sep}%%020s{sep}", {order_by}) tree_ordering, + printf("{sep}%%020s{sep}", T.rank_order) tree_ordering, T."{pk}" tree_pk - FROM {db_table} T + FROM __rank_table T WHERE T."{parent}" IS NULL UNION ALL @@ -165,34 +127,56 @@ class TreeCompiler(SQLCompiler): SELECT __tree.tree_depth + 1, __tree.tree_path || printf("%%s{sep}", T.{pk}), - __tree.tree_ordering || printf("%%020s{sep}", T.{order_by}), + __tree.tree_ordering || printf("%%020s{sep}", T.rank_order), T."{pk}" - FROM {db_table} T + FROM __rank_table T JOIN __tree ON T."{parent}" = __tree.tree_pk ) """ - CTE_SQLITE3_WITH_TEXT_ORDERING = """ - WITH RECURSIVE __tree(tree_depth, tree_path, tree_ordering, tree_pk) AS ( - SELECT - 0 tree_depth, - printf("{sep}%%s{sep}", {pk}) tree_path, - printf("{sep}%%s{sep}", {order_by}) tree_ordering, - T."{pk}" tree_pk - FROM {db_table} T - WHERE T."{parent}" IS NULL + def get_sibling_order_params(self): + ''' + This method uses a simple django queryset to generate sql + that can be used to create the __rank_table that orders + siblings. This is done so that any joins required by order_by + are pre-calculated by django + ''' + sibling_order = self.query.get_sibling_order() + + if isinstance(sibling_order, (list, tuple)): + order_fields = sibling_order + elif isinstance(sibling_order, str): + order_fields = [sibling_order] + else: + raise ValueError("Sibling order must be a string or a list or tuple of strings.") + + # Use Django to make a SQL query whose parts can be repurposed for __rank_table + base_query = _find_tree_model(self.query.model).objects.only("pk", "parent").order_by(*order_fields).query + + # Use the base compiler because we want vanilla sql and want to avoid recursion. + base_compiler = SQLCompiler(base_query, self.connection, None) + base_sql, base_params = base_compiler.as_sql() + result_sql = base_sql % base_params + + # Split the base SQL string on the SQL keywords 'FROM' and 'ORDER BY' + from_split = result_sql.split("FROM") + order_split = from_split[1].split("ORDER BY") + + # Identify the FROM and ORDER BY parts of the base SQL + ordering_params = { + "rank_from": order_split[0].strip(), + "rank_order_by": order_split[1].strip(), + } - UNION ALL + # Identify the primary key field and parent_id field from the SELECT section + base_select = from_split[0][6:] + for field in base_select.split(","): + if "parent_id" in field: # XXX Taking advantage of Hardcoded. + ordering_params["rank_parent"] = field.strip() + else: + ordering_params["rank_pk"] = field.strip() - SELECT - __tree.tree_depth + 1, - __tree.tree_path || printf("%%s{sep}", T.{pk}), - __tree.tree_ordering || printf("%%s{sep}", T.{order_by}), - T."{pk}" - FROM {db_table} T - JOIN __tree ON T."{parent}" = __tree.tree_pk - ) - """ + return ordering_params def as_sql(self, *args, **kwargs): # The general idea is that if we have a summary query (e.g. .count()) @@ -223,10 +207,12 @@ def as_sql(self, *args, **kwargs): "parent": "parent_id", # XXX Hardcoded. "pk": opts.pk.attname, "db_table": opts.db_table, - "order_by": self.query.get_sibling_order(), "sep": SEPARATOR, } + # Add ordering params to params + params.update(self.get_sibling_order_params()) + if "__tree" not in self.query.extra_tables: # pragma: no branch - unlikely tree_params = params.copy() @@ -265,29 +251,16 @@ def as_sql(self, *args, **kwargs): ) if self.connection.vendor == "postgresql": - cte = ( - self.CTE_POSTGRESQL_WITH_INTEGER_ORDERING - if _ordered_by_integer(opts, params) - else self.CTE_POSTGRESQL_WITH_TEXT_ORDERING - ) + cte = self.CTE_POSTGRESQL_WITH_INTEGER_ORDERING elif self.connection.vendor == "sqlite": - cte = ( - self.CTE_SQLITE3_WITH_INTEGER_ORDERING - if _ordered_by_integer(opts, params) - else self.CTE_SQLITE3_WITH_TEXT_ORDERING - ) + cte = self.CTE_SQLITE3_WITH_INTEGER_ORDERING elif self.connection.vendor == "mysql": - cte = ( - self.CTE_MYSQL_WITH_INTEGER_ORDERING - if _ordered_by_integer(opts, params) - else self.CTE_MYSQL_WITH_TEXT_ORDERING - ) - if params["order_by"]: - params["order_by"] = self.connection.ops.quote_name(params["order_by"]) + cte = self.CTE_MYSQL_WITH_INTEGER_ORDERING sql_0, sql_1 = super().as_sql(*args, **kwargs) explain = "" if sql_0.startswith("EXPLAIN "): explain, sql_0 = sql_0.split(" ", 1) + return ("".join([explain, cte.format(**params), sql_0]), sql_1) def get_converters(self, expressions): @@ -315,11 +288,3 @@ def converter(value, expression, connection, context=None): return [int(v) for v in value] # Maybe Field.to_python()? except ValueError: return value - - -def _ordered_by_integer(opts, params): - try: - ordering_field = opts.get_field(params["order_by"]) - return isinstance(ordering_field, models.IntegerField) - except FieldDoesNotExist: - return False diff --git a/tree_queries/query.py b/tree_queries/query.py index 4b2022d..250ad48 100644 --- a/tree_queries/query.py +++ b/tree_queries/query.py @@ -37,12 +37,12 @@ def without_tree_fields(self): """ return self.with_tree_fields(tree_fields=False) - def order_siblings_by(self, order_by): + def order_siblings_by(self, *order_by): """ Sets TreeQuery sibling_order attribute - Pass the name of a single model field as a string - to order tree siblings by that model field + Pass the names of model fields as a list of strings + to order tree siblings by those model fields """ self.query.__class__ = TreeQuery self.query.sibling_order = order_by