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)