import argparse import json import logging import os.path import shutil import sys import time from terroroftinytown.release.botouploader import BotoUploaderBootstrap from terroroftinytown.release.iaupload import IAUploaderBootstrap from terroroftinytown.tracker.bootstrap import Bootstrap from terroroftinytown.tracker.export import ExporterBootstrap from terroroftinytown.tracker.model import Result logger = logging.getLogger(__name__) UPLOADER_CLASS_MAP = { 'ia': IAUploaderBootstrap, 'boto': BotoUploaderBootstrap, } def main(): arg_parser = argparse.ArgumentParser() arg_parser.add_argument('config_path') arg_parser.add_argument('export_dir', default='/home/tinytown/tinytown-export/') arg_parser.add_argument('--verbose', action='store_true') arg_parser.add_argument('--uploader', choices=['ia', 'boto'], default='boto') arg_parser.add_argument('--batch-size', type=int) arg_parser.add_argument('--min-batch-size', type=int) arg_parser.add_argument('--max-batches', type=int, default=1) args = arg_parser.parse_args() if not os.path.exists(args.export_dir): os.mkdir(args.export_dir) log_filename = os.path.join(args.export_dir, 'supervisor.log') logging.basicConfig( level=logging.INFO, filename=log_filename, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) if args.verbose: stream_handler = logging.StreamHandler() stream_handler.setLevel(logging.INFO) formatter = logging.Formatter('%(levelname)s:%(name)s:%(message)s') stream_handler.setFormatter(formatter) logging.getLogger().addHandler(stream_handler) logger.info('Supervisor starting up.') try: wrapper(args) except (Exception, SystemExit): logger.exception('Failure.') raise logger.info('Supervisor done.') def wrapper(args): config_path = args.config_path export_dir = args.export_dir sentinel_file = os.path.join(export_dir, 'tinytown-supervisor-sentinel') if os.path.exists(sentinel_file): raise Exception( 'The sentinel file exists. ' 'Previous supervisor did not exit correctly.' ) else: with open(sentinel_file, 'wb'): pass if sys.version_info[0] != 3: raise Exception('This script expects Python 3') if not os.path.isfile(config_path): raise Exception('Config path is not a file.') for batch_num in range(args.max_batches): if has_results(args): logging.info('Starting batch #%d', batch_num + 1) process_batch(args) time.sleep(2) else: logger.info('No results. Nothing to do.') break os.remove(sentinel_file) logger.info('Done') def process_batch(args): config_path = args.config_path export_dir = args.export_dir uploader_class = UPLOADER_CLASS_MAP[args.uploader] database_locked_sentinel_path = os.path.join( args.export_dir, 'database_is_locked') logger.info('Loading bootstrap.') bootstrap = Bootstrap() bootstrap.setup_args() bootstrap.parse_args(args=[config_path]) bootstrap.load_config() time_struct = time.gmtime() timestamp = bootstrap.config.get( 'iaexporter', 'timestamp', ).format( year=time_struct.tm_year, month=time_struct.tm_mon, day=time_struct.tm_mday, hour=time_struct.tm_hour, minute=time_struct.tm_min, second=time_struct.tm_sec, ) done_directory = os.path.join(export_dir, 'done') if not os.path.exists(done_directory): os.mkdir(done_directory) item_export_directory = os.path.join(export_dir, timestamp) logger.info('Begin export to %s.', item_export_directory) title = bootstrap.config.get('iaexporter', 'title')\ .format(timestamp=timestamp) identifier = bootstrap.config.get('iaexporter', 'item')\ .format(timestamp=timestamp) upload_meta_path = os.path.join(export_dir, 'current.json') upload_meta = { 'identifier': identifier, 'title': title, 'work_directory': item_export_directory, } with open(upload_meta_path, 'w') as out_file: out_file.write(json.dumps(upload_meta)) os.makedirs(item_export_directory) exporter = ExporterBootstrap() exporter_args = [ config_path, '--format', 'beacon', '--include-settings', '--zip', '--dir-length', '0', '--file-length', '0', '--max-right', '8', '--delete', '--zip-filename-infix', '.{}'.format(timestamp), '--database-busy-file', database_locked_sentinel_path, item_export_directory, ] if args.batch_size: exporter_args.extend(['--max-items', str(args.batch_size)]) export_dir_start_size = get_dir_size(item_export_directory) exporter.start(args=exporter_args) logger.info('Export finished') export_dir_end_size = get_dir_size(item_export_directory) if export_dir_start_size == export_dir_end_size: raise Exception('Export directory size did not change: {} bytes' .format(export_dir_end_size)) logger.info('Upload starting') uploader = uploader_class() args = [ config_path, item_export_directory, '--title', upload_meta['title'], '--identifier', upload_meta['identifier'] ] uploader.start(args=args) logger.info('Upload done.') os.remove(upload_meta_path) shutil.move(item_export_directory, done_directory) def has_results(args): bootstrap = Bootstrap() bootstrap.setup_args() bootstrap.parse_args(args=[args.config_path]) bootstrap.load_config() bootstrap.setup_database() if Result.has_results(): if args.min_batch_size: return Result.get_count() >= args.min_batch_size else: return True def get_dir_size(path): total = 0 for root, dirs, files in os.walk(path): total += sum(os.path.getsize(os.path.join(root, name)) for name in files) return total if __name__ == '__main__': main()