import logging import multiprocessing import os import time from datetime import timedelta from django.core.management.base import BaseCommand, CommandError from django.db import close_old_connections, reset_queries from django.utils.encoding import force_str, smart_bytes from django.utils.timezone import now from haystack import connections as haystack_connections from haystack.exceptions import NotHandled from haystack.query import SearchQuerySet from haystack.utils.app_loading import haystack_get_models, haystack_load_apps DEFAULT_BATCH_SIZE = None DEFAULT_AGE = None DEFAULT_MAX_RETRIES = 5 LOG = multiprocessing.log_to_stderr(level=logging.WARNING) def update_worker(args): if len(args) != 10: LOG.error("update_worker received incorrect arguments: %r", args) raise ValueError("update_worker received incorrect arguments") ( model, start, end, total, using, start_date, end_date, verbosity, commit, max_retries, ) = args # FIXME: confirm that this is still relevant with modern versions of Django: # We need to reset the connections, otherwise the different processes # will try to share the connection, which causes things to blow up. from django.db import connections for alias, info in connections.databases.items(): # We need to also tread lightly with SQLite, because blindly wiping # out connections (via ``... = {}``) destroys in-memory DBs. if "sqlite3" not in info["ENGINE"]: try: close_old_connections() if isinstance(connections._connections, dict): del connections._connections[alias] else: delattr(connections._connections, alias) except KeyError: pass # Request that the connection clear out any transient sessions, file handles, etc. haystack_connections[using].reset_sessions() unified_index = haystack_connections[using].get_unified_index() index = unified_index.get_index(model) backend = haystack_connections[using].get_backend() qs = index.build_queryset(using=using, start_date=start_date, end_date=end_date) do_update(backend, index, qs, start, end, total, verbosity, commit, max_retries) return args def do_update( backend, index, qs, start, end, total, verbosity=1, commit=True, max_retries=DEFAULT_MAX_RETRIES, last_max_pk=None, ): # Get a clone of the QuerySet so that the cache doesn't bloat up # in memory. Useful when reindexing large amounts of data. # the query must be ordered by PK in order to get the max PK in each batch small_cache_qs = qs.all().order_by("pk") # If we got the max seen PK from last batch, use it to restrict the qs # to values above; this optimises the query for Postgres as not to # devolve into multi-second run time at large offsets. if last_max_pk is not None: current_qs = small_cache_qs.filter(pk__gt=last_max_pk)[: end - start] else: current_qs = small_cache_qs[start:end] # Remember maximum PK seen so far max_pk = None current_qs = list(current_qs) if current_qs: max_pk = current_qs[-1].pk is_parent_process = hasattr(os, "getppid") and os.getpid() == os.getppid() if verbosity >= 2: if is_parent_process: print(" indexed %s - %d of %d." % (start + 1, end, total)) else: print( " indexed %s - %d of %d (worker PID: %s)." % (start + 1, end, total, os.getpid()) ) retries = 0 while retries < max_retries: try: # FIXME: Get the right backend. backend.update(index, current_qs, commit=commit) if verbosity >= 2 and retries: print( "Completed indexing {} - {}, tried {}/{} times".format( start + 1, end, retries + 1, max_retries ) ) break except Exception as exc: # Catch all exceptions which do not normally trigger a system exit, excluding SystemExit and # KeyboardInterrupt. This avoids needing to import the backend-specific exception subclasses # from pysolr, elasticsearch, whoosh, requests, etc. retries += 1 error_context = { "start": start + 1, "end": end, "retries": retries, "max_retries": max_retries, "pid": os.getpid(), "exc": exc, } error_msg = "Failed indexing %(start)s - %(end)s (retry %(retries)s/%(max_retries)s): %(exc)s" if not is_parent_process: error_msg += " (pid %(pid)s): %(exc)s" if retries >= max_retries: LOG.exception(error_msg, error_context) raise elif verbosity >= 2: LOG.warning(error_msg, error_context, exc_info=True) # If going to try again, sleep a bit before time.sleep(2**retries) # Clear out the DB connections queries because it bloats up RAM. reset_queries() return max_pk class Command(BaseCommand): help = "Freshens the index for the given app(s)." # noqa A003 def add_arguments(self, parser): parser.add_argument( "app_label", nargs="*", help="App label of an application to update the search index.", ) parser.add_argument( "-m", "--minutes", type=int, help="Number of minutes back to consider objects new.", ) parser.add_argument( "-a", "--age", type=int, default=DEFAULT_AGE, help="Number of hours back to consider objects new.", ) parser.add_argument( "-s", "--start", dest="start_date", help="The start date for indexing. Can be any dateutil-parsable string;" " YYYY-MM-DDTHH:MM:SS is recommended to avoid confusion", ) parser.add_argument( "-e", "--end", dest="end_date", help="The end date for indexing. Can be any dateutil-parsable string;" " YYYY-MM-DDTHH:MM:SS is recommended to avoid confusion", ) parser.add_argument( "-b", "--batch-size", dest="batchsize", type=int, help="Number of items to index at once.", ) parser.add_argument( "-r", "--remove", action="store_true", default=False, help="Remove objects from the index that are no longer present in the database.", ) parser.add_argument( "-u", "--using", action="append", default=[], help="Update only the named backend (can be used multiple times). " "By default all backends will be updated.", ) parser.add_argument( "-k", "--workers", type=int, default=0, help="Allows for the use multiple workers to parallelize indexing.", ) parser.add_argument( "--nocommit", action="store_false", dest="commit", default=True, help="Will pass commit=False to the backend.", ) parser.add_argument( "-t", "--max-retries", action="store", dest="max_retries", type=int, default=DEFAULT_MAX_RETRIES, help="Maximum number of attempts to write to the backend when an error occurs.", ) def handle(self, **options): self.verbosity = int(options.get("verbosity", 1)) self.batchsize = options.get("batchsize", DEFAULT_BATCH_SIZE) self.start_date = None self.end_date = None self.remove = options.get("remove", False) self.workers = options.get("workers", 0) self.commit = options.get("commit", True) self.max_retries = options.get("max_retries", DEFAULT_MAX_RETRIES) self.backends = options.get("using") if not self.backends: self.backends = haystack_connections.connections_info.keys() age = options.get("age", DEFAULT_AGE) minutes = options.get("minutes", DEFAULT_AGE) start_date = options.get("start_date") end_date = options.get("end_date") if self.verbosity > 2: LOG.setLevel(logging.DEBUG) elif self.verbosity > 1: LOG.setLevel(logging.INFO) if (minutes and age) or (minutes and start_date) or (age and start_date): raise CommandError( "Minutes / age / start date options are mutually exclusive" ) if minutes is not None: self.start_date = now() - timedelta(minutes=minutes) if age is not None: self.start_date = now() - timedelta(hours=age) if start_date is not None: from dateutil.parser import parse as dateutil_parse try: self.start_date = dateutil_parse(start_date) except ValueError: pass if end_date is not None: from dateutil.parser import parse as dateutil_parse try: self.end_date = dateutil_parse(end_date) except ValueError: pass labels = options.get("app_label") or haystack_load_apps() for label in labels: for using in self.backends: try: self.update_backend(label, using) except Exception: LOG.exception("Error updating %s using %s ", label, using) raise def update_backend(self, label, using): backend = haystack_connections[using].get_backend() unified_index = haystack_connections[using].get_unified_index() for model in haystack_get_models(label): try: index = unified_index.get_index(model) except NotHandled: if self.verbosity >= 2: self.stdout.write("Skipping '%s' - no index." % model) continue if self.workers > 0: # workers resetting connections leads to references to models / connections getting # stale and having their connection disconnected from under them. Resetting before # the loop continues and it accesses the ORM makes it better. close_old_connections() qs = index.build_queryset( using=using, start_date=self.start_date, end_date=self.end_date ) total = qs.count() if self.verbosity >= 1: self.stdout.write( "Indexing %d %s" % (total, force_str(model._meta.verbose_name_plural)) ) batch_size = self.batchsize or backend.batch_size if self.workers > 0: ghetto_queue = [] max_pk = None for start in range(0, total, batch_size): end = min(start + batch_size, total) if self.workers == 0: max_pk = do_update( backend, index, qs, start, end, total, verbosity=self.verbosity, commit=self.commit, max_retries=self.max_retries, last_max_pk=max_pk, ) else: ghetto_queue.append( ( model, start, end, total, using, self.start_date, self.end_date, self.verbosity, self.commit, self.max_retries, ) ) if self.workers > 0: pool = multiprocessing.Pool(self.workers) successful_tasks = pool.map(update_worker, ghetto_queue) if len(ghetto_queue) != len(successful_tasks): self.stderr.write( "Queued %d tasks but only %d completed" % (len(ghetto_queue), len(successful_tasks)) ) for i in ghetto_queue: if i not in successful_tasks: self.stderr.write("Incomplete task: %s" % repr(i)) pool.close() pool.join() if self.remove: if self.start_date or self.end_date or total <= 0: # They're using a reduced set, which may not incorporate # all pks. Rebuild the list with everything. qs = index.index_queryset(using=using).values_list("pk", flat=True) database_pks = {smart_bytes(pk) for pk in qs} else: database_pks = { smart_bytes(pk) for pk in qs.values_list("pk", flat=True) } # Since records may still be in the search index but not the local database # we'll use that to create batches for processing. # See https://github.com/django-haystack/django-haystack/issues/1186 index_total = ( SearchQuerySet(using=backend.connection_alias).models(model).count() ) # Retrieve PKs from the index. Note that this cannot be a numeric range query because although # pks are normally numeric they can be non-numeric UUIDs or other custom values. To reduce # load on the search engine, we only retrieve the pk field, which will be checked against the # full list obtained from the database, and the id field, which will be used to delete the # record should it be found to be stale. index_pks = SearchQuerySet(using=backend.connection_alias).models(model) index_pks = index_pks.values_list("pk", "id") # We'll collect all of the record IDs which are no longer present in the database and delete # them after walking the entire index. This uses more memory than the incremental approach but # avoids needing the pagination logic below to account for both commit modes: stale_records = set() for start in range(0, index_total, batch_size): upper_bound = start + batch_size # If the database pk is no longer present, queue the index key for removal: for pk, rec_id in index_pks[start:upper_bound]: if smart_bytes(pk) not in database_pks: stale_records.add(rec_id) if stale_records: if self.verbosity >= 1: self.stdout.write( " removing %d stale records." % len(stale_records) ) for rec_id in stale_records: # Since the PK was not in the database list, we'll delete the record from the search # index: if self.verbosity >= 2: self.stdout.write(" removing %s." % rec_id) backend.remove(rec_id, commit=self.commit)