folder.py

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


logger = logging.getLogger(__name__)


def delete_folder(request: TokenRequest, title=None, uid=None):
    """
    Deletes a single folder for the organization
    the API token is registered to.

    Folder can be specified by UID or title.

    :param request: TokenRequest object
    :param title: folder title
    :param uid: folder UID
    :return: True if folder is considered deleted, False otherwise
    """
    if uid:
        try:
            r = request.delete(f'api/folders/{uid}').json()
        except HTTPError:
            logger.exception(f'Error when deleting folder: {uid}')
            return False
        return r is not None
    else:
        folder = find_folder(request, title, False)

        if folder is None:
            return True
        try:
            r = request.delete(f'api/folders/{folder.get("uid")}')
            r = r.json()
        except HTTPError:
            logger.exception(f'Error when deleting folder: {title}')
            return False

        logger.info(f'Deleted folder: {title}')
        return r is not None


def find_folder(request: TokenRequest, title, create=True):
    """
    Finds a folder by title. If create is True, creates the folder if it does
    not exist.

    :param request: TokenRequest object
    :param title: folder title
    :param create: create folder if it does not exist
    :return: folder definition
    """
    folders = get_folders(request)
    try:
        folder = next(
            f for f in folders if f['title'].lower() == title.lower())
    except StopIteration:
        folder = None

    if not folder and create:
        folder = create_folder(request, title)
        if not folder:
            return None
        logger.info(f'Created folder: {title}')
    return folder


def get_folders(request: TokenRequest):
    """
    Returns all folders for the organization
    the API token is registered to.

    :param request: TokenRequest object
    :return: list of folder definitions
    """
    try:
        r = request.get('api/folders')
        folders = r.json()
    except HTTPError:
        logger.exception('Error when getting Grafana folders:')
        return []

    return folders


def create_folder(request: TokenRequest, title):
    """
    Creates a folder for the organization
    the API token is registered to.

    :param request: TokenRequest object
    :param title: folder title
    :return: folder definition
    """
    try:
        data = {'title': title, 'uid': title.replace(' ', '_')}
        r = request.post('api/folders', json=data)
    except HTTPError:
        logger.exception(f'Error when creating folder {title}')
        return None
    return r.json()


def delete_unknown_folders(token, folders_to_keep: set):
    """
    Deletes all folders that are not in the folders_to_keep list.

    :param token: TokenRequest object
    :param folders_to_keep: set of folder titles to keep
    """

    all_folders = get_folders(token)

    for folder in all_folders:
        if folder['title'] in folders_to_keep:
            continue
        logger.info(f'Deleting unknown folder: {folder.get("title")}')
        delete_folder(token, uid=folder['uid'])