Merge pull request #122 from tilezen/add_toi_loader

Add a command to load tiles of interest from file

Author: iandees

logging.conf.sample

@@ -1,5 +1,5 @@
 [loggers]
-keys=root,process,seed,intersect,drain,enqueue_tiles_of_interest,dump_tiles_of_interest,wof_process_neighbourhoods,query
+keys=root,process,seed,intersect,drain,enqueue_tiles_of_interest,dump_tiles_of_interest,load_tiles_of_interest,wof_process_neighbourhoods,query
 
 [handlers]
 keys=consoleHandler
@@ -23,6 +23,12 @@ handlers=consoleHandler
 qualName=dump_tiles_of_interest
 propagate=0
 
+[logger_load_tiles_of_interest]
+level=INFO
+handlers=consoleHandler
+qualName=load_tiles_of_interest
+propagate=0
+
 [logger_process]
 level=INFO
 handlers=consoleHandler

tilequeue/command.py

@@ -21,12 +21,13 @@
 from tilequeue.tile import coord_int_zoom_up
 from tilequeue.tile import coord_marshall_int
 from tilequeue.tile import coord_unmarshall_int
-from tilequeue.tile import parse_expired_coord_string
+from tilequeue.tile import parse_expired_coord_string, deserialize_coord
 from tilequeue.tile import seed_tiles
 from tilequeue.tile import tile_generator_for_multiple_bounds
 from tilequeue.tile import tile_generator_for_single_bounds
 from tilequeue.tile import zoom_mask
 from tilequeue.top_tiles import parse_top_tiles
+from tilequeue.utils import grouper
 from tilequeue.worker import DataFetch
 from tilequeue.worker import ProcessAndFormatData
 from tilequeue.worker import QueuePrint
@@ -1214,6 +1215,79 @@ def tilequeue_dump_tiles_of_interest(cfg, peripherals):
     )
 
 
+def tilequeue_load_tiles_of_interest(cfg, peripherals):
+    """
+    Given a newline-delimited file containing tile coordinates in
+    `zoom/column/row` format, load those tiles into a new Redis SET
+    alongside the existing one, then RENAME the new to the old.
+    """
+    logger = make_logger(cfg, 'load_tiles_of_interest')
+    logger.info('Loading tiles of interest')
+
+    toi_filename = "toi.txt"
+    chunk_size = 10000
+    new_toi_key = cfg.redis_cache_set_key + "-new"
+    old_toi_key = cfg.redis_cache_set_key + "-old"
+
+    # Make a raw Redis client so we can do low-level set manipulation
+    redis_client = make_redis_client(cfg)
+
+    logger.info(
+        'Adding tiles of interest from %s to key %s',
+        toi_filename,
+        new_toi_key,
+    )
+
+    if redis_client.exists(new_toi_key):
+        raise Exception(
+            "Can't load new TOI because key {} exists. "
+            "Delete it and try again".format(new_toi_key)
+        )
+
+    n_added = 0
+    with open(toi_filename, 'r') as f:
+        for lines in grouper(f, chunk_size):
+            coords = [
+                coord_marshall_int(deserialize_coord(l))
+                for l in lines if l is not None
+            ]
+            n_added += redis_client.sadd(new_toi_key, *coords)
+            logger.info("%s total tiles added", n_added)
+
+    logger.info(
+        'Finished adding %s tiles of interest from %s to key %s',
+        n_added,
+        toi_filename,
+        new_toi_key,
+    )
+
+    if redis_client.rename(cfg.redis_cache_set_key, old_toi_key):
+        logger.info(
+            'Successfully moved existing TOI at %s to %s',
+            cfg.redis_cache_set_key,
+            old_toi_key,
+        )
+    else:
+        raise Exception(
+            "Moving existing TOI failed. WARNING: There's "
+            "probably no TOI key and this is bad!"
+        )
+
+    if redis_client.rename(new_toi_key, cfg.redis_cache_set_key):
+        logger.info(
+            'Successfully moved new TOI at %s into use at %s',
+            new_toi_key,
+            cfg.redis_cache_set_key,
+        )
+    else:
+        raise Exception(
+            "Moving new TOI failed. WARNING: There's "
+            "probably no TOI key and this is bad!"
+        )
+
+    logger.info("Finished loading TOI.")
+
+
 class TileArgumentParser(argparse.ArgumentParser):
     def error(self, message):
         sys.stderr.write('error: %s\n' % message)
@@ -1235,6 +1309,8 @@ def tilequeue_main(argv_args=None):
         ('intersect', create_command_parser(tilequeue_intersect)),
         ('dump-tiles-of-interest',
             create_command_parser(tilequeue_dump_tiles_of_interest)),
+        ('load-tiles-of-interest',
+            create_command_parser(tilequeue_load_tiles_of_interest)),
         ('enqueue-tiles-of-interest',
          create_command_parser(tilequeue_enqueue_tiles_of_interest)),
         ('tile-size', create_command_parser(tilequeue_tile_sizes)),

tilequeue/utils.py

@@ -12,3 +12,8 @@ def format_stacktrace_one_line(exc_info=None):
     stacktrace = ' | '.join([x.replace('\n', '')
                              for x in exception_lines])
     return stacktrace
+
+
+def grouper(seq, size):
+    """Collect data into fixed-length chunks or blocks"""
+    return (seq[pos:pos + size] for pos in xrange(0, len(seq), size))