apps/search/management/commands/reindex.py

   1 import sys
   2 from django.core.management.base import BaseCommand
   3
   4 from optparse import make_option
   5
   6 def query_yes_no(question, default="yes"):
   7     """Ask a yes/no question via raw_input() and return their answer.
   8
   9     "question" is a string that is presented to the user.
  10     "default" is the presumed answer if the user just hits <Enter>.
  11         It must be "yes" (the default), "no" or None (meaning
  12         an answer is required of the user).
  13
  14     The "answer" return value is one of "yes" or "no".
  15     """
  16     valid = {"yes":True,   "y":True,  "ye":True,
  17              "no":False,     "n":False}
  18     if default == None:
  19         prompt = " [y/n] "
  20     elif default == "yes":
  21         prompt = " [Y/n] "
  22     elif default == "no":
  23         prompt = " [y/N] "
  24     else:
  25         raise ValueError("invalid default answer: '%s'" % default)
  26
  27     while True:
  28         sys.stdout.write(question + prompt)
  29         choice = raw_input().lower()
  30         if default is not None and choice == '':
  31             return valid[default]
  32         elif choice in valid:
  33             return valid[choice]
  34         else:
  35             sys.stdout.write("Please respond with 'yes' or 'no' "\
  36                              "(or 'y' or 'n').\n")
  37
  38 class Command(BaseCommand):
  39     help = 'Reindex everything.'
  40     args = ''
  41
  42     option_list = BaseCommand.option_list + (
  43         make_option('-n', '--book-id', action='store_true', dest='book_id', default=False,
  44             help='book id instead of slugs'),
  45         make_option('-t', '--just-tags', action='store_true', dest='just_tags', default=False,
  46             help='just reindex tags'),
  47     )
  48     def handle(self, *args, **opts):
  49         from catalogue.models import Book
  50         import search
  51         idx = search.Index()
  52
  53         if not opts['just_tags']:
  54             if args:
  55                 books = []
  56                 for a in args:
  57                     if opts['book_id']:
  58                         books += Book.objects.filter(id=int(a)).all()
  59                     else:
  60                         books += Book.objects.filter(slug=a).all()
  61             else:
  62                 books = list(Book.objects.all())
  63
  64             while books:
  65                 try:
  66                     b = books[0]
  67                     print b.title
  68                     idx.index_book(b)
  69                     idx.index.commit()
  70                     books.pop(0)
  71                 except Exception, e:
  72                     print "Error occured: %s" % e
  73                     try:
  74                         # we might not be able to rollback
  75                         idx.index.rollback()
  76                     except:
  77                         pass
  78                     retry = query_yes_no("Retry?")
  79                     if not retry:
  80                         break
  81
  82         print 'Reindexing tags.'
  83         idx.index_tags()
  84         idx.index.commit()