Manage.py fixes

aore/aoutils/aoupdater.py

@@ -18,6 +18,7 @@ class AoUpdater:
         self.db_handler = DbHandler()
         self.mode = source
         self.updalist_generator = None
+        self.tablelist_generator = None
         self.allowed_tables = None
 
     def __get_entries_from_folder(self, path_to_xmls):
@@ -33,10 +34,16 @@ class AoUpdater:
         # TODO: Вычислять версию, если берем данные из каталога
         yield dict(intver=0, textver="Unknown", url=foldername)
 
+    def __get_updates_from_rar(self, url):
+        aorar = AoRar()
+        fname = aorar.download(url)
+        for table_entry in aorar.get_table_entries(fname, allowed_tables):
+            yield table_entry
+
     def __init_update_entries(self, full_base):
         if self.mode == "http":
+            self.tablelist_generator = self.__get_updates_from_rar
             imp = Importer()
-            self.updalist_generator = None
             if full_base:
                 self.updalist_generator = imp.get_full()
             else:
@@ -44,6 +51,7 @@ class AoUpdater:
         else:
             assert path.isdir(self.mode), "Invalid directory {}".format(self.mode)
             self.updalist_generator = self.__get_updates_from_folder(self.mode)
+            self.tablelist_generator = self.__get_entries_from_folder
 
     def process_single_entry(self, table_xmlentry, chunck_size=50000):
         aoparser = AoDataParser(table_xmlentry, chunck_size)
@@ -54,7 +62,7 @@ class AoUpdater:
         self.db_handler.pre_create()
 
         for update_entry in self.updalist_generator:
-            for table_entry in self.__get_entries_from_folder(update_entry['url']):
+            for table_entry in self.tablelist_generator(update_entry['url']):
                 self.process_single_entry(table_entry)
 
         logging.warning("Create success")
@@ -70,9 +78,7 @@ class AoUpdater:
                 logging.warning("Maximum count of updates are processed - exit")
                 break
 
-            aorar = AoRar()
-            fname = aorar.download(update_entry['url'])
-            for table_entry in aorar.get_table_entries(fname, allowed_tables):
+            for table_entry in self.tablelist_generator(update_entry['url']):
                 self.process_single_entry(table_entry)
 
         logging.warning("Update success")

manage.py

@@ -5,31 +5,34 @@ import optparse
 from aore.aoutils.aoupdater import AoUpdater
 
 
-def update_base(updates_count):
-    aoupdater = AoUpdater()
-    aoupdater.update_db(updates_count)
+def update_base(xml_source, updates_count):
+    aoupdater = AoUpdater(xml_source)
+    aoupdater.update(updates_count)
 
 
-def create_base(path_to_xmls):
-    aoupdater = AoUpdater(path_to_xmls)
+def create_base(xml_source):
+    aoupdater = AoUpdater(xml_source)
     aoupdater.create()
 
 
 def main():
     # Parse options
     p = optparse.OptionParser()
-    p.add_option('--create', '-c', help="Create DB from official full XMLs; "
-                                        "CREATE = path to xml source dir")
-    p.add_option('--update', '-u', help="Update DB from official delta archive; "
-                                        "UPDATE = count of updates")
+    p.add_option('--database', '-b', action="store", type="string",
+                 help="Manage database. Value: create - create new DB, update - update existing DB without loose the data")
+    p.add_option('--update-count', '-u', default=1, type="int",
+                 help="Count of updates to process, only for '--database update' option")
+    p.add_option('--source', '-s', default="http",
+                 help="Create/update DB from source. Value: \"http\" or absolute path to folder")
     options, arguments = p.parse_args()
 
-    # create new database
-    if options.create:
-        create_base(options.create)
-    # update database
-    if options.update:
-        update_base(int(options.update))
+    if options.database and options.source:
+        # create new database
+        if options.database == "create":
+            create_base(options.source)
+        # update database
+        if options.database == "update":
+            update_base(options.source, int(options.update_count))
 
 
 if __name__ == '__main__':