marathon-upstreams

marathon-upstreams

#!/usr/bin/python

# -- stdlib --
import argparse
import json
import urllib2

# -- third party --
# -- own --
# -- code --
parser = argparse.ArgumentParser("marathon-upstreams")
parser.add_argument('--marathon', default="http://mesos1:8080")
parser.add_argument('--server-records', action='store_true', default=False)
options = parser.parse_args()

opener = urllib2.build_opener()
opener.addheaders = [('Accept', 'application/json')]

conf = json.loads(opener.open('%s/v2/apps' % options.marathon).read())['apps']

'''
{
    "labels": {
        "nginx-13001": "ip-hash;",
        "nginx-13002": "foo; bar; baz;;;",
        "server-13001": "uluru-api.service",
        "server-listen": 80
    }
}
'''


def upstreams(conf):
    nginx_apps = set()
    upstream_conf = {}

    for app in conf:
        lbls = app['labels']

        for k, v in lbls.items():
            if not k.startswith('nginx-'):
                continue

            app_id = app['id'][1:]
            k = '%s-%s' % (app_id, k[6:])
            lines = [i.strip() for i in v.split(';')]
            lines = [i for i in lines if i]
            upstream_conf[k] = ''.join(['    %s;\n' % i for i in lines])
            nginx_apps.add(k)

    def get_conf(app, srv_port):
        k = '%s-%s' % (app, srv_port)
        return upstream_conf.get(k)

    for srv in urllib2.urlopen('%s/v2/tasks' % options.marathon).readlines():
        srv = srv.split('\t')
        srv = map(str.strip, srv)
        srv = [i for i in srv if i]

        if len(srv) < 2:
            continue

        if not srv[1]:  # srv port
            continue

        app, port = srv[:2]

        if '%s-%s' % (app, port) not in nginx_apps:
            continue

        tasks = list(sorted([i.strip() for i in srv[2:]]))

        if not tasks:
            continue

        print "upstream %s-%s {" % (app, port)
        print get_conf(app, port)

        for task in tasks:
            if not task:
                continue

            print "    server %s;" % task

        print "}\n"


def servers(conf):
    server_conf = []

    tasks = urllib2.urlopen('%s/v2/tasks' % options.marathon).readlines()
    tasks = [i.split('\t') for i in tasks]
    tasks = {i[0]: len([j for j in i[2:] if j.strip()]) for i in tasks}

    for app in conf:
        lbls = app['labels']
        listen = lbls.get('server-listen', 80)

        for k, v in lbls.items():
            if not k.startswith('server-'):
                continue

            if k == 'server-listen':
                continue

            app_id = app['id'][1:]

            if tasks.get(app_id, 0) <= 0:
                continue

            upstream = '%s-%s' % (app_id, k[7:])
            server_conf.append([v, upstream, listen])

    for server, upstream, listen in sorted(server_conf):
        print "server {"
        print "    listen %s;" % listen
        print "    server_name %s;" % server
        print "    location / {"
        print "        proxy_pass http://%s;" % upstream
        print "    }"
        print "}"
        print


if options.server_records:
    servers(conf)
else:
    upstreams(conf)