search index building on import
[wolnelektury.git] / apps / catalogue / management / commands / importbooks.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 os
6 import sys
7 import time
8 from optparse import make_option
9
10 from django.core.management.base import BaseCommand
11 from django.core.management.color import color_style
12 from django.core.files import File
13
14 from catalogue.models import Book
15
16
17 class Command(BaseCommand):
18     option_list = BaseCommand.option_list + (
19         make_option('-q', '--quiet', action='store_false', dest='verbose', default=True,
20             help='Verbosity level; 0=minimal output, 1=normal output, 2=all output'),
21         make_option('-f', '--force', action='store_true', dest='force', default=False,
22             help='Print status messages to stdout'),
23         make_option('-E', '--no-build-epub', action='store_false', dest='build_epub', default=True,
24             help='Don\'t build EPUB file'),
25         make_option('-M', '--no-build-mobi', action='store_false', dest='build_mobi', default=True,
26             help='Don\'t build MOBI file'),
27         make_option('-T', '--no-build-txt', action='store_false', dest='build_txt', default=True,
28             help='Don\'t build TXT file'),
29         make_option('-P', '--no-build-pdf', action='store_false', dest='build_pdf', default=True,
30             help='Don\'t build PDF file'),
31         make_option('-S', '--no-search-index', action='store_false', dest='search_index', default=True,
32             help='Don\'t build PDF file'),
33         make_option('-w', '--wait-until', dest='wait_until', metavar='TIME',
34             help='Wait until specified time (Y-M-D h:m:s)'),
35     )
36     help = 'Imports books from the specified directories.'
37     args = 'directory [directory ...]'
38
39     def handle(self, *directories, **options):
40         from django.db import transaction
41
42         self.style = color_style()
43
44         verbose = options.get('verbose')
45         force = options.get('force')
46         show_traceback = options.get('traceback', False)
47
48         wait_until = None
49         if options.get('wait_until'):
50             wait_until = time.mktime(time.strptime(options.get('wait_until'), '%Y-%m-%d %H:%M:%S'))
51             if verbose > 0:
52                 print "Will wait until %s; it's %f seconds from now" % (
53                     time.strftime('%Y-%m-%d %H:%M:%S', 
54                     time.localtime(wait_until)), wait_until - time.time())
55
56         # Start transaction management.
57         transaction.commit_unless_managed()
58         transaction.enter_transaction_management()
59         transaction.managed(True)
60
61         files_imported = 0
62         files_skipped = 0
63
64         for dir_name in directories:
65             if not os.path.isdir(dir_name):
66                 print self.style.ERROR("%s: Not a directory. Skipping." % dir_name)
67             else:
68                 # files queue
69                 files = sorted(os.listdir(dir_name))
70                 postponed = {}
71                 while files:
72                     file_name = files.pop(0)
73                     file_path = os.path.join(dir_name, file_name)
74                     file_base, ext = os.path.splitext(file_path)
75
76                     # Skip files that are not XML files
77                     if not ext == '.xml':
78                         continue
79
80                     if verbose > 0:
81                         print "Parsing '%s'" % file_path
82                     else:
83                         sys.stdout.write('.')
84                         sys.stdout.flush()
85
86                     # Import book files
87                     try:
88                         book = Book.from_xml_file(file_path, overwrite=force, 
89                                                   build_epub=options.get('build_epub'),
90                                                   build_txt=options.get('build_txt'),
91                                                   build_pdf=options.get('build_pdf'),
92                                                   build_mobi=options.get('build_mobi'),
93                                                   search_index=options.get('search_index'))
94                         files_imported += 1
95
96                         if os.path.isfile(file_base + '.pdf'):
97                             book.pdf_file.save('%s.pdf' % book.slug, File(file(file_base + '.pdf')))
98                             if verbose:
99                                 print "Importing %s.pdf" % file_base
100                         if os.path.isfile(file_base + '.mobi'):
101                             book.mobi_file.save('%s.mobi' % book.slug, File(file(file_base + '.mobi')))
102                             if verbose:
103                                 print "Importing %s.mobi" % file_base
104                         if os.path.isfile(file_base + '.epub'):
105                             book.epub_file.save('%s.epub' % book.slug, File(file(file_base + '.epub')))
106                             if verbose:
107                                 print "Importing %s.epub" % file_base
108                         if os.path.isfile(file_base + '.txt'):
109                             book.txt_file.save('%s.txt' % book.slug, File(file(file_base + '.txt')))
110                             if verbose:
111                                 print "Importing %s.txt" % file_base
112
113                         book.save()
114
115                     except Book.AlreadyExists, msg:
116                         print self.style.ERROR('%s: Book already imported. Skipping. To overwrite use --force.' %
117                             file_path)
118                         files_skipped += 1
119
120                     except Book.DoesNotExist, e:
121                         if file_name not in postponed or postponed[file_name] < files_imported:
122                             # push it back into the queue, maybe the missing child will show up
123                             if verbose:
124                                 print self.style.NOTICE('Waiting for missing children')
125                             files.append(file_name)
126                             postponed[file_name] = files_imported
127                         else:
128                             # we're in a loop, nothing's being imported - some child is really missing
129                             raise e
130
131         # Print results
132         print
133         print "Results: %d files imported, %d skipped, %d total." % (
134             files_imported, files_skipped, files_imported + files_skipped)
135         print
136
137         if wait_until:
138             print 'Waiting...'
139             try:
140                 time.sleep(wait_until - time.time())
141             except IOError:
142                 print "it's already too late"
143
144         transaction.commit()
145         transaction.leave_transaction_management()
146