dashboard.py

"""
Grafana Dashhboard API endpoints wrapper functions.
"""
import logging
import os
import json
import time

from requests.exceptions import HTTPError
from brian_dashboard_manager.grafana.utils.request import TokenRequest

logger = logging.getLogger(__name__)

NUM_RETRIES = 3


def get_dashboard_definitions(dir=None):
    """
    Returns dictionary for each dashboard JSON definition in supplied directory

    :param dir: directory to search for dashboard definitions
    :return: generator of dashboard definitions
    """
    dashboard_dir = dir or os.path.join(
        os.path.dirname(__file__), '../dashboards/')
    for (dirpath, _, filenames) in os.walk(dashboard_dir):
        for file in filenames:
            if file.endswith('.json'):
                filename = os.path.join(dirpath, file)
                dashboard = json.load(open(filename, 'r'))
                yield dashboard


def delete_dashboard(request: TokenRequest, dashboard: dict, folder_id=None):
    """
    Deletes a single dashboard for the organization
    the API token is registered to.

    Dashboard can be specified by UID or title.
    If a folder ID is not supplied, dashboard title should be globally unique.

    :param request: TokenRequest object
    :param dashboard: dashboard object with either a UID or title
    :param folder_id: folder ID to search for dashboard in
    :return: True if dashboard is considered deleted, False otherwise
    """
    try:
        uid = dashboard.get('uid')
        if uid:
            return _delete_dashboard(request, uid)
        elif dashboard.get('title'):
            logger.info(f'Deleting dashboard: {dashboard.get("title")}')
            # if a folder ID is not supplied,
            # dashboard title should be globally unique
            dash = _search_dashboard(request, dashboard, folder_id)
            if dash is None:
                return True
            uid = dash.get('uid', '')
            if uid:
                return _delete_dashboard(request, uid)
            else:
                return True
        return False

    except HTTPError as e:
        if e.response is not None and e.response.status_code == 404:
            return True
        title = dashboard.get('title')
        logger.exception(
            f'Error when deleting dashboard: {title or ""}')
        return False


def _delete_dashboard(request: TokenRequest, uid: int):
    """
    Deletes a single dashboard for the organization
    the API token is registered to.

    :param request: TokenRequest object
    :param uid: dashboard UID
    :return: True if dashboard is considered deleted, False otherwise
    """
    try:
        r = request.delete(f'api/dashboards/uid/{uid}')
        resp = r.json()
        if resp and 'deleted' in resp.get('message', ''):
            return True
    except HTTPError as e:
        if e.response is not None and e.response.status_code == 404:
            return True
        raise e
    return False


def delete_dashboards(request: TokenRequest):
    """
    Deletes all dashboards for the organization
    the API token is registered to.

    :param request: TokenRequest object
    :return: True if all dashboards are considered deleted, False otherwise
    """
    r = request.get('api/search')
    dashboards = r.json()
    if dashboards and len(dashboards) > 0:
        for dash in dashboards:
            try:
                _delete_dashboard(request, dash['uid'])
            except HTTPError:
                logger.exception(
                    f'Error when deleting dashboard with UID #{dash["uid"]}')
    return True


# Searches for a dashboard with given title
def list_dashboards(request: TokenRequest, title=None, folder_id=None):
    """
    Searches for dashboard(s) with given title.
    If no title is provided, all dashboards are returned,
    filtered by folder ID if provided.

    :param request: TokenRequest object
    :param title: optional dashboard title to search for
    :param folder_id: optional folder ID to search for dashboards in
    :return: list of dashboards matching the search criteria
    """
    param = {
        **({'query': title} if title else {}),
        'type': 'dash-db',
        'limit': 5000,
        'page': 1
    }
    if folder_id is not None:
        param['folderIds'] = folder_id

    dashboards = []

    while True:
        r = request.get('api/search', params=param)
        page = r.json()
        if page:
            dashboards.extend(page)
            if len(page) < param['limit']:
                break
            param['page'] += 1
        else:
            break

    return dashboards


# Searches Grafana for a dashboard
# matching the title of the provided dashboard.
def _search_dashboard(request: TokenRequest, dashboard: dict, folder_id=None):
    """
    Searches Grafana for a dashboard with given title from the supplied dict.
    Primarily used to get the provisioned dashboard definition if it exists

    :param request: TokenRequest object
    :param dashboard: dashboard dictionary with a title
    :param folder_id: optional folder ID to search for dashboards in
    :return: dashboard definition if found, None otherwise
    """
    try:
        title = dashboard['title']
        dashboards = list_dashboards(request, title, folder_id)
        if dashboards and isinstance(dashboards, list):
            if len(dashboards) >= 1:
                for dash in dashboards:
                    if dash['title'] == dashboard['title']:
                        definition = _get_dashboard(request, dash['uid'])
                        return definition
        return None
    except HTTPError:
        return None


def _get_dashboard(request: TokenRequest, uid):
    """
    Fetches the dashboard with supplied UID for the token's organization.

    :param request: TokenRequest object
    :param uid: dashboard UID
    :return: dashboard definition if found, None otherwise
    """

    try:
        r = request.get(f'api/dashboards/uid/{uid}')
    except HTTPError:
        return None
    return r.json()['dashboard']


def create_dashboard(request: TokenRequest, dashboard: dict, folder_id=None):
    """
    Creates the given dashboard for the organization tied to the token.
    If the dashboard already exists, it will be updated.

    :param request: TokenRequest object
    :param dashboard: dashboard dictionary
    :param folder_id: optional folder ID to search for the dashboard in
    :return: dashboard definition if dashboard was created, None otherwise
    """

    title = dashboard['title']
    existing_dashboard = None
    has_uid = dashboard.get('uid') is not None
    if has_uid:
        existing_dashboard = _get_dashboard(request, uid=dashboard['uid'])

    # The title might not match the one that's provisioned with that UID.
    # Try to find it by searching for the title instead.
    if existing_dashboard is not None:
        grafana_title = existing_dashboard['title']
        different = grafana_title != title
    else:
        different = False

    if existing_dashboard is None or different:
        existing_dashboard = _search_dashboard(request, dashboard, folder_id)

    if existing_dashboard:
        dashboard['uid'] = existing_dashboard['uid']
        dashboard['id'] = existing_dashboard['id']
        dashboard['version'] = existing_dashboard['version']
    else:
        # We are creating a new dashboard, delete ID if it exists.
        dashboard.pop('id', None)

    payload = {
        'dashboard': dashboard,
        'overwrite': False
    }
    if folder_id:
        payload['folderId'] = folder_id

    # retry up to NUM_RETRIES times
    for _ in range(NUM_RETRIES):
        try:
            r = request.post('api/dashboards/db', json=payload)
            return r.json()
        except HTTPError as e:
            message = ''
            if e.response is not None:
                # log the error message from Grafana
                try:
                    message = e.response.json()
                except json.JSONDecodeError:
                    message = e.response.text
            logger.exception(f"Error when provisioning dashboard {title}: {message}")

            # only retry on server side errors
            if e.response is not None and e.response.status_code < 500:
                break

        time.sleep(1)  # sleep for 1 second before retrying
    return None