authentik.lib.utils.db

authentik database utilities

View Source

 1"""authentik database utilities"""
 2
 3import gc
 4from collections.abc import Generator
 5
 6from django.db import reset_queries
 7from django.db.models import Model, QuerySet
 8
 9
10def chunked_queryset[T: Model](queryset: QuerySet[T], chunk_size: int = 1_000) -> Generator[T]:
11    if not queryset.exists():
12        return []
13
14    def get_chunks(qs: QuerySet) -> Generator[QuerySet[T]]:
15        qs = qs.order_by("pk")
16        pks = qs.values_list("pk", flat=True)
17        start_pk = pks[0]
18        while True:
19            try:
20                end_pk = pks.filter(pk__gte=start_pk)[chunk_size]
21            except IndexError:
22                break
23            yield qs.filter(pk__gte=start_pk, pk__lt=end_pk)
24            start_pk = end_pk
25        yield qs.filter(pk__gte=start_pk)
26
27    for chunk in get_chunks(queryset):
28        reset_queries()
29        gc.collect()
30        yield from chunk.iterator(chunk_size=chunk_size)

def chunked_queryset( queryset: django.db.models.query.QuerySet, chunk_size: int = 1000) -> Generator[T]: View Source

11def chunked_queryset[T: Model](queryset: QuerySet[T], chunk_size: int = 1_000) -> Generator[T]:
12    if not queryset.exists():
13        return []
14
15    def get_chunks(qs: QuerySet) -> Generator[QuerySet[T]]:
16        qs = qs.order_by("pk")
17        pks = qs.values_list("pk", flat=True)
18        start_pk = pks[0]
19        while True:
20            try:
21                end_pk = pks.filter(pk__gte=start_pk)[chunk_size]
22            except IndexError:
23                break
24            yield qs.filter(pk__gte=start_pk, pk__lt=end_pk)
25            start_pk = end_pk
26        yield qs.filter(pk__gte=start_pk)
27
28    for chunk in get_chunks(queryset):
29        reset_queries()
30        gc.collect()
31        yield from chunk.iterator(chunk_size=chunk_size)