prometheus-exporter.py

import argparse
import random
import re
import time
from datetime import datetime, timedelta

from prometheus_client import generate_latest, start_http_server, Gauge

COMMENT_RE = r"^\s*\#.*$"
# LETTER = random.choice(string.ascii_letters)
IS_SERVER = False
LETTER = "I am here"
BEACON_STR = 'uniqlo'


def attach_timestamp(line, timestamp):
    if re.match(COMMENT_RE, line):
        return line
    return "{0} {1}".format(line, timestamp)


def format_lines(lines, remove_comments=True, delta=0):
    t_delta = datetime.now() + timedelta(seconds=delta)
    timestamp = int(round(t_delta.timestamp() * 1000))
    if remove_comments:
        return [attach_timestamp(line, timestamp=timestamp) for line in lines if
                line and not re.match(COMMENT_RE, line)]
    else:
        return [attach_timestamp(line, timestamp=timestamp) for line in lines if line]


def setup(num_metrics=10):
    dummies = []
    label_names = ['foo', BEACON_STR]
    for idx in range(num_metrics):
        c = Gauge('dummy_events_{0}'.format(idx), 'This is dummy events counter number {0}'.format(idx),
                  labelnames=label_names)
        dummies.append(c)
    return dummies


def pulse(dummies, norand=False, delta=0):
    metric_stats = []
    labels = [LETTER, BEACON_STR]
    for idx, c in enumerate(dummies):
        if norand:
            c.labels(*labels).inc(1)
        else:
            c.labels(*labels).set(random.randint(0, 10000))
        stats = generate_latest(c).decode("utf-8").split("\n")
        lines = format_lines(stats, remove_comments=False, delta=delta)
        metric_stats.extend(lines)
    return metric_stats


def main():
    global BEACON_STR
    global IS_SERVER
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", help="Run in server mode", action="store_true", default=False)
    parser.add_argument("--port", help="Port for the dummy server", type=int, default=8000)
    parser.add_argument("--norand", help="Do not populate random values for metrics", action="store_true",
                        default=False)
    parser.add_argument("--entries", help="Number of unique entries to create", type=int, default=10)
    parser.add_argument("--delta", help="Time delta in seconds", type=int, default=15)
    parser.add_argument("--metrics", help="Number of metrics to create", type=int, default=10)
    args = parser.parse_args()

    is_server = args.server
    port = args.port
    norand = args.norand
    num_entries = args.entries
    num_metrics = args.metrics
    delta = args.delta

    IS_SERVER = is_server

    if is_server:
        print("Running as server!")
        print("Running on port {0}".format(port))
        BEACON_STR = "federer"
    else:
        print("Running as one-time exporter!")
        print("Number of entries {0}".format(num_entries))
        print("Time delta in seconds {0}".format(delta))
    print("No random numbers? {0}".format(norand))
    print("Number of metrics {0}".format(num_metrics))

    dummies = setup(num_metrics=num_metrics)

    if is_server:
        start_http_server(port)
        while True:
            _ = pulse(dummies, norand=norand)
            time.sleep(1)
    else:
        all_stats = []
        for idx in range(num_entries):
            stats = pulse(dummies, norand=norand, delta=delta * idx)
            if delta < 0:
                all_stats.insert(0, stats)
            else:
                all_stats.append(stats)
        with open('dummy_prometheus_metrics.dat', 'w+') as fi:
            for stats in all_stats:
                for stat in stats:
                    fi.write(stat)
                    fi.write("\n")
    return


if __name__ == "__main__":
    main()