Issues with `delete_stale_model_entries`

[jsma]

I've encountered a couple issues with delete_stale_model_entries.

django-modelsearch/modelsearch/backends/database/postgres/postgres.py

Lines 222 to 230 in 53d1292

	def delete_stale_model_entries(self, model):
	existing_pks = model._default_manager.annotate(
	object_id=Cast("pk", TextField())
	).values("object_id")
	content_types_pks = get_descendants_content_types_pks(model)
	stale_entries = self.entries.filter(
	content_type_id__in=content_types_pks
	).exclude(object_id__in=existing_pks)
	stale_entries.delete()

First, it doesn't consult get_indexed_objects when determining the "existing" objects that belong in the index.
This leads to an issue where if get_indexed_objects is later modified to reduce (filter/exclude) the number of indexed objects, delete_stale_model_entries won't actually remove them from the index since it fetches all PKs. It only prunes records that have been deleted from the database.

Second, it passes a potentially huge list of object IDs into a QuerySet filter which may result in the database taking a significant amount of time to run the query (I've seen it run for longer than 40 minutes on an AWS RDS m5.large instance before manually cancelling the index rebuild).

For example, let's assume a standard User model and that there are 250k+ users in the database:

class User(index.Indexed, AbstractUser):
    search_fields = [
        index.SearchField("first_name"),
        index.SearchField("last_name"),
        index.SearchField("email"),
    ]

When the index is built, there will be 250k+ users indexed.
If it is then modified to only index staff (e.g. 10 users), it will not prune the existing records. The database will also cry doing an IN query on 250k+ PKs:

class User(index.Indexed, AbstractUser):
    search_fields = [
        index.SearchField("first_name"),
        index.SearchField("last_name"),
        index.SearchField("email"),
    ]
    @classmethod
    def get_indexed_objects(cls):
        return super().get_indexed_objects().filter(is_staff=True)

Below is my initial attempt at fixing these issues but it is admittedly incomplete and inadequate.

I'm using this in a Wagtail site, and we only use modelsearch to power custom choosers in the Wagtail admin (we have a separate custom integration with Elasticsearch for end-user facing search).

What causes a bit of a wrinkle in my approach is that in the specific case of Wagtail, Wagtail must index all pages so that drafts etc. can be searched in the Wagtail admin. get_indexed_objects cannot be consulted since it may make some pages impossible to find in the Wagtail admin.

I also do not happen to have any other multi-table inheritance models in my project to test, so for now I just throw an exception to think through this bit at a later time. My guess is that this would need to walk through each child model (e.g., Dog and Cat) to consult get_indexed_objects to get the correct list of indexable PKs for each model, and then only include "pure" parent model PKs (e.g., Pet).

class CustomIndex(PostgresIndex):
    def delete_stale_model_entries(self, model):
        content_types_pks = get_descendants_content_types_pks(model)

        # self.delete_stale_entries() filters out `if not model._meta.parents` which
        # means all subclasses will be excluded and only the base Page class will appear
        # here. This in turn means we have to fetch all Page pks since it will be
        # filtering based on the content_type_pks of all Page subclasses.
        # Page.get_indexed_objects() can't be used because it only returns pages that
        # are pure Page instances (not a page created through a Page subclass), which in
        # practice means it would only return the root page.
        if model == Page:
            indexable_qs = model._default_manager.all()
        # Currently, there aren't any other indexed models that will have any subclass
        # content types, but scream if we do.
        elif len(content_types_pks) > 1:
            raise Exception(
                f"{model._meta.label} has descendents. "
                f"Please review to determine how to best `delete_stale_model_entries`."
            )
        # For all other "normal" models, any PK that has previously been indexed but no longer
        # appears in `get_indexed_objects()` should be pruned.
        else:
            indexable_qs = model.get_indexed_objects()

        # Rather than attempt an extremely inefficient `exclude(object_id__in=existing_pks)`,
        # find just the PKs that have been indexed that are no longer indexable.
        indexable_pks = set(
            indexable_qs.annotate(object_id=Cast("pk", TextField())).values_list(
                "object_id", flat=True
            )
        )
        indexed_pks = set(
            self.entries.filter(content_type_id__in=content_types_pks).values_list(
                "object_id", flat=True
            )
        )
        stale_pks = list(indexed_pks - indexable_pks)
        if stale_pks:
            # Delete the stale items in chunks to avoid transmitting a large number of
            # pks in the `__in` filter
            chunk_size = 100
            i = 0
            while True:
                chunked_pks = stale_pks[i * chunk_size :][:chunk_size]
                if not chunked_pks:
                    break
                stale_entries = self.entries.filter(
                    content_type_id__in=content_types_pks).filter(object_id__in=chunked_pks)
                stale_entries.delete()
                i += 1


class SearchBackend(PostgresSearchBackend):
    def get_index_for_model(self, model):
        return CustomIndex(self)

FYI, I've noticed there is quite a bit of duplication across all of the database index sub-classes. E.g., delete_stale_model_entries is exactly the same across the Postgres, MySQL, and SQLite sub-classes. Any appetite for creating a BaseDatabaseIndex class that puts all of the common functionality in a base class so the sub-classes only contain database-specific logic?