edit-prowjobs.py

#!/usr/bin/env python3

# Copyright 2021 The Kubernetes Authors.
#
# 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.

"""Edit prowjobs en-masse by round-tripping them through ruamel.yaml

This is not intended for general usage, because:
- not all jobs can or should be edited
- many jobs have different formatting, and we're not at a point where
  we can enforce formatting standards, so this is almost guaranteed
  to introduce formatting change noise
- the idea is to manually edit this file with the specific edit to be
  done, rather that developing a general purpose language to do this
"""

import argparse
import glob
import re

from os import path

import ruamel.yaml

# Prow files that will be ignored
EXCLUDED_JOB_CONFIGS = [
    # Ruamel won't be able to successfully dump fejta-bot-periodics
    # See https://bitbucket.org/ruamel/yaml/issues/258/applying-json-patch-breaks-comment
    "fejta-bot-periodics.yaml",
    # generated.yaml is generated by generate_tests.py, and will be overwritten.
    "generated.yaml",
]
# A hilariously large line length to ensure we never line-wrap
MAX_WIDTH = 2000000000



def setup_yaml():
    # Setup the ruamel.yaml parser
    yaml = ruamel.yaml.YAML(typ='rt')
    yaml.preserve_quotes = True
    # GoogleCloudPlatform/ - no yaml.indent
    # bazelbuild/ - no yaml.indent
    # cadvisor/ - no yaml.indent
    # containerd/ - no yaml.indent
    # image-pushing/ - yaml.indent(mapping=2, sequence=4, offset=2)
    # kubernetes/ - yaml.indent(mapping=2) seems to cause the least change
    # kubernetes-client - TBD
    # kubernetes-csi - TBD
    # kubernetes-sigs - TBD
    yaml.indent(mapping=2, sequence=4, offset=2)
    yaml.width = MAX_WIDTH
    return yaml

def edit_job_config(yaml, prow_job_file_name, force_rewrite=False):
    with open(prow_job_file_name, "r") as job_fp:
        prow_config = yaml.load(job_fp)

    def edit(job):
        edited = False
        name = job["name"]
        print(f'  handling job: {name}')
        annotations = job["annotations"]
        dashboard_list = re.split('[, ]+', annotations["testgrid-dashboards"])
        if 'sig-k8s-infra-gcb' not in dashboard_list:
            dashboard_list.append('sig-k8s-infra-gcb')
            annotations["testgrid-dashboards"] = ", ".join(dashboard_list)
            edited = True
        return edited

    should_rewrite = force_rewrite

    # For each presubmit, postsubmit, and periodic
    # presubmits -> <any repository> -> [{name: prowjob}]
    if "presubmits" in prow_config:
        for _, jobs in prow_config["presubmits"].items():
            for job in jobs:
                if edit(job):
                    should_rewrite = True

    # postsubmits -> <any repository> -> [{name: prowjob}]
    if "postsubmits" in prow_config:
        for _, jobs in prow_config["postsubmits"].items():
            for job in jobs:
                if edit(job):
                    should_rewrite = True

    # periodics -> [{name: prowjob}]
    if "periodics" in prow_config:
        for job in prow_config["periodics"]:
            if edit(job):
                should_rewrite = True

    # Dump ProwConfig to prowJobFile
    if should_rewrite:
        print(f'  writing {prow_job_file_name}')
        with open(prow_job_file_name, "w") as job_fp:
            yaml.dump(prow_config, job_fp)
            job_fp.truncate()

def main(prow_job_dir, force_rewrite):
    yaml = setup_yaml()
    for f in glob.glob(f'{prow_job_dir}/**/*.yaml', recursive=True):
        if path.basename(f) not in EXCLUDED_JOB_CONFIGS:
            try:
                print(f'processing config: {f}')
                edit_job_config(yaml, f, force_rewrite)
            except Exception as e:  # pylint: disable=broad-except
                print(f'ERROR: could not edit {f}: {e}')

if __name__ == '__main__':
    PARSER = argparse.ArgumentParser(
        description='Does things to prowjob configs')
    PARSER.add_argument(
        '--prow-job-dir',
        default='../config/jobs',
        help='Path to Prow Job Directory')
    PARSER.add_argument(
        '--force',
        default=True,
        help='Force rewrite of all job configs')
    ARGS = PARSER.parse_args()

    main(ARGS.prow_job_dir, ARGS.force)