Add a parent filter on ListCorpusEntities

arkindex/documents/api/entities.py

@@ -336,6 +336,15 @@ class ElementEntities(RetrieveAPIView):
     get=extend_schema(
         operation_id='ListCorpusEntities',
         tags=['entities'],
+        parameters=[
+            OpenApiParameter(
+                'parent',
+                type=UUID,
+                description='Restrict entities to those linked to all transcriptions of an element '
+                            'and all its descendants. Note that links to metadata are ignored.',
+                required=False,
+            )
+        ]
     )
 )
 class CorpusEntities(CorpusACLMixin, ListAPIView):
@@ -345,7 +354,29 @@ class CorpusEntities(CorpusACLMixin, ListAPIView):
     serializer_class = BaseEntitySerializer
 
     def get_queryset(self):
-        return self.get_corpus(self.kwargs['pk']).entities.order_by('name', 'id')
+        corpus = self.get_corpus(self.kwargs['pk'])
+        queryset = corpus.entities.order_by('name', 'id')
+
+        if self.request.query_params.get('parent'):
+            try:
+                parent_id = UUID(self.request.query_params.get('parent'))
+            except (TypeError, ValueError):
+                raise serializers.ValidationError({'parent': ['Parent ID is not a valid UUID']})
+
+            parent_query = corpus.elements.filter(id=parent_id)
+            if not parent_query.exists():
+                raise NotFound
+
+            # We want to include entities that are either related to the selected element, or to its descendants
+            # Element.objects.get_descending() does not include the parent element, so we add it manually using a union.
+            # We remove the ordering from get_descending as it would add an extra column to the query and cause
+            # the union to fail.
+            # Filtering using Q(element_id=…) | Q(element_id__in=…) is 100 times slower than filtering using a union.
+
+            element_ids = parent_query.values('id').union(Element.objects.get_descending(parent_id).order_by())
+            queryset = queryset.filter(transcriptions__element_id__in=element_ids)
+
+        return queryset
 
 
 @extend_schema_view(get=extend_schema(operation_id='ListElementLinks', tags=['entities']))

arkindex/documents/tests/test_entities_api.py

@@ -722,6 +722,59 @@ class TestEntitiesAPI(FixtureAPITestCase):
             response = self.client.get(reverse('api:corpus-entities', kwargs={'pk': str(self.private_corpus.id)}))
             self.assertEqual(response.status_code, status.HTTP_403_FORBIDDEN)
 
+    def test_list_corpus_entities_parent(self):
+        with self.assertNumQueries(4):
+            response = self.client.get(
+                reverse('api:corpus-entities', kwargs={'pk': str(self.corpus.id)}),
+                {'parent': str(self.element.id)},
+            )
+            self.assertEqual(response.status_code, status.HTTP_200_OK)
+        self.assertDictEqual(response.json(), {
+            'count': 1,
+            'previous': None,
+            'next': None,
+            'number': 1,
+            'results': [
+                {
+                    'id': str(self.entity_bis.id),
+                    'name': 'entity 2',
+                    'type': 'location',
+                    'validated': False,
+                    'dates': [],
+                    'metas': None,
+                    'worker_version_id': str(self.worker_version_2.id),
+                },
+            ]
+        })
+
+    def test_list_corpus_entities_invalid_parent_uuid(self):
+        TranscriptionEntity.objects.create(
+            transcription=self.transcription,
+            entity=self.entity,
+            offset=4,
+            length=len(self.entity.name)
+        )
+        with self.assertNumQueries(1):
+            response = self.client.get(reverse('api:corpus-entities', kwargs={'pk': str(self.corpus.id)}), {'parent': 'a'})
+            self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
+        self.assertDictEqual(response.json(), {
+            'parent': ['Parent ID is not a valid UUID']
+        })
+
+    def test_list_corpus_entities_parent_not_found(self):
+        TranscriptionEntity.objects.create(
+            transcription=self.transcription,
+            entity=self.entity,
+            offset=4,
+            length=len(self.entity.name)
+        )
+        with self.assertNumQueries(2):
+            response = self.client.get(
+                reverse('api:corpus-entities', kwargs={'pk': str(self.corpus.id)}),
+                {'parent': 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'},
+            )
+            self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
+
     def test_list_element_entities(self):
         """
         List an element's entities sorted depending on their origin (metadata or transcriptions)