prom-extract

#!/usr/bin/env python3

# Copyright 2021 Robert Krawitz/Red Hat
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
import sys
import time
import os
from datetime import datetime, timezone, timedelta
import json
import subprocess
import urllib3
import uuid
import yaml
try:
    from prometheus_api_client import PrometheusConnect
    import argparse
    from jinja2 import Template
    import selectors
except ModuleNotFoundError as exc:
    print(f"prom-extract failed: {exc}", file=sys.stderr)
    print("{}")
    sys.exit(1)

# Both openshift and openshift_client provide packages named openshift.
# Make sure that we have the correct one.
try:
    import openshift
    _ = openshift.project
except (ModuleNotFoundError, AttributeError) as exc:
    print(f"prom-extract failed: {exc} (need to install 'openshift_client', not 'openshift')", file=sys.stderr)
    print("{}")
    sys.exit(1)


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


def efail(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)
    sys.exit(1)


def run_command(cmd, fail_on_bad_status=True, report_stderr_async=True,
                report_stdout_async=False):
    """ Run specified command, capturing stdout and stderr as array of timestamped lines.
        Optionally fail if return status is non-zero.  Also optionally report
        stdout and/or stderr to the appropriate file descriptors
    """
    with subprocess.Popen(cmd, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE, stderr=subprocess.PIPE) as command:
        stdout_data = []
        stderr_data = []

        sel = selectors.DefaultSelector()
        sel.register(command.stdout, selectors.EVENT_READ)
        sel.register(command.stderr, selectors.EVENT_READ)
        while True:
            # Keep reading until we reach EOF on both channels.
            # command.poll() is not a good criterion because the process
            # might complete before everything has been read.
            foundSomething = False
            for key, _ in sel.select():
                data = key.fileobj.readline()
                if len(data) > 0:
                    foundSomething = True
                    data = data.decode().rstrip()
                    if key.fileobj is command.stdout:
                        stdout_data.append([datetime.now(timezone.utc).isoformat(), data])
                        if report_stdout_async:
                            print(data)
                    elif key.fileobj is command.stderr:
                        stderr_data.append([datetime.now(timezone.utc).isoformat(), data])
                        if report_stderr_async:
                            print(data, file=sys.stderr)
            if not foundSomething:
                while command.poll() is None:
                    time.sleep(1)
                if fail_on_bad_status and command.poll() != 0:
                    raise RuntimeError('Command %s failed: exit status %d' % (' '.join(cmd), command.poll()))
                return (stdout_data, stderr_data, command.poll())


def get_prometheus_default_url():
    try:
        with openshift.project('openshift-monitoring'):
            return 'https://%s' % openshift.selector(['route/prometheus-k8s']).objects()[0].as_dict()['spec']['host']
    except Exception as err:
        efail("Unable to retrieve prometheus-k8s route: %s" % err)


def get_prometheus_token_by_version():
    with openshift.project('openshift-monitoring'):
        openshift_version = openshift.get_server_version().split('-')[0].split('.')
        if int(openshift_version[0]) > 4 or (int(openshift_version[0]) == 4 and int(openshift_version[1]) > 10):
            return openshift.invoke('sa', ['new-token', '-n', 'openshift-monitoring', 'prometheus-k8s']).out().strip()
        else:
            return openshift.get_serviceaccount_auth_token('prometheus-k8s')


def get_prometheus_token():
    try:
        return 'Bearer %s' % get_prometheus_token_by_version()
    except Exception as err:
        efail("Unable to retrieve prometheus-k8s token: %s" % err)


def get_prometheus_timestamp():
    retries = 5
    while retries > 0:
        retries = retries - 1
        try:
            with openshift.project('openshift-monitoring'):
                result = openshift.selector('pod/prometheus-k8s-0').object().execute(['date', '+%s.%N'],
                                                                                     container_name='prometheus')
                return datetime.utcfromtimestamp(float(result.out()))
        except Exception as err:
            if retries <= 0:
                efail("Unable to retrieve date: %s" % err)
            else:
                time.sleep(5)


def get_nodes():
    try:
        return json.loads([n.as_dict() for n in openshift.selector(['node']).objects()])
    except Exception as err:
        efail("Unable to retrieve cluster version: %s" % err)


def generate_uuid():
    return str(uuid.uuid4())


def get_object(kind, name):
    return openshift.selector('%s/%s' % (kind, name)).objects()[0].as_dict()


try:
    parser = argparse.ArgumentParser(description='Scrape data from Prometheus')
    parser.add_argument('-u', '--url', '--prometheus-url', type=str,
                        help='Prometheus URL', metavar='URL',
                        default=get_prometheus_default_url())
    parser.add_argument('-s', '--step', type=int, default=30, metavar='seconds',
                        help='Step duration')
    parser.add_argument('-t', '--token', type=str,
                        help='Prometheus authentication token', metavar='token',
                        default=get_prometheus_token())
    parser.add_argument('-m', '--metrics-profile', type=str, metavar='file',
                        help='Metrics profile file or URL', default='metrics.yaml')
    parser.add_argument('--metrics-only', action='store_true',
                        help='Generate metrics for specified start time and optional end time only')
    parser.add_argument('--start_time', type=int, metavar='time',
                        help='Metrics start time in seconds from epoch', default=None)
    parser.add_argument('--end_time', type=int, metavar='time',
                        help='Metrics end time in seconds from epoch', default=None)
    parser.add_argument('--epoch', type=int, default=60, metavar='seconds',
                        help='Start of metrics relative to job start')
    parser.add_argument('--post-settling-time', type=int, default=60, metavar='seconds',
                        help='Time to continue collecting metrics after job completion')
    parser.add_argument('--json-from-command', action='store_true',
                        help='Interpret command stdout as JSON')
    parser.add_argument('--uuid', type=str, metavar='UUID', default=generate_uuid(),
                        help='Index results by UUID (generate if not provided)')
    parser.add_argument('--job_type', '--job-type', help='Type of job (fio, uperf, etc)',
                        metavar='command', type=str)
    parser.add_argument('-D', '--define', help='Define template variable for metrics YAML; namespace_re to specify namespace regex',
                        metavar='name=value', action='append')
    parser.add_argument('command', metavar='command', help='command [args...]',
                        type=str, nargs='*')
    args = parser.parse_args()

    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
    authorization = {}

    if args.token != '':
        authorization['Authorization'] = args.token

    prom = PrometheusConnect(url=args.url, disable_ssl=True, headers=authorization)

    try:
        template_dict = dict()
        # Default namespace.
        template_dict['namespace_re'] = "clusterbuster-.*"
        failed_args = False
        if args.define:
            for arg in args.define:
                if '=' not in arg:
                    eprint(f'Definition {arg} invalid (must be name=value)')
                    failed_args = True
                else:
                    key, value = arg.split('=', 1)
                    template_dict[key] = value
            if failed_args:
                sys.exit(1)
        try:
            with open(args.metrics_profile, 'r') as metrics_yaml:
                baretext = metrics_yaml.read()
                yamltxt = Template(baretext).render(template_dict)
                yaml = yaml.safe_load(yamltxt)
        except FileNotFoundError as err:
            efail(f'Cannot open metrics profile: {err}')
    except Exception as err:
        efail(f'Cannot load metrics definition: {err}')

    if args.start_time:
        startTime = datetime.fromtimestamp(args.start_time)
    else:
        startTime = get_prometheus_timestamp()
    metricsStartTime = startTime + timedelta(seconds=-abs(args.epoch))

    if not args.metrics_only:
        stdout_data = []
        stderr_data = []

        os.environ['_BENCH_ARMY_KNIFE_EXTRACT_DATA'] = '1'
        try:
            if args.command:
                stdout_data, stderr_data, cmd_exit_status = run_command(args.command)
            else:
                def _read_stdin():
                    readline = sys.stdin.readline()
                    while readline:
                        yield readline
                        readline = sys.stdin.readline()
                for line in _read_stdin():
                    stdout_data.append([datetime.now(timezone.utc).isoformat(), line.decode().rstrip()])
        except KeyboardInterrupt:
            efail("Interrupted")
        except Exception as err:
            efail(err)
        endTime = get_prometheus_timestamp()

        json_output = None
        if args.json_from_command and len(stdout_data) > 0:
            try:
                json_output = json.loads("\n".join(a[1] for a in stdout_data))
            except json.decoder.JSONDecodeError as err:
                eprint("Cannot decode command output as JSON: %s" % err)

        json_results = {}
        json_api_objects = []
        if json_output is not None:
            if 'results' in json_output:
                json_results = json_output['results']
                del json_output['results']

            if 'api_objects' in json_output:
                for api_object in json_output['api_objects']:
                    bad_object = False
                    for tag in ['kind', 'namespace', 'name']:
                        if tag not in api_object:
                            eprint("API object %s does not contain a %s" % (api_object, tag))
                            bad_object = True
                    if not bad_object:
                        try:
                            with openshift.project(api_object['namespace']):
                                try:
                                    apiobj = get_object(api_object['kind'], api_object['name'])
                                    if apiobj is not None:
                                        json_api_objects.append(apiobj)
                                except Exception as get_err:
                                    eprint("Unable to retrieve object %s/%s in namespace %s: %s" %
                                           (api_object['kind'], api_object['name'], api_object['namespace'], get_err))
                        except Exception as ns_err:
                            eprint("Unable to set namespace %s: %s" % (api_object['namespace'], ns_err))
                del json_output['api_objects']

            json_output.pop('results', None)

    if args.post_settling_time > 0:
        eprint("Waiting %d seconds for complete metrics results" % args.post_settling_time)
        time.sleep(args.post_settling_time)

    if args.end_time:
        metricsEndTime = datetime.utcfromtimestamp(args.end_time)
    else:
        metricsEndTime = get_prometheus_timestamp()

    metric_results = {}

    for metric in yaml['metrics']:
        if 'query' not in metric:
            continue
        metric_data = []
        try:
            if 'instant' not in metric or metric['instant'] is not True:
                metric_data = prom.custom_query_range(metric['query'], start_time=metricsStartTime,
                                                      end_time=metricsEndTime, step=args.step)
            else:
                metric_data = prom.custom_query(metric['query'])
        except Exception as err:
            eprint(f"Query {metric['metricName']} ({metric['query']}) failed: {err}")
        name = metric['query']
        if 'metricName' in metric:
            name = metric['metricName']
        metric_results[name] = {
            'query': metric['query'],
            'name': name,
            'data': metric_data
        }

    if args.metrics_only:
        results = metric_results
    else:
        results = {
            'metadata': {
                'jobStart': startTime.isoformat(timespec='seconds'),
                'jobEnd': endTime.isoformat(timespec='seconds'),
                'uuid': args.uuid,
                'cluster_version': openshift.get_server_version(),
                'nodes': [n.as_dict() for n in openshift.selector('nodes').objects()],
                'command': args.command,
            },
            'rundata': {
                'metrics': metric_results,
                'stderr': stderr_data
            }
        }

        if json is not None:
            if json_results:
                results['rundata']['results'] = json_results
            if json_api_objects:
                results['rundata']['api_objects'] = json_api_objects
            if json_output:
                results['rundata']['run_data'] = json_output
        else:
            if stdout_data:
                results['rundata']['stdout'] = stdout_data

    print(json.dumps(results, indent=4))
except (KeyboardInterrupt, BrokenPipeError):
    sys.exit(1)