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


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


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


try:
    import sys
    from prometheus_api_client import PrometheusConnect
    import argparse
    from jinja2 import Template
    import selectors
    import time
    import os
    from datetime import datetime, timezone, timedelta
    import json
    import subprocess
    import urllib3
    import uuid
    import yaml
    import openshift_client
except (KeyboardInterrupt, BrokenPipeError):
    efail('')
except ModuleNotFoundError as exc:
    efail(f"prom-extract failed: {exc}")


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_client.project('openshift-monitoring'):
            return 'https://%s' % openshift_client.selector(['route/prometheus-k8s']).objects()[0].as_dict()['spec']['host']
    except (KeyboardInterrupt, BrokenPipeError):
        efail('')
    except Exception as err:
        efail("Unable to retrieve prometheus-k8s route: %s" % err)


def get_prometheus_token_by_version():
    with openshift_client.project('openshift-monitoring'):
        openshift_version = openshift_client.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_client.invoke('sa', ['new-token', '-n', 'openshift-monitoring', 'prometheus-k8s']).out().strip()
        else:
            return openshift_client.get_serviceaccount_auth_token('prometheus-k8s')


def get_prometheus_token():
    try:
        return 'Bearer %s' % get_prometheus_token_by_version()
    except (KeyboardInterrupt, BrokenPipeError):
        efail('')
    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_client.project('openshift-monitoring'):
                result = openshift_client.selector('pod/prometheus-k8s-0').object().execute(['date', '+%s.%N'],
                                                                                     container_name='prometheus')
                return datetime.fromtimestamp(float(result.out()), timezone.utc)
        except (KeyboardInterrupt, BrokenPipeError):
            efail('')
        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_client.selector(['node']).objects()])
    except (KeyboardInterrupt, BrokenPipeError):
        efail('')
    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_client.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('--indent', help='Indentation')
    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 (KeyboardInterrupt, BrokenPipeError):
        efail('')
    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, BrokenPipeError):
            efail('')
        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 (KeyboardInterrupt, BrokenPipeError):
                efail('')
            except json.decoder.JSONDecodeError as err:
                efail("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_client.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 (KeyboardInterrupt, BrokenPipeError):
                                    efail('')
                                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 (KeyboardInterrupt, BrokenPipeError):
                            efail('')
                        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 (KeyboardInterrupt, BrokenPipeError):
            efail('')
        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_client.get_server_version(),
                'nodes': [n.as_dict() for n in openshift_client.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

    s = json.dumps(results, indent=args.indent)
    print(s)
except (KeyboardInterrupt, BrokenPipeError):
    efail('')
except Exception as exc:
    efail(f"prom-extract failed: {exc}")