src/search/management/commands/reindex.py

   1 # This file is part of Wolnelektury, licensed under GNU Affero GPLv3 or later.
   2 # Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
   3 #
   4 import sys
   5 import traceback
   6
   7 from django.core.management.base import BaseCommand
   8
   9
  10 def query_yes_no(question, default="yes"):
  11     """Ask a yes/no question via raw_input() and return their answer.
  12
  13     "question" is a string that is presented to the user.
  14     "default" is the presumed answer if the user just hits <Enter>.
  15         It must be "yes" (the default), "no" or None (meaning
  16         an answer is required of the user).
  17
  18     The "answer" return value is one of "yes" or "no".
  19     """
  20     valid = {"yes": True, "y": True, "ye": True,
  21              "no": False, "n": False}
  22     if default is None:
  23         prompt = " [y/n] "
  24     elif default == "yes":
  25         prompt = " [Y/n] "
  26     elif default == "no":
  27         prompt = " [y/N] "
  28     else:
  29         raise ValueError("invalid default answer: '%s'" % default)
  30
  31     while True:
  32         sys.stdout.write(question + prompt)
  33         choice = raw_input().lower()
  34         if default is not None and choice == '':
  35             return valid[default]
  36         elif choice in valid:
  37             return valid[choice]
  38         else:
  39             sys.stdout.write("Please respond with 'yes' or 'no' (or 'y' or 'n').\n")
  40
  41
  42 class Command(BaseCommand):
  43     help = 'Reindex everything.'
  44
  45     def add_arguments(self, parser):
  46         parser.add_argument(
  47                 '-n', '--book-id', action='store_true', dest='book_id',
  48                 default=False, help='book id instead of slugs')
  49         parser.add_argument(
  50                 '-t', '--just-tags', action='store_true', dest='just_tags',
  51                 default=False, help='just reindex tags')
  52         parser.add_argument(
  53                 '--start', dest='start_from', default=None,
  54                 help='start from this slug')
  55         parser.add_argument(
  56                 '--stop', dest='stop_after', default=None,
  57                 help='stop after this slug')
  58         parser.add_argument('args', nargs='*', metavar='slug/id')
  59
  60     def handle(self, **opts):
  61         from catalogue.models import Book
  62         from search.index import Index
  63         idx = Index()
  64
  65         if not opts['just_tags']:
  66             if opts['args']:
  67                 books = []
  68                 for a in opts['args']:
  69                     if opts['book_id']:
  70                         books += Book.objects.filter(id=int(a)).all()
  71                     else:
  72                         books += Book.objects.filter(slug=a).all()
  73             else:
  74                 books = list(Book.objects.order_by('slug'))
  75             start_from = opts.get('start_from')
  76             stop_after = opts.get('stop_after')
  77             if start_from:
  78                 start_from = start_from.replace('-', '')
  79             if stop_after:
  80                 stop_after = stop_after.replace('-', '')
  81             while books:
  82                 try:
  83                     b = books[0]
  84                     slug = b.slug.replace('-', '')
  85                     if stop_after and slug > stop_after:
  86                         break
  87                     if not start_from or slug >= start_from:
  88                         print(b.slug)
  89                         idx.index_book(b)
  90                         idx.index.commit()
  91                     books.pop(0)
  92                 except:
  93                     traceback.print_exc()
  94                     try:
  95                         # we might not be able to rollback
  96                         idx.index.rollback()
  97                     except:
  98                         pass
  99                     retry = query_yes_no("Retry?")
 100                     if not retry:
 101                         break
 102
 103         print('Reindexing tags.')
 104         idx.index_tags()
 105         idx.index.commit()