src/search/management/commands/reindex.py

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