tb-copy

#!/usr/bin/env python

from itertools import imap
import logging
from optparse import OptionParser
import os.path
import random
import sys

from tilebox import BoundingPyramid, TileStore, consume
from tilebox.filter.benchmark import Benchmark, Statsd, StatsdCountErrors, StatsdCountTiles
from tilebox.filter.consistenthash import EveryNth
from tilebox.filter.contenttype import ContentTypeAdder
from tilebox.filter.error import DropErrors, LogErrors, MaximumConsecutiveErrors, MaximumErrorRate, MaximumErrors
from tilebox.filter.logger import Logger
from tilebox.filter.rate import RateLimit
from tilebox.store.boundingpyramid import BoundingPyramidTileStore


def main(argv):
    logger = logging.getLogger(os.path.basename(argv[0]))
    option_parser = OptionParser()
    option_parser.add_option('--benchmark', action='store_true')
    option_parser.add_option('-b', '--bounding-pyramid', metavar='BOUNDING-PYRAMID')
    option_parser.add_option('--add-content-type', action='store_true')
    option_parser.add_option('-i', metavar='I', type=int)
    option_parser.add_option('-g', '--generate', metavar='TILE-STORE', action='append')
    option_parser.add_option('--limit', metavar='N', type=int)
    option_parser.add_option('-m', '--move', action='store_true')
    option_parser.add_option('--maximum-consecutive-errors', metavar='N', type=int)
    option_parser.add_option('--maximum-errors', metavar='N', type=int)
    option_parser.add_option('--maximum-error-rate', metavar='FLOAT', type=float)
    option_parser.add_option('-n', metavar='N', type=int)
    option_parser.add_option('-o', '--overwrite', action='store_true')
    option_parser.add_option('-r', '--rate-limit', metavar='HZ', type=float)
    option_parser.add_option('--randomize', action='store_true')
    option_parser.add_option('--statsd', action='store_true')
    option_parser.add_option('--statsd-host', default='127.0.0.1', metavar='HOST')
    option_parser.add_option('--statsd-port', default=8125, metavar='PORT', type=int)
    option_parser.add_option('--statsd-prefix', default='tb-copy-', metavar='PREFIX')
    option_parser.add_option('-v', '--verbose', action='store_true')
    options, args = option_parser.parse_args(argv[1:])
    if options.verbose:
        logging.basicConfig(level=logging.INFO)
    else:
        logging.basicConfig(level=logging.WARNING)
    assert len(args) >= 2
    if options.bounding_pyramid:
        bounding_pyramid = BoundingPyramid.from_string(options.bounding_pyramid)
    else:
        bounding_pyramid = None
    if options.statsd:
        statsd = Statsd(options.statsd_prefix, options.statsd_host, options.statsd_port)
    else:
        statsd = None
    benchmark = Benchmark(statsd=statsd) if options.benchmark else None
    if options.generate:
        generate = map(TileStore.load, options.generate)
    else:
        generate = ()
    try:
        output_tilestore = TileStore.load(args[-1])
        for arg in args[:-1]:
            input_tilestore = TileStore.load(arg)
            if bounding_pyramid:
                tilestream = BoundingPyramidTileStore(bounding_pyramid).list()
            else:
                tilestream = input_tilestore.list()
            if options.i is not None and options.n is not None:
                tilestream = imap(EveryNth(options.n, options.i), tilestream)
            if options.randomize:
                tilestream = list(tilestream)
                random.shuffle(tilestream)
            if not options.overwrite:
                tilestream = (tile for tile in tilestream if tile not in output_tilestore)
            if options.rate_limit:
                tilestream = imap(RateLimit(options.rate_limit), tilestream)
            if benchmark:
                tilestream = imap(benchmark.sample(), tilestream)
            tilestream = input_tilestore.get(tilestream)
            if benchmark:
                tilestream = imap(benchmark.sample('get'), tilestream)
            for i, g in enumerate(generate):
                tilestream = g.get(tilestream)
                if options.benchmark:
                    tilestream = imap(benchmark.sample('generate-%d' % (i,)), tilestream)
            tilestream = imap(LogErrors(logger, logging.ERROR, '%(tilecoord)s: %(error)s'), tilestream)
            if statsd:
                tilestream = imap(StatsdCountErrors(statsd), tilestream)
            if options.maximum_consecutive_errors:
                tilestream = imap(MaximumConsecutiveErrors(options.maximum_consecutive_errors), tilestream)
            if options.maximum_error_rate:
                tilestream = imap(MaximumErrorRate(options.maximum_error_rate), tilestream)
            if options.maximum_errors:
                tilestream = imap(MaximumErrors(options.maximum_errors), tilestream)
            tilestream = imap(DropErrors(), tilestream)
            if options.add_content_type:
                tilestream = imap(ContentTypeAdder(), tilestream)
            tilestream = output_tilestore.put(tilestream)
            if benchmark:
                tilestream = imap(benchmark.sample('put'), tilestream)
            if options.move:
                tilestream = input_tilestore.delete(tilestream)
                if benchmark:
                    tilestream = imap(benchmark.sample('delete'), tilestream)
            if options.verbose:
                tilestream = imap(Logger(logger, logging.INFO, '%(tilecoord)s'), tilestream)
            if statsd:
                tilestream = imap(StatsdCountTiles(statsd), tilestream)
            consume(tilestream, options.limit)
    finally:
        logging.basicConfig(level=logging.INFO)
        if benchmark:
            keys = ['get']
            keys.extend('generate-%i' % (i,) for i in xrange(0, len(generate)))
            keys.extend(['put', 'delete'])
            for key in keys:
                if key in benchmark.statisticss:
                    statistics = benchmark.statisticss[key]
                    logger.info('%s: %s%s' % (key, statistics, ' (%.1f tiles/s)' % (1.0 / statistics.mean,) if statistics.n else ''))


if __name__ == '__main__':
    sys.exit(main(sys.argv))