image-refresh

#!/usr/bin/env python3

# Copyright (C) 2016-2024 Red Hat, Inc.
# SPDX-License-Identifier: GPL-3.0-or-later

import argparse
import json
import os
import platform
import shlex
import subprocess
import sys
import time

from lib import git, testmap
from lib.constants import BOTS_DIR, SCRIPTS_DIR
from lib.github import NOT_TESTED_DIRECT, GitHub


def run(
    *cmd: str,
    check: bool = True,
    verbose: bool = False,
    dry_run: bool = False,
    env: dict[str, str] | None = None,
) -> int:
    if dry_run:
        print('\n+ #', shlex.join(cmd))
        return 0
    if verbose:
        print('\n+', shlex.join(cmd), file=sys.stderr)
    result = subprocess.run(cmd, check=check, env=env)
    return result.returncode


def image_refresh(
    github: GitHub, image: str, issue_nr: int, *, verbose: bool = False, dry_run: bool = False,
) -> str:
    """Run the image refresh. Returns the branch name."""
    log_url = os.environ.get('COCKPIT_CI_LOG_URL', None)
    github.comment(
        issue_nr,
        f"image-refresh {image} started on {platform.node()}. Log: {log_url or 'unknown'}",
        dry_run=dry_run,
    )

    # Cleanup any extraneous disk usage elsewhere
    run('./vm-reset', verbose=verbose)

    # download the current image, for comparing them; that may not exist yet for newly introduced images
    if run('./image-download', image, check=False, verbose=verbose) == 0:
        old_image = os.path.realpath(f'{BOTS_DIR}/images/{image}')
    else:
        old_image = None

    # create the new image
    verbose_args = ['--verbose'] if verbose else []
    run(
        './image-create',
        *verbose_args,
        image,
        env={**os.environ, 'VIRT_BUILDER_NO_CACHE': "yes"},
        verbose=verbose,
    )

    git.add(f'images/{image}')

    if git.changes_staged():
        # diff against the old image, upload the new one
        diff_text = ''
        if old_image and os.path.exists(f'{SCRIPTS_DIR}/{image}.setup'):
            print(f'+ ./image-diff {old_image} {image}', file=sys.stderr)
            diff_text = subprocess.check_output(['./image-diff', old_image, image], text=True)
            print(diff_text)

        run('./image-upload', '--prune-s3', image, verbose=verbose, dry_run=dry_run)
        git.commit(f'images: Update {image} image\n\n{diff_text}\nCloses #{issue_nr}')
    else:
        git.commit(f'images: Update {image} image (unchanged)\n\nCloses #{issue_nr}', allow_empty=True)
    branch = git.push(github, f"image-refresh-{image}", dry_run=dry_run)

    # Check if the issue is already a pull request (and do nothing if so)
    issue = github.get_obj(f"issues/{issue_nr}", {})
    if "pull_request" in issue:
        return branch

    # ... else, create a PR and trigger tests.
    triggers = testmap.tests_for_image(image)
    pull = github.convert_issue_to_pull(branch, issue_nr, dry_run=dry_run)
    if pull is not None:
        head = git.get_current_head()

        # Create a synthetic status to forward the log URL to the new HEAD
        if log_url:
            github.post(
                f"statuses/{head}",
                {
                    'state': 'success',
                    'context': f'image-refresh/{image}',
                    'description': 'Forwarded status',
                    'target_url': log_url,
                },
            )

        for trigger in triggers:
            github.post(
                f"statuses/{head}",
                {"state": "pending", "context": trigger, "description": NOT_TESTED_DIRECT},
            )
    else:
        print('\n** Would trigger tests:', json.dumps(triggers, indent=4))

    return branch


def main() -> None:
    parser = argparse.ArgumentParser(description="Refresh image")
    parser.add_argument("--verbose", "-v", action="store_true", help="Verbose output")
    parser.add_argument("--issue", required=True, type=int, help="GitHub issue number to act on")
    parser.add_argument("--dry-run", "-n", action="store_true", help="Dry run to validate this task if supported")
    parser.add_argument("image")
    args = parser.parse_args()

    github = GitHub()
    start_time = time.time()
    try:
        branch = image_refresh(github, args.image, args.issue, verbose=args.verbose, dry_run=args.dry_run)
    except BaseException as exc:
        result = f"image-refresh {args.image} failed: {exc}"
        sys.stderr.write(f"\n# {result}\n# Duration: {time.time() - start_time}s\n")
        if not args.dry_run:
            github.comment(args.issue, result)
        raise
    else:
        result = f"image-refresh {args.image} succeeded: https://github.com/{github.repo}/commits/{branch}"
        sys.stderr.write(f"\n# {result}\n# Duration: {time.time() - start_time}s\n")
        if not args.dry_run:
            github.comment(args.issue, result)


if __name__ == '__main__':
    main()