codex 1.6.19__py3-none-any.whl → 1.7.0a0__py3-none-any.whl

codex/librarian/search/update.py

@@ -1,321 +1,258 @@
 """Search Index update."""
 
 from datetime import datetime
-from math import ceil
-from multiprocessing import Pool, cpu_count
 from time import time
-from typing import TYPE_CHECKING
 from zoneinfo import ZoneInfo
 
-from django.core.exceptions import ObjectDoesNotExist
-from django.db import DatabaseError
+from django.db.models.aggregates import Max
+from django.db.models.expressions import F
+from django.db.models.functions.datetime import Now
 from humanize import naturaldelta
 
-from codex._vendor.haystack.exceptions import SearchFieldError
 from codex.librarian.search.remove import RemoveMixin
 from codex.librarian.search.status import SearchIndexStatusTypes
-from codex.librarian.search.tasks import SearchIndexRemoveStaleTask
-from codex.memory import get_mem_limit
 from codex.models import Comic, Library
-from codex.settings.settings import (
-    CHUNK_PER_GB,
-    CPU_MULTIPLIER,
-    MAX_CHUNK_SIZE,
-    MMAP_RATIO,
-    WRITER_MEMORY_PERCENT,
-)
+from codex.models.comic import ComicFTS
+from codex.models.functions import GroupConcat
+from codex.serializers.fields import CountryField, LanguageField, PyCountryField
+from codex.settings.settings import SEARCH_INDEX_BATCH_SIZE
 from codex.status import Status
 
-if TYPE_CHECKING:
-    from codex.search.backend import CodexSearchBackend
+_COMICFTS_UPDATE_FIELDS = (
+    # ForeignKeys
+    "publisher",
+    "imprint",
+    "series",
+    "volume",
+    # Attributes
+    "age_rating",
+    "file_type",
+    "country",
+    "issue",
+    "language",
+    "name",
+    "notes",
+    "original_format",
+    "reading_direction",
+    "review",
+    "scan_info",
+    "summary",
+    "tagger",
+    # ManyToMany
+    "characters",
+    "contributors",
+    "genres",
+    "locations",
+    "series_groups",
+    "stories",
+    "story_arcs",
+    "tags",
+    "teams",
+    # Base
+    "updated_at",
+)
+_MIN_UTC_DATE = datetime.min.replace(tzinfo=ZoneInfo("UTC"))
 
 
-class UpdateMixin(RemoveMixin):
+class FTSUpdateMixin(RemoveMixin):
     """Search Index update methods."""
 
     _STATUS_FINISH_TYPES = (
         SearchIndexStatusTypes.SEARCH_INDEX_CLEAR,
+        SearchIndexStatusTypes.SEARCH_INDEX_CREATE,
         SearchIndexStatusTypes.SEARCH_INDEX_UPDATE,
     )
-    _MIN_UTC_DATE = datetime.min.replace(tzinfo=ZoneInfo("UTC"))
-    _MIN_BATCH_SIZE = 1
-    # A larger batch size might be slightly faster for very large
-    # indexes and require less optimization later, but steady progress
-    # updates are better UX.
-    _MAX_BATCH_SIZE = 640
-    _EXPECTED_EXCEPTIONS = (
-        DatabaseError,
-        IndexError,
-        ObjectDoesNotExist,
-        SearchFieldError,
-    )
-    _MAX_RETRIES = 8
-    _CPU_MULTIPLIER = CPU_MULTIPLIER
-    _UPDATE_SORT_BY = ("-updated_at",)
 
     def _init_statuses(self, rebuild):
         """Initialize all statuses order before starting."""
         statii = []
         if rebuild:
             statii += [Status(SearchIndexStatusTypes.SEARCH_INDEX_CLEAR)]
-        statii += [Status(SearchIndexStatusTypes.SEARCH_INDEX_UPDATE)]
-        if not rebuild:
+        else:
             statii += [Status(SearchIndexStatusTypes.SEARCH_INDEX_REMOVE)]
+        statii += [
+            Status(SearchIndexStatusTypes.SEARCH_INDEX_UPDATE),
+            Status(SearchIndexStatusTypes.SEARCH_INDEX_CREATE),
+        ]
         self.status_controller.start_many(statii)
 
-    def _get_queryset_from_search_index(self, backend, qs, suffix):
-        """Get the date of the last updated item in the search index."""
-        search_results = backend.search("*", sort_by=self._UPDATE_SORT_BY)
-        results = search_results.get("results")
-        if results:
-            most_recent_result = results[0]
-            most_recent_updated_at = most_recent_result.updated_at
-            qs = qs.filter(updated_at__gt=most_recent_updated_at)
-            suffix = f"since {most_recent_updated_at}"
-
-        return qs, suffix
-
-    def _get_queryset(self, backend, rebuild):
-        """Rebuild or set up update."""
-        qs = Comic.objects.all()
-
-        suffix = "with all comics"
-        if not rebuild:
-            qs, suffix = self._get_queryset_from_search_index(backend, qs, suffix)
-
-        self.log.info(f"Updating search index {suffix}...")
-        return qs
-
-    def _get_throttle_params(self, num_comics):
-        """Get params based on memory and total number of comics.
-
-        >4GB is normal.
-        <=2GB is constrained.
-        """
-        mem_limit_gb = get_mem_limit("g")
-
-        # max procs
-        # throttle multiprocessing in lomem environments.
-        # each process running has significant memory overhead.
-        cpu_max = ceil(mem_limit_gb * self._CPU_MULTIPLIER)
-        max_procs = min(cpu_count(), cpu_max)
-
-        batch_size = int(
-            max(
-                self._MIN_BATCH_SIZE,
-                min(num_comics / max_procs, self._MAX_BATCH_SIZE),
-            )
-        )
-
-        num_procs = int(min(max(1, num_comics / batch_size), max_procs))
-
-        opts = {
-            "comics": num_comics,
-            "memgb": mem_limit_gb,
-            "procs": num_procs,
-            "batch_size": batch_size,
-        }
-        self.log.debug(
-            f"{MMAP_RATIO=}, {WRITER_MEMORY_PERCENT=}, {CPU_MULTIPLIER=},"
-            f" {CHUNK_PER_GB=}, {MAX_CHUNK_SIZE=}"
-        )
-        self.log.debug(f"Search Index update opts: {opts}")
-        return num_procs, batch_size
-
-    def _halve_batches(self, batches):
-        """Half the size of retried batches.
-
-        Helps if it was a memory issue.
-        If it's a data issue then it binary searches down to the real problem.
-        """
-        old_num_batches = len(batches)
-        half_batches = {}
-        for batch_num, batch_pks in batches.items():
-            if not batch_pks:
-                continue
-            batch_num_a = batch_num * 2
-            if len(batch_pks) == 1:
-                half_batches[batch_num_a] = batch_pks
-                continue
-            halfway = int(len(batch_pks) / 2)
-            half_batches[batch_num_a] = batch_pks[:halfway]
-            batch_num_b = batch_num_a + 1
-            half_batches[batch_num_b] = batch_pks[halfway:]
-        batches = half_batches
-        self.log.debug(f"Split {old_num_batches} batches into {len(batches)}.")
-        return batches
-
-    def _apply_batches(self, pool, batches, backend):
-        """Apply the batches to the process pool."""
-        results = {}
-        for batch_num, batch_pks in batches.items():
-            args = (None, batch_pks)
-            kwd = {"batch_num": batch_num, "abort_event": self.abort_event}
-            result = pool.apply_async(backend.update, args, kwd)
-            results[batch_num] = (result, batch_pks)
-        pool.close()
-        self.log.debug(f"Search index update queued {len(results)} batches...")
-        return results
-
-    def _collect_results(self, results, complete, status):
-        """Collect results from the process pool."""
-        retry_batches = {}
-        retry_batch_num = 0
-        if self.abort_event.is_set():
-            return retry_batches, retry_batch_num
-
-        num_results = len(results)
-        for batch_num, (result, batch_pks) in results.items():
-            try:
-                complete += result.get()
-                self.log.debug(
-                    f"Search index batch {batch_num}/{num_results} complete: "
-                    f"{status.complete}/{status.total} comics"
-                )
-                status.complete = complete
-                self.status_controller.update(status)
-            except self._EXPECTED_EXCEPTIONS:
-                pass
-            except Exception:
-                reason = f"Search Index Update collect result batch {batch_num}"
-                self.log.exception(reason)
-            else:
-                # success
-                continue
-
-            # failure
-            self.log.debug(f"Search index update will retry batch {batch_num}")
-            retry_batches[retry_batch_num] = batch_pks
-            retry_batch_num += 1
-
-        return retry_batches, complete
-
-    def _try_update_batch(self, backend, batch_info, status):
-        """Attempt to update batches, with reursive retry."""
-        (
-            batches,
-            num_procs,
-            num_comics,
-            batch_size,
-            complete,
-            attempt,
-        ) = batch_info
-        if self.abort_event.is_set():
-            return complete
-
-        if attempt > 1:
-            batches = self._halve_batches(batches)
-
-        num_batches = len(batches)
-        if not num_batches:
-            self.log.debug("Search index nothing to update.")
-            return complete
-        self.log.debug(
-            f"Search index updating {num_batches} batches, attempt {attempt}"
+    def _update_search_index_clean(self, rebuild):
+        """Clear or clean the search index."""
+        if rebuild:
+            self.log.info("Rebuilding search index...")
+            self.clear_search_index()
+        else:
+            self.remove_stale_records()
+
+    @classmethod
+    def _annotate_fts_query(cls, qs):
+        return qs.annotate(
+            fts_publisher=F("publisher__name"),
+            fts_imprint=F("imprint__name"),
+            fts_series=F("series__name"),
+            fts_volume=F("volume__name"),
+            fts_country=F("country__name"),
+            fts_language=F("language__name"),
+            fts_scan_info=F("scan_info__name"),
+            fts_tagger=F("tagger__name"),
+            fts_characters=GroupConcat("characters__name", distinct=True),
+            fts_contributors=GroupConcat("contributors__person__name", distinct=True),
+            fts_genres=GroupConcat("genres__name", distinct=True),
+            fts_locations=GroupConcat("locations__name", distinct=True),
+            fts_series_groups=GroupConcat("series_groups__name", distinct=True),
+            fts_stories=GroupConcat("stories__name", distinct=True),
+            fts_story_arcs=GroupConcat(
+                "story_arc_numbers__story_arc__name", distinct=True
+            ),
+            fts_tags=GroupConcat("tags__name", distinct=True),
+            fts_teams=GroupConcat("teams__name", distinct=True),
         )
-        procs = min(num_procs, num_batches)
-        pool = Pool(procs, maxtasksperchild=1)
-        results = self._apply_batches(pool, batches, backend)
-        retry_batches, complete = self._collect_results(results, complete, status)
-        pool.join()
 
-        num_successful_batches = num_batches - len(retry_batches)
-
-        ratio = 100 * (num_successful_batches / num_batches)
-        self.log.debug(
-            f"Search Index attempt {attempt} batch success ratio: {round(ratio)}%"
+    @staticmethod
+    def _get_pycountry_fts_field(field_instance: PyCountryField, iso_code):
+        if not iso_code:
+            return ""
+        return ",".join((iso_code, field_instance.to_representation(iso_code)))
+
+    def _get_comicts_fts_list(
+        self, comics, country_field, language_field, obj_list, create
+    ):
+        """Prepare a batch of search entries."""
+        for comic in comics:
+            country = self._get_pycountry_fts_field(country_field, comic.fts_country)
+            language = self._get_pycountry_fts_field(language_field, comic.fts_language)
+
+            now = Now()
+            comicfts = ComicFTS(
+                comic_id=comic.pk,
+                updated_at=now,
+                publisher=comic.fts_publisher,
+                imprint=comic.fts_imprint,
+                series=comic.fts_series,
+                volume=comic.fts_volume,
+                issue=comic.issue,
+                name=comic.name,
+                age_rating=comic.age_rating,
+                country=country,
+                file_type=comic.file_type,
+                language=language,
+                notes=comic.notes,
+                original_format=comic.original_format,
+                reading_direction=comic.reading_direction,
+                review=comic.review,
+                scan_info=comic.fts_scan_info,
+                summary=comic.summary,
+                tagger=comic.fts_tagger,
+                # ManyToMany
+                characters=comic.fts_characters,
+                contributors=comic.fts_contributors,
+                genres=comic.fts_genres,
+                locations=comic.fts_locations,
+                series_groups=comic.fts_series_groups,
+                stories=comic.fts_stories,
+                story_arcs=comic.fts_story_arcs,
+                tags=comic.fts_tags,
+                teams=comic.fts_teams,
+            )
+            if create:
+                comicfts.created_at = now
+            obj_list.append(comicfts)
+
+    def _get_comicfts_list(self, comics, create=False):
+        """Create a ComicFTS object for bulk_create or bulk_update."""
+        country_field = CountryField()
+        language_field = LanguageField()
+        obj_list = []
+        comics = self._annotate_fts_query(comics)
+        self._get_comicts_fts_list(
+            comics, country_field, language_field, obj_list, create
         )
-        if not retry_batches:
-            return complete
+        return obj_list
 
-        if attempt < self._MAX_RETRIES:
-            batch_info = (
-                retry_batches,
-                num_procs,
-                num_comics,
-                batch_size,
-                complete,
-                attempt + 1,
-            )
-            complete = self._try_update_batch(backend, batch_info, status)
+    def _update_search_index_finish(self, count, verb, status):
+        verb = verb.capitalize() + "d"
+        if count:
+            self.log.info(f"{verb} {count} search entries.")
         else:
-            total = len(retry_batches) * batch_size
-            self.log.error(
-                f"Search Indexer failed to update {total} comics"
-                f"in {len(retry_batches)} batches."
-            )
-        return complete
+            self.log.debug(f"{verb} no search entries.")
+        self.status_controller.finish(status)
 
-    @staticmethod
-    def _get_update_batch_end_index(start, batch_size, batch_num, procs):
-        if batch_num < procs:
-            # Gets early results to the status update
-            # by making small batches in the beginning.
-            end = start + int(batch_size * (batch_num / procs))
-        else:
-            end = start + batch_size
-        return end
+    def _update_search_index_operate(self, comics_qs, create=False):
+        count = comics_qs.count()
+        verb = "create" if create else "update"
+        if not count:
+            self.log.info(f"No search entries to {verb}.")
 
-    @staticmethod
-    def _get_update_batches(qs, batch_size, num_comics):
-        all_pks = qs.order_by("updated_at", "pk").values_list("pk", flat=True)
-        batch_num = 0
-        start = 0
-        end = start + batch_size
+        if count > SEARCH_INDEX_BATCH_SIZE:
+            self.log.debug(
+                f"Batching this search engine {verb} operation in to chunks of {SEARCH_INDEX_BATCH_SIZE}."
+                f" Search engine {verb}s run much faster as one large batch but then there's no progress updates."
+                " You may adjust the batch size with the environment variable CODEX_SEARCH_INDEX_BATCH_SIZE."
+            )
 
-        batches = {}
-        while start < num_comics:
-            batches[batch_num] = all_pks[start:end]
-            batch_num += 1
-            start = end
-            end = start + batch_size
-        return batches
+        status_type = (
+            SearchIndexStatusTypes.SEARCH_INDEX_CREATE
+            if create
+            else SearchIndexStatusTypes.SEARCH_INDEX_UPDATE
+        )
+        complete = None if count <= SEARCH_INDEX_BATCH_SIZE else 0
+        operate_status = Status(status_type, complete, count)
+        self.status_controller.start(operate_status)
 
-    def _mp_update(self, backend, qs):
-        # Init
-        start_time = time()
-        num_comics = qs.count()
-        if not num_comics or self.abort_event.is_set():
-            return
-        status = Status(SearchIndexStatusTypes.SEARCH_INDEX_UPDATE, 0, num_comics)
+        batch_from = 0
         try:
-            self.status_controller.start(status)
-
-            num_procs, batch_size = self._get_throttle_params(num_comics)
-            batches = self._get_update_batches(qs, batch_size, num_comics)
-            batch_info = (batches, num_procs, num_comics, batch_size, 0, 1)
-            complete = self._try_update_batch(
-                backend,
-                batch_info,
-                status,
-            )
-
-            # Log performance
-            elapsed_time = time() - start_time
-            elapsed = naturaldelta(elapsed_time)
-            cps = int(complete / elapsed_time)
-            self.log.info(
-                f"Search engine updated {complete} comics"
-                f" in {elapsed} at {cps} comics per second."
-            )
-        except Exception:
-            self.log.exception("Update search index with multiprocessing")
+            while batch_from < count:
+                batch_start = time()
+                batch_to = batch_from + SEARCH_INDEX_BATCH_SIZE
+                comics_batch = comics_qs[batch_from:batch_to]
+
+                if self.abort_event.is_set():
+                    break
+                operate_comicfts = self._get_comicfts_list(comics_batch, create)
+                operate_comicfts_count = len(operate_comicfts)
+                if self.abort_event.is_set():
+                    break
+                if create:
+                    ComicFTS.objects.bulk_create(operate_comicfts)
+                else:
+                    ComicFTS.objects.bulk_update(
+                        operate_comicfts, _COMICFTS_UPDATE_FIELDS
+                    )
+                operate_status.add_complete(operate_comicfts_count)
+                self.status_controller.update(operate_status)
+                batch_time = time() - batch_start
+                eps = round(count / batch_time)
+                self.log.debug(
+                    f"{verb} {count}/{operate_comicfts_count} search entries in {batch_time}, {eps} per second."
+                )
+                batch_from = batch_to
         finally:
-            until = time() + 1
-            self.status_controller.finish(status, until=until)
+            self._update_search_index_finish(count, verb, operate_status)
+
+    def _update_search_index_update(self, all_indexed_comic_ids):
+        """Update out of date search entries."""
+        fts_watermark = ComicFTS.objects.aggregate(max=Max("updated_at"))["max"]
+        if not fts_watermark:
+            fts_watermark = _MIN_UTC_DATE
+            since = "the fracturing of the multiverse"
+        else:
+            since = fts_watermark
+        self.log.info(f"Looking for search entries to update since {since}...")
+        out_of_date_comics = Comic.objects.filter(
+            pk__in=all_indexed_comic_ids, updated_at__gt=fts_watermark
+        )
+        self._update_search_index_operate(out_of_date_comics, create=False)
 
-    def clear_search_index(self):
-        """Clear the search index."""
-        clear_status = Status(SearchIndexStatusTypes.SEARCH_INDEX_CLEAR)
-        self.status_controller.start(clear_status)
-        backend: CodexSearchBackend = self.engine.get_backend()  # type: ignore
-        backend.clear(commit=True)
-        self.status_controller.finish(clear_status)
-        self.log.info("Old search index cleared.")
+    def _update_search_index_create(self, all_indexed_comic_ids):
+        """Create missing search entries."""
+        self.log.info("Looking for missing search entries to create...")
+        missing_comics = Comic.objects.all()
+        if len(all_indexed_comic_ids):
+            missing_comics = missing_comics.exclude(pk__in=all_indexed_comic_ids)
+        self._update_search_index_operate(missing_comics, create=True)
 
     def _update_search_index(self, start_time, rebuild):
         """Update or Rebuild the search index."""
-        remove_stale = False
         any_update_in_progress = Library.objects.filter(
             covers_only=False, update_in_progress=True
         ).exists()
@@ -323,49 +260,34 @@ class UpdateMixin(RemoveMixin):
             self.log.debug(
                 "Database update in progress, not updating search index yet."
             )
-            return remove_stale
-
-        if not rebuild and not self.is_search_index_uuid_match():
-            rebuild = True
+            return
 
         self._init_statuses(rebuild)
 
-        # Clear
-        if rebuild:
-            self.log.info("Rebuilding search index...")
-            self.clear_search_index()
-
-        # Update
-        backend: CodexSearchBackend = self.engine.get_backend()  # type: ignore
-        backend.setup(False)
         if self.abort_event.is_set():
-            self.log.debug("Abort update search index.")
-            return remove_stale
-        qs = self._get_queryset(backend, rebuild)
-        self._mp_update(backend, qs)
-
-        # Finish
-        if rebuild:
-            self.set_search_index_version()
-        else:
-            remove_stale = True
+            return
+        self._update_search_index_clean(rebuild)
+        if self.abort_event.is_set():
+            return
+        all_indexed_comic_ids = ComicFTS.objects.values_list("comic_id", flat=True)
+        if self.abort_event.is_set():
+            return
+        self._update_search_index_update(all_indexed_comic_ids)
+        if self.abort_event.is_set():
+            return
+        self._update_search_index_create(all_indexed_comic_ids)
 
         elapsed_time = time() - start_time
         elapsed = naturaldelta(elapsed_time)
         self.log.info(f"Search index updated in {elapsed}.")
-        return remove_stale
 
     def update_search_index(self, rebuild=False):
         """Update or Rebuild the search index."""
         start_time = time()
         self.abort_event.clear()
-        remove_stale = False
         try:
-            remove_stale = self._update_search_index(start_time, rebuild)
+            self._update_search_index(start_time, rebuild)
         except Exception:
             self.log.exception("Update search index")
         finally:
             self.status_controller.finish_many(self._STATUS_FINISH_TYPES)
-            if remove_stale:
-                task = SearchIndexRemoveStaleTask()
-                self.librarian_queue.put(task)

codex/librarian/watchdog/observers.py

@@ -165,8 +165,9 @@ class LibraryPollingObserver(UatuMixin):
             paths = frozenset(qs.values_list("path", flat=True))
 
             for emitter in self.emitters:
+                polling_emitter: DatabasePollingEmitter = emitter  # type: ignore
                 if emitter.watch.path in paths:
-                    emitter.poll(force)  # type: ignore
+                    polling_emitter.poll(force)
         except Exception:
             self.log.exception(
                 f"{self.__class__.__name__}.poll({library_pks}, {force})"

codex/migrations/0029_comicfts.py

@@ -0,0 +1,58 @@
+"""Generated by Django 5.0.8 on 2024-08-07 22:07."""
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    """Run migrations."""
+
+    dependencies = [
+        ("codex", "0028_telemeter"),
+    ]
+
+    operations = [
+        migrations.SeparateDatabaseAndState(
+            state_operations=[
+                migrations.CreateModel(
+                    name="ComicFTS",
+                    fields=[
+                        ("created_at", models.DateTimeField(auto_now_add=True)),
+                        ("updated_at", models.DateTimeField(auto_now=True)),
+                        (
+                            "comic",
+                            models.OneToOneField(
+                                on_delete=django.db.models.deletion.CASCADE,
+                                primary_key=True,
+                                serialize=False,
+                                to="codex.comic",
+                            ),
+                        ),
+                        ("body", models.TextField()),
+                    ],
+                    options={
+                        "get_latest_by": "updated_at",
+                        "managed": False,
+                    },
+                ),
+            ],
+            database_operations=[
+                migrations.RunSQL(
+                    sql=(
+                        "CREATE VIRTUAL TABLE codex_comicfts USING fts5("
+                        "comic_id UNINDEXED, created_at UNINDEXED, "
+                        "updated_at UNINDEXED, "
+                        "publisher, imprint, series, volume, issue, name, age_rating, "
+                        "country, language, "
+                        "notes, original_format, review, scan_info, summary, "
+                        "tagger, "
+                        "characters, contributors, genres,"
+                        "locations, roles, series_groups, stories, "
+                        "story_arcs, tags, teams, "
+                        "reading_direction, file_type)"
+                    ),
+                    reverse_sql="DROP TABLE IF EXISTS codex_comicfts",
+                ),
+            ],
+        )
+    ]