Machina preview.
[wolnelektury.git] / 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()