klbr/khoj
1
0
Fork 0
mirror of https://github.com/khoaliber/khoj.git synced 2026-03-08 05:39:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix async batch delete of indexed entries

Browse Source
This commit is contained in:
Debanjum Singh Solanky
2024-07-06 22:45:26 +05:30
parent 08b379c2ab
commit a31cd0dec1
1 changed files with 10 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/khoj/database/adapters/__init__.py
View File

@@ -1012,7 +1012,7 @@ class EntryAdapters:
return deleted_count return deleted_count
@staticmethod @staticmethod
def get_entries_by_batch(user: KhojUser, batch_size: int, file_type: str = None, file_source: str = None): def get_filtered_entries(user: KhojUser, file_type: str = None, file_source: str = None):
queryset = Entry.objects.filter(user=user) queryset = Entry.objects.filter(user=user)
if file_type is not None: if file_type is not None:
@@ -1021,14 +1021,15 @@ class EntryAdapters:
if file_source is not None: if file_source is not None:
queryset = queryset.filter(file_source=file_source) queryset = queryset.filter(file_source=file_source)
while queryset.exists(): return queryset
batch_ids = list(queryset.values_list("id", flat=True)[:batch_size])
yield Entry.objects.filter(id__in=batch_ids)
@staticmethod @staticmethod
def delete_all_entries(user: KhojUser, file_type: str = None, file_source: str = None, batch_size=1000): def delete_all_entries(user: KhojUser, file_type: str = None, file_source: str = None, batch_size=1000):
deleted_count = 0 deleted_count = 0
for batch in EntryAdapters.get_entries_by_batch(user, batch_size, file_type, file_source): queryset = EntryAdapters.get_filtered_entries(user, file_type, file_source)
while queryset.exists():
batch_ids = list(queryset.values_list("id", flat=True)[:batch_size])
batch = Entry.objects.filter(id__in=batch_ids, user=user)
count, _ = batch.delete() count, _ = batch.delete()
deleted_count += count deleted_count += count
return deleted_count return deleted_count
@@ -1036,7 +1037,10 @@ class EntryAdapters:
@staticmethod @staticmethod
async def adelete_all_entries(user: KhojUser, file_type: str = None, file_source: str = None, batch_size=1000): async def adelete_all_entries(user: KhojUser, file_type: str = None, file_source: str = None, batch_size=1000):
deleted_count = 0 deleted_count = 0
async for batch in EntryAdapters.get_entries_by_batch(user, batch_size, file_type, file_source): queryset = EntryAdapters.get_filtered_entries(user, file_type, file_source)
while await queryset.aexists():
batch_ids = await sync_to_async(list)(queryset.values_list("id", flat=True)[:batch_size])
batch = Entry.objects.filter(id__in=batch_ids, user=user)
count, _ = await batch.adelete() count, _ = await batch.adelete()
deleted_count += count deleted_count += count
return deleted_count return deleted_count