Update bulk_generate_documents to use object_keys

wagtail · Sep 13, 2024 · 8b526ed · 8b526ed
1 parent dc7072f
commit 8b526ed
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/wagtail_vector_index/storage/django.py b/src/wagtail_vector_index/storage/django.py
@@ -315,10 +315,11 @@ def generate_documents(
     @transaction.atomic
     def bulk_generate_documents(self, objects, *, embedding_backend):
         objects_by_key = {ModelKey.from_instance(obj): obj for obj in objects}
-        documents = Document.objects.filter(object_key__in=objects_by_key.keys())
+        documents = Document.objects.for_keys(objects_by_key.keys())
+
         documents_by_object_key = defaultdict(list)
         for document in documents:
-            documents_by_object_key[document.object_key].append(document)
+            documents_by_object_key[document.object_keys[0]].append(document)
 
         objects_to_rebuild = {}