diff --git a/core/common/models.py b/core/common/models.py index 527f76ca..de6aa7dd 100644 --- a/core/common/models.py +++ b/core/common/models.py @@ -505,12 +505,13 @@ def index_batch(ids): batch_size = INDEX_BATCH_SIZE start = 0 id_qs = queryset.order_by('-id').values_list('id', flat=True) - while True: - batch = list(id_qs[start:start + batch_size]) - if not batch: - break + # Keyset paging, so each page costs the same; a union can't be filtered, so it pages by offset + keyset = not id_qs.query.combinator + batch = list(id_qs[start:start + batch_size]) + while batch: run.attempt(start, batch, index_batch) start += batch_size + batch = list(id_qs.filter(id__lt=batch[-1])[:batch_size] if keyset else id_qs[start:start + batch_size]) if refresh: run.refresh(doc) diff --git a/core/common/tasks.py b/core/common/tasks.py index aacccd26..d3a02976 100644 --- a/core/common/tasks.py +++ b/core/common/tasks.py @@ -878,12 +878,14 @@ def index_in_parts(*parts): def exclude_source_members(queryset, source_id): """ - Leaves out the queryset's rows that are members of source_id. An explicit id__in: Django's exclude(sources=...) - correlates on the join table's own row id rather than the resource id, so it leaves the wrong rows out. + Leaves out the queryset's rows that are members of source_id. Correlated explicitly on the resource id: + Django's exclude(sources=...) correlates on the join table's own row id, so it leaves the wrong rows out. + NOT EXISTS, not NOT IN: Postgres can't turn NOT IN into an anti-join, and once the member list outgrows + work_mem it rescans that list for every row (hours on a large release). """ through = queryset.model.sources.through member_field = queryset.model._meta.model_name + '_id' # pylint: disable=protected-access - return queryset.exclude(id__in=through.objects.filter(source_id=source_id).values(member_field)) + return queryset.filter(~Exists(through.objects.filter(source_id=source_id, **{member_field: OuterRef('id')}))) def get_concepts_to_index(source, locales=None, exclude_locale=None, exclude_members_of=None): diff --git a/core/common/tests.py b/core/common/tests.py index deb2f9ad..f9ceb121 100644 --- a/core/common/tests.py +++ b/core/common/tests.py @@ -1528,7 +1528,9 @@ def test_batch_index_full_refresh_error_does_not_fail_the_run(self, _): @patch('core.common.models.parallel_bulk', return_value=[]) def test_batch_index_partial_by_ids_refreshes_once_after_the_last_batch(self, parallel_bulk_mock): id_queryset = MagicMock() - id_queryset.__getitem__.side_effect = lambda batch_slice: [1, 2] if batch_slice.start == 0 else [] + id_queryset.query.combinator = None + id_queryset.__getitem__.return_value = [2, 1] + id_queryset.filter.return_value.__getitem__.return_value = [] queryset = Mock() queryset.order_by.return_value.values_list.return_value = id_queryset doc_instance = Mock() @@ -1540,6 +1542,34 @@ def test_batch_index_partial_by_ids_refreshes_once_after_the_last_batch(self, pa self.assertNotIn('refresh', parallel_bulk_mock.call_args[1]) doc_instance._index.refresh.assert_called_once_with() # pylint: disable=protected-access + id_queryset.filter.assert_called_once_with(id__lt=1) + + def _paged_ids(self, queryset): + pages = [] + + def get_actions(ids): + pages.append(list(ids)) + return iter([]) + + with override_settings(TEST_MODE=False), patch('core.common.models.INDEX_BATCH_SIZE', 2), \ + patch('core.common.models.BatchIndexRun.bulk'): + BaseModel.batch_index_partial_by_ids(queryset, ConceptDocument, get_actions, refresh=False) + return pages + + def test_batch_index_partial_by_ids_pages_by_keyset(self): + source = OrganizationSourceFactory() + concepts = [ConceptFactory(parent=source) for _ in range(5)] + ids = sorted((concept.id for concept in concepts), reverse=True) + + self.assertEqual(self._paged_ids(Concept.objects.filter(id__in=ids)), [ids[0:2], ids[2:4], ids[4:]]) + + def test_batch_index_partial_by_ids_pages_a_union_by_offset(self): + source = OrganizationSourceFactory() + concepts = [ConceptFactory(parent=source) for _ in range(5)] + ids = sorted((concept.id for concept in concepts), reverse=True) + queryset = Concept.objects.filter(id__in=ids[:3]).union(Concept.objects.filter(id__in=ids[3:])) + + self.assertEqual(self._paged_ids(queryset), [ids[0:2], ids[2:4], ids[4:]]) @override_settings(TEST_MODE=False) def test_batch_index_full_continues_remaining_batches_after_error(self):