def test_recursive_cte_query(self):
def make_regions_cte(cte):
return Region.objects.filter(
parent__isnull=True
).values(
"name",
path=F("name"),
depth=Value(0, output_field=int_field),
).union(
cte.join(Region, parent=cte.col.name).values(
"name",
path=Concat(
cte.col.path, Value("\x01"), F("name"),
output_field=text_field,
),
depth=cte.col.depth + Value(1, output_field=int_field),
),
all=True,
)
cte = With.recursive(make_regions_cte)
regions = cte.join(Region, name=cte.col.name).with_cte(cte).annotate(
path=cte.col.path,
depth=cte.col.depth,
).filter(depth=2).order_by("path")
data = [(r.name, r.path.split("\x01"), r.depth) for r in regions]
self.assertEqual(data, [
('moon', ['sun', 'earth', 'moon'], 2),
('deimos', ['sun', 'mars', 'deimos'], 2),
('phobos', ['sun', 'mars', 'phobos'], 2),
])
def get_recursive_subgroups(user_group: UserGroup) -> "QuerySet[UserGroup]":
cte = With.recursive(
lambda cte: UserGroup.objects.filter(id=user_group.id)
.values("id")
.union(cte.join(UserGroup, direct_supergroups=cte.col.id).values("id"))
)
return cte.join(UserGroup, id=cte.col.id).with_cte(cte)
def get_recursive_membership_groups(user_profile: UserProfile) -> "QuerySet[UserGroup]":
cte = With.recursive(
lambda cte: user_profile.direct_groups.values("id").union(
cte.join(UserGroup, direct_subgroups=cte.col.id).values("id")
)
)
return cte.join(UserGroup, id=cte.col.id).with_cte(cte)
def test_attname_should_not_mask_col_name(self):
def make_regions_cte(cte):
return Region.objects.filter(
name="moon"
).values(
"name",
"parent_id",
).union(
cte.join(Region, name=cte.col.parent_id).values(
"name",
"parent_id",
),
all=True,
)
cte = With.recursive(make_regions_cte)
regions = (
Region.objects.all()
.with_cte(cte)
.annotate(_ex=Exists(
cte.queryset()
.values(value=Value("1", output_field=int_field))
.filter(name=OuterRef("name"))
))
.filter(_ex=True)
.order_by("name")
)
print(regions.query)
data = [r.name for r in regions]
self.assertEqual(data, ['earth', 'moon', 'sun'])
def test_recursive_cte_reference_in_condition(self):
def make_regions_cte(cte):
return Region.objects.filter(
parent__isnull=True
).values(
"name",
path=F("name"),
depth=Value(0, output_field=int_field),
is_planet=Value(0, output_field=int_field),
).union(
cte.join(
Region, parent=cte.col.name
).annotate(
# annotations for filter and CASE/WHEN conditions
parent_name=ExpressionWrapper(
cte.col.name,
output_field=text_field,
),
parent_depth=ExpressionWrapper(
cte.col.depth,
output_field=int_field,
),
).filter(
~Q(parent_name="mars"),
).values(
"name",
path=Concat(
cte.col.path, Value("\x01"), F("name"),
output_field=text_field,
),
depth=cte.col.depth + Value(1, output_field=int_field),
is_planet=Case(
When(parent_depth=0, then=Value(1)),
default=Value(0),
output_field=int_field,
),
),
all=True,
)
cte = With.recursive(make_regions_cte)
regions = cte.join(Region, name=cte.col.name).with_cte(cte).annotate(
path=cte.col.path,
depth=cte.col.depth,
is_planet=cte.col.is_planet,
).order_by("path")
data = [(r.path.split("\x01"), r.is_planet) for r in regions]
print(data)
self.assertEqual(data, [
(["bernard's star"], 0),
(['proxima centauri'], 0),
(['proxima centauri', 'proxima centauri b'], 1),
(['sun'], 0),
(['sun', 'earth'], 1),
(['sun', 'earth', 'moon'], 0),
(['sun', 'mars'], 1), # mars moons excluded: parent_name != 'mars'
(['sun', 'mercury'], 1),
(['sun', 'venus'], 1),
])
def test_circular_ref_error(self):
def make_bad_cte(cte):
# NOTE: not a valid recursive CTE query
return cte.join(Region, parent=cte.col.name).values(
depth=cte.col.depth + 1,
)
cte = With.recursive(make_bad_cte)
regions = cte.join(Region, name=cte.col.name).with_cte(cte)
with self.assertRaises(ValueError) as context:
print(regions.query)
self.assertIn("Circular reference:", str(context.exception))
def get_ancestors(self, node, ascending=False, include_self=False):
"""Query node ancestors
:param node: A model instance or a QuerySet or Q object querying
the adjacency list model. If a QuerySet, it should query a
single value with something like `.values('id')`. If Q the
`include_self` argument will be ignored.
:param ascending: Order of results. The default (`False`) gets
results in descending order (root ancestor first, immediate
parent last).
:param include_self:
:returns: A `QuerySet` instance.
"""
if isinstance(node, Q):
where = node
elif include_self:
if isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(id__in=node.order_by())
else:
where = Q(id=node.id)
elif isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(id__in=node.order_by().values("parent_id"))
else:
where = Q(id=node.parent_id)
def make_cte_query(cte):
return self.filter(where).order_by().annotate(
_depth=Value(0, output_field=field),
).union(
cte.join(
self.all().order_by(),
id=cte.col.parent_id,
).annotate(
_depth=cte.col._depth + Value(1, output_field=field),
),
)
cte = With.recursive(make_cte_query)
return (
cte.queryset()
.with_cte(cte)
.order_by(("" if ascending else "-") + "_depth")
)
def test_alias_as_subquery(self):
# This test covers CTEColumnRef.relabeled_clone
def make_regions_cte(cte):
return KeyPair.objects.filter(
parent__key="level 1",
).annotate(
rank=F('value'),
).union(
cte.join(
KeyPair.objects.all().order_by(),
parent_id=cte.col.id,
).annotate(
rank=F('value'),
),
all=True,
)
cte = With.recursive(make_regions_cte)
children = cte.queryset().with_cte(cte)
xdups = With(cte.queryset().filter(
parent__key="level 1",
).annotate(
rank=F('value')
).values('id', 'rank'), name='xdups')
children = children.annotate(
_exclude=Exists(
(
xdups.queryset().filter(
id=OuterRef("id"),
rank=OuterRef("rank"),
)
)
)
).filter(_exclude=True).with_cte(xdups)
print(children.query)
query = KeyPair.objects.filter(parent__in=children)
print(query.query)
print(children.query)
self.assertEqual(query.get().key, 'level 3')
# Tests the case in which children's query was modified since it was
# used in a subquery to define `query` above.
self.assertEqual(
list(c.key for c in children),
['level 2', 'level 2']
)
def test_recursive_cte_with_empty_union_part(self):
def make_regions_cte(cte):
return Region.objects.none().union(
cte.join(Region, parent=cte.col.name),
all=True,
)
cte = With.recursive(make_regions_cte)
regions = cte.join(Region, name=cte.col.name).with_cte(cte)
print(regions.query)
try:
self.assertEqual(regions.count(), 0)
except DatabaseError:
raise SkipTest(
"Expected failure: QuerySet omits `EmptyQuerySet` from "
"UNION queries resulting in invalid CTE SQL"
)
def test_pickle_recursive_cte_queryset(self):
def make_regions_cte(cte):
return Region.objects.filter(parent__isnull=True).annotate(
depth=Value(0, output_field=int_field), ).union(
cte.join(Region, parent=cte.col.name).annotate(
depth=cte.col.depth +
Value(1, output_field=int_field), ),
all=True,
)
cte = With.recursive(make_regions_cte)
regions = cte.queryset().with_cte(cte).filter(depth=2).order_by("name")
pickled_qs = pickle.loads(pickle.dumps(regions))
data = [(r.name, r.depth) for r in pickled_qs]
self.assertEqual(data, [(r.name, r.depth) for r in regions])
self.assertEqual(data, [('deimos', 2), ('moon', 2), ('phobos', 2)])
def test_2_django_cte(self):
def make_cte(cte):
shinano_gold = Apple.objects.get(name='シナノゴールド')
return Apple.objects.filter(id=shinano_gold.pk).annotate(
node=Value(0, output_field=IntegerField()), ).union(
cte.join(Apple, id=cte.col.parent_id).annotate(
node=cte.col.node +
Value(1, output_field=IntegerField())),
all=True,
)
cte = With.recursive(make_cte)
apples = (cte.queryset().with_cte(cte).annotate(
node=cte.col.node).order_by('node'))
self.assertCte(apples)
self.print_query()
def test_4_cte_to_dict(self):
def make_cte(cte):
shinano_gold = Apple.objects.get(name='シナノゴールド')
return Apple.objects.filter(id=shinano_gold.pk).values(
'id',
'parent',
'name',
node=Value(0, output_field=IntegerField()),
).union(
cte.join(Apple, id=cte.col.parent_id).values(
'id',
'parent',
'name',
node=cte.col.node + Value(1, output_field=IntegerField())),
all=True,
)
cte = With.recursive(make_cte)
apples = (cte.queryset().with_cte(cte).annotate(
node=cte.col.node).order_by('node'))
# 途中で values() を使って dict 化しているため、共通の assertCte() は使えない
# 件数
self.assertEqual(len(apples), 3)
# シナノゴールド自身があること
own = apples[0]
self.assertEqual(own['node'], 0)
self.assertEqual(own['name'], 'シナノゴールド')
# シナノゴールドの親(千秋)
own = apples[1]
self.assertEqual(own['node'], 1)
self.assertEqual(own['name'], '千秋')
# 千秋の親(東光)
own = apples[2]
self.assertEqual(own['node'], 2)
self.assertEqual(own['name'], '東光')
self.print_query()
def _cte_get_ancestors(self, node, ascending=False, include_self=False):
"""Query node ancestors
:param node: A model instance or a QuerySet or Q object querying
the adjacency list model. If a QuerySet, it should query a
single value with something like `.values('id')`. If Q the
`include_self` argument will be ignored.
:param ascending: Order of results. The default (`False`) gets
results in descending order (root ancestor first, immediate
parent last).
:param include_self:
:returns: A `QuerySet` instance.
"""
if isinstance(node, Q):
where = node
elif include_self:
if isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(id__in=node.order_by())
else:
where = Q(id=node.id)
elif isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(id__in=node.order_by().values("parent_id"))
else:
where = Q(id=node.parent_id)
def make_cte_query(cte):
return self.filter(where).order_by().annotate(_depth=Value(
0, output_field=field), ).union(
cte.join(
self.all().order_by(),
id=cte.col.parent_id,
).annotate(_depth=cte.col._depth +
Value(1, output_field=field), ), )
cte = With.recursive(make_cte_query)
return (
cte.queryset().with_cte(cte).order_by(("" if ascending else "-") +
"_depth"))
def test_3_django_cte_root(self):
def make_cte(cte):
kokko = Apple.objects.get(name='国光')
return Apple.objects.filter(id=kokko.pk).annotate(node=Value(
0, output_field=IntegerField()), ).union(
cte.join(Apple, id=cte.col.parent_id).annotate(
node=cte.col.node +
Value(1, output_field=IntegerField())),
all=True,
)
cte = With.recursive(make_cte)
apples = (cte.queryset().with_cte(cte).annotate(
node=cte.col.node).order_by('node'))
self.assertEqual(len(apples), 1)
apple = apples.get()
self.assertEqual(apple.node, 0)
self.assertEqual(apple.name, '国光')
self.print_query()
def test_alias_change_in_annotation(self):
def make_regions_cte(cte):
return Region.objects.filter(
parent__name="sun",
).annotate(
value=F('name'),
).union(
cte.join(
Region.objects.all().annotate(
value=F('name'),
),
parent_id=cte.col.name,
),
all=True,
)
cte = With.recursive(make_regions_cte)
query = cte.queryset().with_cte(cte)
exclude_leaves = With(cte.queryset().filter(
parent__name='sun',
).annotate(
value=Concat(F('name'), F('name'))
), name='value_cte')
query = query.annotate(
_exclude_leaves=Exists(
exclude_leaves.queryset().filter(
name=OuterRef("name"),
value=OuterRef("value"),
)
)
).filter(_exclude_leaves=True).with_cte(exclude_leaves)
print(query.query)
# Nothing should be returned.
self.assertFalse(query)
def test_named_ctes(self):
def make_paths_cte(paths):
return Region.objects.filter(parent__isnull=True).values(
"name",
path=F("name"),
).union(
paths.join(Region, parent=paths.col.name).values(
"name",
path=Concat(
paths.col.path,
Value(" "),
F("name"),
output_field=text_field,
),
),
all=True,
)
paths = With.recursive(make_paths_cte, name="region_paths")
def make_groups_cte(groups):
return paths.join(Region, name=paths.col.name).values(
"name",
parent_path=paths.col.path,
parent_name=F("name"),
).union(
groups.join(Region, parent=groups.col.name).values(
"name",
parent_path=groups.col.parent_path,
parent_name=groups.col.parent_name,
),
all=True,
)
groups = With.recursive(make_groups_cte, name="region_groups")
region_totals = With(
groups.join(Order, region_id=groups.col.name).values(
name=groups.col.parent_name,
path=groups.col.parent_path,
).annotate(
orders_count=Count("id"),
region_total=Sum("amount"),
),
name="region_totals",
)
regions = (
region_totals.join(Region, name=region_totals.col.name).with_cte(
paths).with_cte(groups).with_cte(region_totals).annotate(
path=region_totals.col.path,
# count of orders in this region and all subregions
orders_count=region_totals.col.orders_count,
# sum of order amounts in this region and all subregions
region_total=region_totals.col.region_total,
).order_by("path"))
data = [(r.name, r.orders_count, r.region_total) for r in regions]
self.assertEqual(data, [
('proxima centauri', 4, 2033),
('proxima centauri b', 3, 33),
('sun', 18, 1374),
('earth', 7, 132),
('moon', 3, 6),
('mars', 3, 123),
('mercury', 3, 33),
('venus', 4, 86),
])
def _cte_get_descendants(self, node, include_self=False):
"""Query node descendants
:param node: A model instance or a QuerySet or Q object querying
the adjacency list model. If a QuerySet, it should query a
single value with something like `.values('id')`. If Q the
`include_self` argument will be ignored.
:returns: A `QuerySet` instance.
"""
ordering_col = self.model.ordering_col_attr
discard_dups = False
if isinstance(node, Q):
where = node
discard_dups = True
elif include_self:
if isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(id__in=node.order_by())
discard_dups = True
else:
where = Q(id=node.id)
elif isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(parent_id__in=node.order_by())
discard_dups = True
else:
where = Q(parent_id=node.id)
def make_cte_query(cte):
return self.filter(where).order_by().annotate(
_cte_ordering=str_array(ordering_col), ).union(
cte.join(
self.all().order_by(),
parent_id=cte.col.id,
).annotate(_cte_ordering=array_append(
cte.col._cte_ordering,
F(ordering_col),
)),
all=True,
)
cte = With.recursive(make_cte_query)
query = cte.queryset().with_cte(cte)
if discard_dups:
# Remove duplicates when the supplied Queryset or Q object
# may contain/match both parents and children. For a given
# id, retain the row with the longest path. TODO remove this
# and ensure duplicates do not matter or the criteria never
# matches both parents and children in all calling code.
xdups = With(cte.queryset().annotate(max_len=array_length(
F("_cte_ordering"),
output_field=field), ).distinct("id").order_by(
"id",
"-max_len",
).values(
"id",
"_cte_ordering",
),
name="xdups")
query = query.annotate(
_exclude_dups=Exists(xdups.queryset().filter(
id=OuterRef("id"),
_cte_ordering=OuterRef("_cte_ordering"),
))).filter(_exclude_dups=True).with_cte(xdups)
return query.order_by(cte.col._cte_ordering)
def ancestors(self):
cte = With.recursive(self._get_ancestors)
model = self.__class__
return cte.join(model.tree.all(), pk=cte.col.parent_id).with_cte(cte).annotate(
depth=cte.col.depth
).order_by('depth')
def descendants(self):
cte = With.recursive(self._get_descendants)
model = self.__class__
return cte.join(model.tree.all(), pk=cte.col.pk).with_cte(cte).annotate(
depth=cte.col.depth
).order_by('depth')
def get_descendants(self, node, include_self=False):
"""Query node descendants
:param node: A model instance or a QuerySet or Q object querying
the adjacency list model. If a QuerySet, it should query a
single value with something like `.values('id')`. If Q the
`include_self` argument will be ignored.
:returns: A `QuerySet` instance.
"""
ordering_col = self.model.ordering_col_attr
discard_dups = False
if isinstance(node, Q):
where = node
discard_dups = True
elif include_self:
if isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(id__in=node.order_by())
discard_dups = True
else:
where = Q(id=node.id)
elif isinstance(node, QuerySet):
if _is_empty(node):
return self.none()
where = Q(parent_id__in=node.order_by())
discard_dups = True
else:
where = Q(parent_id=node.id)
def make_cte_query(cte):
return self.filter(where).order_by().annotate(
_cte_ordering=str_array(ordering_col),
).union(
cte.join(
self.all().order_by(),
parent_id=cte.col.id,
).annotate(
_cte_ordering=array_append(
cte.col._cte_ordering,
F(ordering_col),
)
),
all=True,
)
cte = With.recursive(make_cte_query)
query = cte.queryset().with_cte(cte)
if discard_dups:
# Remove duplicates when the supplied Queryset or Q object
# may contain/match both parents and children. For a given
# id, retain the row with the longest path. TODO remove this
# and ensure duplicates do not matter or the criteria never
# matches both parents and children in all calling code.
xdups = With(
cte.queryset().annotate(
max_len=array_length(
F("_cte_ordering"),
output_field=field
),
).distinct("id").order_by(
"id",
"-max_len",
).values(
"id",
"_cte_ordering",
),
name="xdups"
)
query = query.annotate(
_exclude_dups=Exists(SubQueryset(xdups.queryset().filter(
id=OuterRef("id"),
_cte_ordering=OuterRef("_cte_ordering"),
)))
).filter(_exclude_dups=True).with_cte(xdups)
return query.order_by(cte.col._cte_ordering)