update_marathon_app_definitions.py

#!/usr/bin/env python
import re
import json
import pprint
import argparse
import logging
import urlparse
import requests
from copy import deepcopy

logger = logging.getLogger()
logger.setLevel(logging.INFO)
handler = logging.StreamHandler()
logger.addHandler(handler)

def pp(json_object):
    logger.info(pprint.pformat(json_object, indent=4))

def get_labels(app, env_vars):
    labels = {"consul": ""}
    name_key, tags_key = None, None
    logger.info("Current labels for %s: %s" % (app["id"], app["labels"]))
    try:
        [name_key] = [k for k in env_vars.keys() if k.endswith("_NAME")]
    except ValueError, e:
        logger.info("name key doesn't exist")
    try:
        [tags_key] = [k for k in env_vars.keys() if k.endswith("_TAGS")]
    except ValueError, e:
        logger.info("tags key doesn't exist")
    if name_key:
        if env_vars[name_key] != app["id"].replace("/", ""):
            labels["consul"] = env_vars[name_key]
    if tags_key:
        for tag in env_vars[tags_key].split(","):
            labels[tag] = "tag"
    logger.info("New labels: %s" % labels)
    return labels

def get_env_vars(app):
    env_vars = {}
    pattern = "^SERVICE_[0-9]+\_"
    for k, v in app["env"].iteritems():
        if re.search(pattern, k):
            env_vars[k] = v
    return env_vars

def get_json(url):
    r = requests.get(url)
    return r.json()

def put_app(url, app_data):
    r = requests.put(url,
        data=json.dumps(app_data), headers={"Content-Type": "application/json"})
    r.raise_for_status()
    logger.info("%s - %s" % (r.status_code, r.text))

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("env", help="Environment name (eg. 'dev-core')")
    return parser.parse_args()

def main():
    args = parse_args()
    base_url = "http://marathon.service.consul.eu-%s-1.%s.hubrick.net:8080/" % (
        "central" if args.env is not "production" else "west", args.env)
    apps = get_json(urlparse.urljoin(base_url, "/v2/apps"))["apps"]
    for app in apps:
        env_vars = get_env_vars(app)
        new_labels = get_labels(app, env_vars)
        answer = raw_input("Apply these new labels to '%s' (y/n)? "
            % app["id"].replace("/", ""))
        if answer == "y":
            put_app(urlparse.urljoin(base_url, "/v2/apps%s" %app["id"]),
                {"labels": new_labels})

if __name__ == "__main__":
    main()