X-Git-Url: https://git.mdrn.pl/wolnelektury.git/blobdiff_plain/73ce961f14509aabfa26536f847afd28111029c6..ccf1868db403f603a87db2194a8ced22a05766f1:/apps/search/management/commands/reindex.py

diff --git a/apps/search/management/commands/reindex.py b/apps/search/management/commands/reindex.py
index bce47080d..4941354ce 100755
--- a/apps/search/management/commands/reindex.py
+++ b/apps/search/management/commands/reindex.py
@@ -1,25 +1,88 @@
+# -*- coding: utf-8 -*-
+# This file is part of Wolnelektury, licensed under GNU Affero GPLv3 or later.
+# Copyright © Fundacja Nowoczesna Polska. See NOTICE for more information.
+#
+import sys
 from django.core.management.base import BaseCommand
 
+from optparse import make_option
+
+def query_yes_no(question, default="yes"):
+    """Ask a yes/no question via raw_input() and return their answer.
+
+    "question" is a string that is presented to the user.
+    "default" is the presumed answer if the user just hits <Enter>.
+        It must be "yes" (the default), "no" or None (meaning
+        an answer is required of the user).
+
+    The "answer" return value is one of "yes" or "no".
+    """
+    valid = {"yes":True,   "y":True,  "ye":True,
+             "no":False,     "n":False}
+    if default == None:
+        prompt = " [y/n] "
+    elif default == "yes":
+        prompt = " [Y/n] "
+    elif default == "no":
+        prompt = " [y/N] "
+    else:
+        raise ValueError("invalid default answer: '%s'" % default)
+
+    while True:
+        sys.stdout.write(question + prompt)
+        choice = raw_input().lower()
+        if default is not None and choice == '':
+            return valid[default]
+        elif choice in valid:
+            return valid[choice]
+        else:
+            sys.stdout.write("Please respond with 'yes' or 'no' "\
+                             "(or 'y' or 'n').\n")
+
 class Command(BaseCommand):
     help = 'Reindex everything.'
     args = ''
-
+    
+    option_list = BaseCommand.option_list + (
+        make_option('-n', '--book-id', action='store_true', dest='book_id', default=False,
+            help='book id instead of slugs'),
+        make_option('-t', '--just-tags', action='store_true', dest='just_tags', default=False,
+            help='just reindex tags'),
+    )
     def handle(self, *args, **opts):
         from catalogue.models import Book
-        import search
-        idx = search.ReusableIndex()
-        idx.open()
-
-        if args:
-            books = []
-            for a in args:
-                books += Book.objects.filter(slug=a).all()
-        else:
-            books = Book.objects.all()
-            
-        for b in books:
-            print b.title
-            idx.index_book(b, None)
+        from search.index import Index
+        idx = Index()
+        
+        if not opts['just_tags']:
+            if args:
+                books = []
+                for a in args:
+                    if opts['book_id']:
+                        books += Book.objects.filter(id=int(a)).all()
+                    else:
+                        books += Book.objects.filter(slug=a).all()
+            else:
+                books = list(Book.objects.all())
+
+            while books:
+                try:
+                    b = books[0]
+                    print b.title
+                    idx.index_book(b)
+                    idx.index.commit()
+                    books.pop(0)
+                except Exception, e:
+                    print "Error occured: %s" % e
+                    try:
+                        # we might not be able to rollback
+                        idx.index.rollback()
+                    except:
+                        pass
+                    retry = query_yes_no("Retry?")
+                    if not retry:
+                        break
+
         print 'Reindexing tags.'
         idx.index_tags()
-        idx.close()
+        idx.index.commit()