Migrate to docker cluster in Tsuru 0.6.0

migrate.py

import redis
import os

redis_prefix = os.environ.get('REDIS_PREFIX', 'docker-cluster')
redis_host = os.environ.get('REDIS_HOST', 'localhost')
redis_port = int(os.environ.get('REDIS_PORT', '6379'))
mongo_uri = os.environ.get('MONGO_URI')
mongo_db = os.environ.get('MONGO_DB')


def redis_to_redis():
    r = redis.Redis(host=redis_host, port=redis_port)
    images = r.keys(redis_prefix + ":image:*")
    for img_key in images:
        if r.type(img_key) == "set":
            continue
        hosts = r.lrange(img_key, 0, -1)
        tmp_key = "tmp_" + img_key
        r.delete(tmp_key)
        for h in hosts:
            r.sadd(tmp_key, h)
        r.rename(tmp_key, img_key)


def redis_to_mongo():
    import pymongo
    client = pymongo.MongoClient(mongo_uri)
    db = client[mongo_db]
    r = redis.Redis(host=redis_host, port=redis_port)
    keys = r.keys(redis_prefix + "*")
    for key in keys:
        parts = key.split(':')
        if len(parts) == 1:
            continue
        if len(parts) == 2:
            if parts[1] == 'nodes':
                to_mongo_nodes(db, r, key)
            else:
                to_mongo_container(db, r, key)
        elif parts[1] == 'image':
            to_mongo_image(db, r, key)


def to_mongo_image(db, r, key):
    image = "".join(key.split(":")[2:])
    hosts = list(r.smembers(key))
    db.images.update({"_id": image}, {"$addToSet": {"hosts": {"$each": hosts}}}, upsert=True)


def to_mongo_container(db, r, key):
    cont = key.split(":")[-1]
    host = r.get(key)
    db.containers.update({"_id": cont}, {"$set": {"host": host}}, upsert=True)


def to_mongo_nodes(db, r, key):
    nodes = r.smembers(key)
    for node in nodes:
        healing = r.get(redis_prefix + ":node:healing:" + node) == "1"
        healing_is_fail = r.get(redis_prefix + ":node:healing:isfailure:" + node) == "1"
        metadata = r.hgetall(redis_prefix + ":node:metadata:" + node)
        db.nodes.update({"_id": node}, {"$set":{
            "metadata": metadata,
            "healing": {"locked": healing, "isfailure": healing_is_fail}
        }}, upsert=True)


def main():
    redis_to_redis()
    if mongo_uri:
        redis_to_mongo()


if __name__ == '__main__':
    main()