truenas/update_dependencies

#!/usr/bin/env python3
import argparse
import errno
import os
import subprocess

from collections import defaultdict

ITEMS = defaultdict(lambda: {'success': [], 'error': []})


class ValidationException(Exception):
    def __init__(self, error_msg, error_no=errno.EFAULT):
        self.errmsg = error_msg
        self.errno = error_no

    def get_error_name(self):
        return errno.errorcode.get(self.errno) or 'EUNKNOWN'

    def __str__(self):
        return f'[{self.get_error_name()}] {self.errmsg}'


class NotFoundException(ValidationException):
    def __init__(self, error):
        super().__init__(error, errno.ENOENT)


class TrainNotFoundException(NotFoundException):
    def __init__(self):
        super(TrainNotFoundException, self).__init__('Failed to find train')


class CatalogItemNotFoundException(NotFoundException):
    def __init__(self, path):
        super(CatalogItemNotFoundException, self).__init__(f'Failed to find {path!r} catalog item')


def report_result():
    print('[\033[94mINFO\x1B[0m]\tExecution Complete')
    for index, item in enumerate(ITEMS):
        index += 1
        data = ITEMS[item]
        print(f'\n[\033[94mINFO\x1B[0m]\t{index}) {item}')
        if data['success']:
            print(
                f'[\033[92mOK\x1B[0m]\t - Successfully updated dependencies for {", ".join(data["success"])} versions'
            )

        for i_v, version in enumerate(data['error']):
            v_name, error = version
            print(
                f'[\033[91mFAILED\x1B[0m]\t ({i_v + 1}) Failed to update dependencies for {v_name!r} version: {error}'
            )


def update_train_charts(train_path, commit):
    # We will gather all charts in the train and then for each chart all it's versions will be updated
    if not os.path.exists(train_path):
        raise TrainNotFoundException()

    print(f'[\033[94mINFO\x1B[0m]\tProcessing {train_path!r} train')
    for item in os.listdir(train_path):
        process_catalog_item(os.path.join(train_path, item))

    report_result()

    if commit and any(ITEMS[item]['success'] for item in ITEMS):
        if any(ITEMS[item]['error'] for item in ITEMS):
            print(f'[\033[91mFAILED\x1B[0m]\tNot committing changes as failures detected')
        else:
            commit_msg = f'Updated catalog item dependencies ({train_path.rsplit("/", 1)[-1]} train)\n' \
                         'Following items were updated:\n'
            for item in ITEMS:
                commit_msg += f'Updated {item} ({", ".join(ITEMS[item]["success"])} versions)\n\n'

            for cmd in (
                ['git', '-C', train_path, 'add', train_path],
                ['git', '-C', train_path, 'commit', '-m', commit_msg]
            ):
                cp = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
                stderr = cp.communicate()[1]
                if cp.returncode:
                    print(f'[\033[91mFAILED\x1B[0m]\tFailed to execute {" ".join(cmd)}: {stderr.decode()}')
                    exit(1)

            print('[\033[92mOK\x1B[0m]\tChanges committed successfully')


def process_catalog_item(item_path):
    if not os.path.exists(item_path):
        raise CatalogItemNotFoundException(item_path)

    item_name = item_path.rsplit('/', 1)[-1]
    print(f'[\033[94mINFO\x1B[0m]\tProcessing {item_name!r} catalog item')
    for item_version in os.listdir(item_path):
        if os.path.isdir(os.path.join(item_path, item_version)):
            update_item_version(item_name, item_version, os.path.join(item_path, item_version))


def update_item_version(item_name, version, version_path):
    cp = subprocess.Popen(
        ['helm', 'dependency', 'update', version_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    stdout, stderr = cp.communicate()
    if cp.returncode:
        ITEMS[item_name]['error'].append((version, stderr.decode()))
    else:
        ITEMS[item_name]['success'].append(version)


def main():
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(help='sub-command help', dest='action')

    parser_setup = subparsers.add_parser('update', help='Update dependencies for specified train')
    parser_setup.add_argument('--train', help='Specify train path to update dependencies', required=True)
    parser_setup.add_argument('--commit', help='Commit after updating dependencies', default=False)

    args = parser.parse_args()
    if args.action == 'update':
        update_train_charts(args.train, args.commit)
    else:
        parser.print_help()


if __name__ == '__main__':
    main()