diff --git a/aore/aoutils/aoupdater.py b/aore/aoutils/aoupdater.py index 790dd5a..9c3c6a4 100644 --- a/aore/aoutils/aoupdater.py +++ b/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") diff --git a/manage.py b/manage.py index 4dc0a1f..81b9e74 100644 --- a/manage.py +++ b/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__':