Rework cmd-prune

* Move prune functions is cosalib.prune
* Run bump-timestamp command during build
* Create a new subcommand to prune local builds
* Create different subcommand for s3 bucket pruning

Author: vrutkovs

src/bump-timestamp

@@ -0,0 +1,74 @@
+#!/usr/bin/python3 -u
+
+'''
+    This script is called by `build` command and updates builds timestamp.
+'''
+
+import argparse
+import collections
+import os
+import sys
+
+from cosalib.builds import Builds
+from cosalib.cmdlib import get_timestamp
+
+Build = collections.namedtuple('Build', ['id', 'timestamp', 'basearches'])
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--workdir", default='.', help="Path to workdir")
+args = parser.parse_args()
+
+builds = Builds(args.workdir)
+
+scanned_builds = []
+builds_dir = os.path.join(args.workdir, "builds")
+
+# first, pick up all the builds from the dir itself
+with os.scandir(builds_dir) as it:
+    for entry in it:
+        # ignore non-dirs
+        if not entry.is_dir(follow_symlinks=False):
+            # those are really the only two non-dir things we expect there
+            if entry.name not in ['builds.json', 'latest']:
+                print(f"Ignoring non-directory {entry.path}")
+            continue
+
+        # scan all per-arch builds, pick up the most recent build of those as
+        # the overall "build" timestamp for pruning purposes
+        with os.scandir(entry.path) as basearch_it:
+            multiarch_build = None
+            for basearch_entry in basearch_it:
+                # ignore non-dirs
+                if not basearch_entry.is_dir(follow_symlinks=False):
+                    print(f"Ignoring non-directory {basearch_entry.path}")
+                    continue
+                ts = get_timestamp(basearch_entry)
+                if not ts:
+                    continue
+                if not multiarch_build:
+                    multiarch_build = Build(id=entry.name, timestamp=ts,
+                                            basearches=[basearch_entry.name])
+                else:
+                    multiarch_build.basearches += [basearch_entry.name]
+                    multiarch_build.timestamp = max(
+                        multiarch_build.timestamp, ts)
+            if multiarch_build:
+                scanned_builds.append(multiarch_build)
+
+# just get the trivial case out of the way
+if len(scanned_builds) == 0:
+    print("No builds found!")
+    sys.exit(0)
+
+# sort by timestamp, newest first
+scanned_builds = sorted(scanned_builds,
+                        key=lambda x: x.timestamp,
+                        reverse=True)
+
+builds.raw()['builds'] = []
+for build in reversed(scanned_builds):
+    for basearch in build.basearches:
+        builds.insert_build(build.id, basearch)
+
+builds.bump_timestamp()
+print("Build timestamp was updated")

src/cmd-build

@@ -25,7 +25,6 @@ EOF
 # Parse options
 FORCE=
 FORCE_IMAGE=
-SKIP_PRUNE=0
 VERSION=
 PARENT=
 rc=0
@@ -47,9 +46,6 @@ while true; do
         --force-image)
             FORCE_IMAGE=1
             ;;
-        --skip-prune)
-            SKIP_PRUNE=1
-            ;;
         --version)
             shift
             VERSION=$1
@@ -403,15 +399,7 @@ mkdir -p "${builddir}"
 mv -T "${tmp_builddir}" "${builddir}"
 # Replace the latest link
 ln -Tsf "${buildid}" builds/latest
-# Update builds.json
-# the variables passed to `prune_builds` end up single quoted and
-# python treats them as literals, so we workaround this by duplicating
-# the command  ¯\_(ツ)_/¯
-if [ "${SKIP_PRUNE}" == 1 ]; then
-    "${dn}"/prune_builds --workdir "${workdir}" --insert-only "${buildid}"
-else
-    "${dn}"/prune_builds --workdir "${workdir}"
-fi
+"${dn}"/bump-timestamp --workdir "${workdir}"
 rm builds/.build-commit
 
 # and finally, build the specified targets

src/cmd-cleanup-bucket

@@ -0,0 +1,49 @@
+#!/usr/bin/python3 -u
+
+'''
+    This script removes unreferenced builds from s3 bucket
+'''
+
+import argparse
+import sys
+
+from cosalib.builds import Builds
+from cosalib.prune import fetch_build_meta, get_unreferenced_s3_builds, delete_build
+
+parser = argparse.ArgumentParser(prog="coreos-assembler cleanup-bucket")
+parser.add_argument("--workdir", default='.', help="Path to workdir")
+parser.add_argument("--dry-run", help="Don't actually delete anything",
+                    action='store_true')
+parser.add_argument("--bucket", help="S3 bucket")
+parser.add_argument("--prefix", help="S3 prefix")
+args = parser.parse_args()
+
+builds = Builds(args.workdir)
+
+scanned_builds = []
+for build in builds.raw()["builds"]:
+    for arch in build['arches']:
+        build = fetch_build_meta(build['id'], arch)
+        if build:
+            scanned_builds.append(build)
+
+new_builds = []
+builds_to_delete = []
+
+# Find unreferenced builds in the bucket and remove them
+buildids = [x['id'] for x in scanned_builds]
+unreferenced_s3_builds = get_unreferenced_s3_builds(buildids, args.bucket, args.prefix)
+
+error_during_pruning = False
+for unmatched_build_id in unreferenced_s3_builds:
+    # TODO: fetch arches from s3
+    build = fetch_build_meta(unmatched_build_id, 'x86_64')
+    if build and not args.dry_run:
+        try:
+            delete_build(build)
+        except Exception as e:
+            error_during_pruning = True
+            print(f"{e}")
+
+if error_during_pruning:
+    sys.exit(1)

src/cmd-prune

@@ -1,78 +1,125 @@
-#!/usr/bin/env bash
-set -euo pipefail
-
-# This is just a thin wrapper around prune_builds. That way we still get the
-# preflight checks to make sure the workdir looks sane.
-
-dn=$(dirname "$0")
-# shellcheck source=src/cmdlib.sh
-. "${dn}"/cmdlib.sh
-
-print_help() {
-    cat 1>&2 <<'EOF'
-Usage: coreos-assembler prune --help
-       coreos-assembler prune [--keep=N] [--keep-last-days=N]
-
-  Delete older untagged build artifacts. By default, only the last 3 untagged
-  builds are kept.  This can be overridden with the `--keep` option.
-EOF
-}
-
-# Parse options
-KEEP_LAST_N=
-KEEP_LAST_DAYS=
-rc=0
-options=$(getopt --options h --longoptions help,keep:,keep-last-days: -- "$@") || rc=$?
-[ $rc -eq 0 ] || {
-    print_help
-    exit 1
-}
-eval set -- "$options"
-while true; do
-    case "$1" in
-        -h | --help)
-            print_help
-            exit 0
-            ;;
-        --keep)
-            shift
-            KEEP_LAST_N="$1"
-            ;;
-        --keep-last-days)
-            shift
-            KEEP_LAST_DAYS="$1"
-            ;;
-        --)
-            shift
-            break
-            ;;
-        *)
-            fatal "$0: unrecognized option: $1"
-            exit 1
-            ;;
-    esac
-    shift
-done
-
-if [ $# -ne 0 ]; then
-    print_help
-    fatal "ERROR: Too many arguments"
-    exit 1
-fi
-
-# just support one of the two for now
-if [ -n "${KEEP_LAST_N:-}" ] && [ -n "${KEEP_LAST_DAYS:-}" ]; then
-    fatal "ERROR: Only one of --keep or --keep-last-days allowed"
-elif [ -z "${KEEP_LAST_N:-}" ] && [ -z "${KEEP_LAST_DAYS:-}" ]; then
-    KEEP_LAST_N=3
-fi
-
-if [ -n "${KEEP_LAST_DAYS:-}" ]; then
-    set -- --keep-last-days "${KEEP_LAST_DAYS}"
-else
-    set -- --keep-last-n "${KEEP_LAST_N}"
-fi
-
-prepare_build
-
-"${dn}"/prune_builds --workdir "${workdir:?}" "$@"
+#!/usr/bin/python3 -u
+
+'''
+    This script removes previous builds. DO NOT USE on production pipelines
+'''
+
+import os
+import sys
+import argparse
+
+
+from datetime import timedelta, datetime, timezone
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from cosalib.builds import Builds
+from cosalib.prune import fetch_build_meta, delete_build
+
+# Let's just hardcode this here for now
+DEFAULT_KEEP_LAST_N = 3
+DEFAULT_KEEP_LAST_DAYS = 7
+
+
+parser = argparse.ArgumentParser(prog="coreos-assembler prune")
+parser.add_argument("--workdir", default='.', help="Path to workdir")
+parser.add_argument("--dry-run", help="Don't actually delete anything",
+                    action='store_true')
+parser.add_argument("--bucket", help="S3 bucket")
+parser.add_argument("--prefix", help="S3 prefix")
+parser.add_argument("--insert-only", metavar="BUILDID", action='store',
+                    help="Append a new latest build, do not prune")
+keep_options = parser.add_mutually_exclusive_group()
+keep_options.add_argument("--keep-last-n", type=int, metavar="N",
+                          default=DEFAULT_KEEP_LAST_N,
+                          help="Number of untagged builds to keep (0 for all)")
+keep_options.add_argument("--keep-last-days", metavar="N", type=int,
+                          default=DEFAULT_KEEP_LAST_DAYS,
+                          help="Keep untagged builds within number of days")
+args = parser.parse_args()
+
+keep_younger_than = None
+if args.keep_last_days is not None:
+    if args.keep_last_days <= 0:
+        raise argparse.ArgumentTypeError("value must be positive: %d" %
+                                         args.keep_last_days)
+    keep_younger_than = (datetime.now(timezone.utc) -
+                         timedelta(days=args.keep_last_days))
+
+skip_pruning = (not keep_younger_than and args.keep_last_n == 0)
+print("prune: skip_pruning: {skip_pruning}")
+
+builds = Builds(args.workdir)
+# collect all builds being pointed to by tags
+tagged_builds = set([tag['target'] for tag in builds.raw().get('tags', [])])
+
+# Handle --insert-only
+if args.insert_only:
+    builds.insert_build(args.insert_only)
+    builds.flush()
+    print("prune: --insert-only completed")
+    sys.exit(0)
+
+scanned_builds = []
+for build in builds.raw()["builds"]:
+    for arch in build['arches']:
+        build = fetch_build_meta(build['id'], arch)
+        if build:
+            scanned_builds.append(build)
+
+new_builds = []
+builds_to_delete = []
+
+# Don't prune known builds
+if skip_pruning:
+    new_builds = scanned_builds
+else:
+    if keep_younger_than:
+        for build in scanned_builds:
+            if build.id in tagged_builds:
+                print(f"Skipping tagged build {build.id}")
+                new_builds.append(build)
+                continue
+
+            if build.timestamp < keep_younger_than:
+                builds_to_delete.append(build)
+            else:
+                new_builds.append(build)
+    else:
+        n = args.keep_last_n
+        assert(n > 0)
+        for build in scanned_builds:
+            if n == 0:
+                builds_to_delete.append(build)
+            else:
+                new_builds.append(build)
+                n = n - 1
+
+print(f"prune: new builds: {new_builds}")
+
+# create a new builds list
+builds.raw()['builds'] = []
+for build in reversed(new_builds):
+    for arch in build['arches']:
+        builds.insert_build(build['id'], arch)
+
+builds.bump_timestamp()
+
+if len(builds_to_delete) == 0:
+    print("prune: not removing any builds")
+else:
+    buildids = [x['id'] for x in builds_to_delete]
+    print(f"prune: removing {' '.join(buildids)}")
+
+# now delete other build dirs not in the manifest
+error_during_pruning = False
+for build in builds_to_delete:
+    print(f"Pruning {build}")
+    if not args.dry_run:
+        try:
+            delete_build(build)
+        except Exception as e:
+            error_during_pruning = True
+            print(f"{e}")
+
+if error_during_pruning:
+    sys.exit(1)

src/coreos-assembler

@@ -39,7 +39,7 @@ build_commands="init fetch build run prune clean"
 # commands more likely to be used in a prod pipeline only
 advanced_build_commands="buildprep buildupload oscontainer"
 buildextend_commands="qemu aws azure gcp openstack installer live vmware metal"
-utility_commands="tag sign compress koji-upload kola aws-replicate"
+utility_commands="tag sign compress koji-upload kola aws-replicate cleanup-bucket"
 other_commands="shell"
 if [ -z "${cmd}" ]; then
     echo Usage: "coreos-assembler CMD ..."