From ec320f3b5c498d1941b7804928fbd53ab2e02c81 Mon Sep 17 00:00:00 2001
From: Jonathan Lebon <jonathan@jlebon.com>
Date: Tue, 2 Jul 2019 16:11:11 -0400
Subject: [PATCH] Use new builds/ dir helpers

Prep for switching those over to the new build layout.

Most commands have pretty straightforward changes. Some of the more
complex ones that deal with the build layout a lot were refactored a bit
as prep for multi-arch.
---
 src/cmd-build                 |  21 +++---
 src/cmd-buildextend-aws       |   6 +-
 src/cmd-buildextend-gcp       |   6 +-
 src/cmd-buildextend-installer |  10 +--
 src/cmd-buildextend-metal     |   7 +-
 src/cmd-buildextend-openstack |  10 +--
 src/cmd-buildextend-vmware    |   9 +--
 src/cmd-buildprep             |  24 ++++---
 src/cmd-buildupload           |  23 +++----
 src/cmd-compress              | 123 ++++++++++++++++++----------------
 src/cmd-tag                   |  48 +++----------
 src/cosalib/build.py          |  12 ++--
 src/prune_builds              |  78 ++++++++++-----------
 13 files changed, 188 insertions(+), 189 deletions(-)

diff --git a/src/cmd-build b/src/cmd-build
index d831e9c55c..4acf5a04dc 100755
--- a/src/cmd-build
+++ b/src/cmd-build
@@ -80,10 +80,9 @@ prepare_build
 ostree --version
 rpm-ostree --version
 
-previous_build=
-if [ -L "${workdir}"/builds/latest ]; then
-    previous_build=$(readlink "${workdir}"/builds/latest)
-    previous_builddir="${workdir}/builds/${previous_build}"
+previous_build=$(get_latest_build)
+if [ -n "${previous_build}" ]; then
+    previous_builddir=$(get_build_dir "${previous_build}")
     echo "Previous build: ${previous_build}"
 fi
 
@@ -360,18 +359,20 @@ saved_build_tmpdir="${workdir}/tmp/last-build-tmp"
 rm -rf "${saved_build_tmpdir}"
 mv -T tmp "${saved_build_tmpdir}"
 ostree prune --repo="${tmprepo}" --refs-only
-# Back to the toplevel build directory, so we can rename this one
-cd "${workdir}"/builds
+# Back to the toplevel work directory, so we can rename this one
+cd "${workdir}"
 # We create a .build-commit file to note that we're in the
 # middle of a "commit".  This may be useful in the future
 # for having things be transactional.  If for example we
 # were interrupted between the rename() and linkat() below,
 # things would be inconsistent and future builds would fail
 # on the `mv`.
-touch .build-commit
-mv -T "${tmp_builddir}" "${buildid}"
+touch builds/.build-commit
+builddir=$(get_build_dir "${buildid}")
+mkdir -p "${builddir}"
+mv -T "${tmp_builddir}" "${builddir}"
 # Replace the latest link
-ln -Tsfr "${buildid}" latest
+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
@@ -381,4 +382,4 @@ if [ "${SKIP_PRUNE}" == 1 ]; then
 else
     "${dn}"/prune_builds --workdir "${workdir}"
 fi
-rm .build-commit
+rm builds/.build-commit
diff --git a/src/cmd-buildextend-aws b/src/cmd-buildextend-aws
index b48db744ad..2f32c59289 100755
--- a/src/cmd-buildextend-aws
+++ b/src/cmd-buildextend-aws
@@ -8,7 +8,7 @@ import os,sys,json,yaml,shutil,argparse,subprocess,re,collections
 import tempfile,hashlib,gzip
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import run_verbose, write_json, sha256sum_file
+from cmdlib import run_verbose, write_json, sha256sum_file, Builds
 
 # Parse args and dispatch
 parser = argparse.ArgumentParser()
@@ -23,8 +23,8 @@ parser.add_argument("--grant-user", help="Grant user launch permission",
                     nargs="*", default=[])
 args = parser.parse_args()
 
-
-builddir = os.path.join('builds', args.build)
+builds = Builds()
+builddir = builds.get_build_dir(args.build)
 buildmeta_path = os.path.join(builddir, 'meta.json')
 with open(buildmeta_path) as f:
     buildmeta = json.load(f)
diff --git a/src/cmd-buildextend-gcp b/src/cmd-buildextend-gcp
index 0135a790ec..276affc03a 100755
--- a/src/cmd-buildextend-gcp
+++ b/src/cmd-buildextend-gcp
@@ -18,7 +18,8 @@ sys.path.insert(0, cosa_dir)
 from cmdlib import (
         run_verbose,
         sha256sum_file,
-        write_json)
+        write_json,
+        Builds)
 
 # Parse args and dispatch
 parser = argparse.ArgumentParser()
@@ -46,7 +47,8 @@ if args.project is None:
     raise Exception(arg_exp_str.format("project", "GCP_PROJECT"))
 
 # Identify the builds
-builddir = os.path.join('builds', args.build)
+builds = Builds()
+builddir = builds.get_build_dir(args.build)
 buildmeta_path = os.path.join(builddir, 'meta.json')
 with open(buildmeta_path) as f:
     buildmeta = json.load(f)
diff --git a/src/cmd-buildextend-installer b/src/cmd-buildextend-installer
index daa62de973..fdbe1958c3 100755
--- a/src/cmd-buildextend-installer
+++ b/src/cmd-buildextend-installer
@@ -14,7 +14,7 @@ import tempfile
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 from cmdlib import run_verbose, write_json, sha256sum_file
-from cmdlib import import_ostree_commit
+from cmdlib import import_ostree_commit, Builds
 
 # Parse args and dispatch
 parser = argparse.ArgumentParser()
@@ -23,16 +23,16 @@ parser.add_argument("--force", action='store_true', default=False,
                     help="Overwrite previously generated installer")
 args = parser.parse_args()
 
+builds = Builds()
+
 # default to latest build if not specified
 if not args.build:
-    with open('builds/builds.json') as f:
-        j = json.load(f)
-        args.build = j['builds'][0]
+    args.build = builds.get_latest()
 
 print(f"Targeting build: {args.build}")
 
 workdir = os.path.abspath(os.getcwd())
-builddir = os.path.join(workdir, 'builds', args.build)
+builddir = builds.get_build_dir(args.build)
 buildmeta_path = os.path.join(builddir, 'meta.json')
 with open(buildmeta_path) as f:
     buildmeta = json.load(f)
diff --git a/src/cmd-buildextend-metal b/src/cmd-buildextend-metal
index f5d42e42a7..c4a735ad99 100755
--- a/src/cmd-buildextend-metal
+++ b/src/cmd-buildextend-metal
@@ -68,14 +68,13 @@ export LIBGUESTFS_BACKEND=direct
 prepare_build
 
 if [ -z "${build}" ]; then
-    if [ -L "${workdir}"/builds/latest ]; then
-        build=$(readlink "${workdir}"/builds/latest)
-    else
+    build=$(get_latest_build)
+    if [ -z "${build}" ]; then
         fatal "No build found."
     fi
 fi
 
-builddir="${workdir}/builds/${build}"
+builddir=$(get_build_dir "$build")
 if [ ! -d "${builddir}" ]; then
     fatal "Build dir ${builddir} does not exist."
 fi
diff --git a/src/cmd-buildextend-openstack b/src/cmd-buildextend-openstack
index 6bb177eb72..1c80dede01 100755
--- a/src/cmd-buildextend-openstack
+++ b/src/cmd-buildextend-openstack
@@ -11,22 +11,22 @@ import shutil
 import argparse
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import run_verbose, write_json, sha256sum_file
+from cmdlib import run_verbose, write_json, sha256sum_file, Builds
 
 # Parse args and dispatch
 parser = argparse.ArgumentParser()
 parser.add_argument("--build", help="Build ID")
 args = parser.parse_args()
 
+builds = Builds()
+
 # default to latest build if not specified
 if not args.build:
-    with open('builds/builds.json') as f:
-        j = json.load(f)
-        args.build = j['builds'][0]
+    args.build = builds.get_latest()
 
 print(f"Targeting build: {args.build}")
 
-builddir = os.path.join('builds', args.build)
+builddir = builds.get_build_dir(args.build)
 buildmeta_path = os.path.join(builddir, 'meta.json')
 with open(buildmeta_path) as f:
     buildmeta = json.load(f)
diff --git a/src/cmd-buildextend-vmware b/src/cmd-buildextend-vmware
index 3e3e8c4420..d0e6edfd92 100755
--- a/src/cmd-buildextend-vmware
+++ b/src/cmd-buildextend-vmware
@@ -11,19 +11,20 @@ import argparse
 import tarfile
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import run_verbose, load_json, write_json, sha256sum_file
+from cmdlib import run_verbose, load_json, write_json, sha256sum_file, Builds
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--build", help="Build ID",
                     required=False)
 args = parser.parse_args()
 
+builds = Builds()
+
 # default to latest build if not specified
 if not args.build:
-    j = load_json('builds/builds.json')
-    args.build = j['builds'][0]
+    args.build = builds.get_latest()
 
-builddir = os.path.join('builds', args.build)
+builddir = builds.get_build_dir(args.build)
 buildmeta_path = os.path.join(builddir, 'meta.json')
 buildmeta = load_json(buildmeta_path)
 base_name = buildmeta['name']
diff --git a/src/cmd-buildprep b/src/cmd-buildprep
index fa06ba02ab..8606576638 100755
--- a/src/cmd-buildprep
+++ b/src/cmd-buildprep
@@ -14,7 +14,7 @@ import boto3
 import shutil
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import load_json, rm_allow_noent  # noqa: E402
+from cmdlib import load_json, rm_allow_noent, Builds  # noqa: E402
 
 
 def main():
@@ -29,18 +29,26 @@ def main():
     else:
         raise Exception("Invalid scheme: only file://, s3://, and http(s):// supported")
 
-    builds = []
+    builds = None
     if fetcher.exists('builds.json'):
-        builds = fetcher.fetch_json('builds.json')['builds']
+        fetcher.fetch_json('builds.json')
+        builds = Builds()
 
-    if len(builds) == 0:
+    if not builds or builds.is_empty():
         print("Remote has no builds!")
         return
 
-    buildid = builds[0]
-    os.makedirs(f'builds/{buildid}', exist_ok=True)
+
+    buildid = builds.get_latest()
+    builddir = builds.get_build_dir(buildid)
+    os.makedirs(builddir, exist_ok=True)
+
+    # trim out the leading builds/
+    assert builddir.startswith("builds/")
+    builddir = builddir[len("builds/"):]
+
     for f in ['meta.json', 'ostree-commit-object']:
-        fetcher.fetch(f'{buildid}/{f}')
+        fetcher.fetch(f'{builddir}/{f}')
 
     # and finally the symlink
     rm_allow_noent('builds/latest')
@@ -48,7 +56,7 @@ def main():
 
     # also nuke the any local matching OStree ref, since we want to build on
     # top of this new one
-    buildmeta = load_json('builds/latest/meta.json')
+    buildmeta = load_json(f'builds/{builddir}/meta.json')
     if 'ref' in buildmeta and os.path.isdir('tmp/repo'):
         subprocess.check_call(['ostree', 'refs', '--repo', 'tmp/repo',
                                '--delete', buildmeta['ref']],
diff --git a/src/cmd-buildupload b/src/cmd-buildupload
index 01f85cc0df..da25600f00 100755
--- a/src/cmd-buildupload
+++ b/src/cmd-buildupload
@@ -11,7 +11,7 @@ import sys
 import tempfile
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import load_json  # noqa: E402
+from cmdlib import load_json, Builds  # noqa: E402
 
 
 def main():
@@ -44,18 +44,17 @@ def parse_args():
 
 def cmd_upload_s3(args):
     if not args.freshen:
-        s3_upload_build(args)
+        builds = Builds()
+        if args.build == 'latest':
+            args.build = builds.get_latest()
+        print(f"Targeting build: {args.build}")
+        s3_upload_build(args, builds.get_build_dir(args.build), args.build)
     s3_cp(args, 'builds/builds.json', 'builds.json',
           '--cache-control=max-age=60')
 
 
-def s3_upload_build(args):
-    builddir = f'builds/{args.build}'
+def s3_upload_build(args, builddir, dest):
     build = load_json(f'{builddir}/meta.json')
-    buildid = build['buildid']
-
-    print(f"Targeting build: {buildid}")
-    print(f"  OSTree commit: {build['ostree-commit']}")
 
     # Upload images with special handling for gzipped data.
     uploaded = set()
@@ -71,11 +70,11 @@ def s3_upload_build(args):
                 args.enable_gz_peel):
             nogz = bn[:-3]
             img['path'] = nogz
-            s3_cp(args, path, f'{buildid}/{nogz}',
+            s3_cp(args, path, f'{dest}/{nogz}',
                   '--content-encoding=gzip',
                   f'--content-disposition=inline; filename={nogz}')
         else:
-            s3_cp(args, path, f'{buildid}/{bn}')
+            s3_cp(args, path, f'{dest}/{bn}')
         uploaded.add(bn)
 
     for f in os.listdir(builddir):
@@ -83,7 +82,7 @@ def s3_upload_build(args):
         if f in uploaded or f == 'meta.json':
             continue
         path = os.path.join(builddir, f)
-        s3_cp(args, path, f'{buildid}/{f}')
+        s3_cp(args, path, f'{dest}/{f}')
 
     # Now upload a modified version of the meta.json which has the fixed
     # filenames without the .gz suffixes. We don't want to modify the local
@@ -91,7 +90,7 @@ def s3_upload_build(args):
     with tempfile.NamedTemporaryFile('w') as f:
         json.dump(build, f, indent=4)
         f.flush()
-        s3_cp(args, f.name, f'{buildid}/meta.json',
+        s3_cp(args, f.name, f'{dest}/meta.json',
               '--content-type=application/json')
 
 
diff --git a/src/cmd-compress b/src/cmd-compress
index 5063e417dc..6f6afcb4c7 100755
--- a/src/cmd-compress
+++ b/src/cmd-compress
@@ -10,76 +10,87 @@ import shutil
 import argparse
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import run_verbose, write_json, sha256sum_file, rm_allow_noent
+from cmdlib import run_verbose, write_json, sha256sum_file
+from cmdlib import rm_allow_noent, Builds
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--build", help="Build ID")
 args = parser.parse_args()
 
+builds = Builds()
+
 # default to latest build if not specified
 if args.build:
     build = args.build
 else:
-    with open('builds/builds.json') as f:
-        j = json.load(f)
-        build = j['builds'][0]
+    build = builds.get_latest()
 
 print(f"Targeting build: {build}")
 
-buildmeta_path = os.path.join('builds', build, 'meta.json')
-with open(buildmeta_path) as f:
-    buildmeta = json.load(f)
+# Don't compress certain images
+imgs_to_skip = ["iso", "vmware", "initramfs", "kernel"]
 
-tmpdir = 'tmp/compress'
-if os.path.isdir(tmpdir):
-    shutil.rmtree(tmpdir)
-os.mkdir(tmpdir)
 
-# Note we mutate the build dir in place, similarly to the buildextend commands.
-# One cool approach here might be to `cp -al` the whole build dir, mutate it,
-# then RENAME_EXCHANGE the two... though it doesn't seem Python exposes it yet
-# so that'd require some hacking around. For now, we just guarantee that
-# `compress` is idempotent and can resume from failures.
+def compress_one_builddir(builddir):
+    print(f"Compressing: {builddir}")
+    buildmeta_path = os.path.join(builddir, 'meta.json')
+    with open(buildmeta_path) as f:
+        buildmeta = json.load(f)
 
-at_least_one = False
-imgs_to_compress = []
-imgs_to_skip = ["iso", "vmware", "initramfs", "kernel"]
+    tmpdir = 'tmp/compress'
+    if os.path.isdir(tmpdir):
+        shutil.rmtree(tmpdir)
+    os.mkdir(tmpdir)
 
-for img_format in buildmeta['images']:
-
-    # Don't compress certain images
-    if img_format in imgs_to_skip:
-        print(f"Skipping {img_format}")
-        continue
-
-    img = buildmeta['images'][img_format]
-    file = img['path']
-    filepath = os.path.join('builds', build, file)
-    if not file.endswith('.gz'):
-        tmpfile = os.path.join(tmpdir, (file + '.gz'))
-        # SHA256 + size for uncompressed image was already calculated during 'build'
-        img['uncompressed-sha256'] = img['sha256']
-        img['uncompressed-size'] = img['size']
-        with open(tmpfile, 'wb') as f:
-            run_verbose(['gzip', '-c', filepath], stdout=f)
-        file_gz = file + '.gz'
-        filepath_gz = filepath + '.gz'
-        compressed_size = os.path.getsize(tmpfile)
-        img['path'] = file_gz
-        img['sha256'] = sha256sum_file(tmpfile)
-        img['size'] = compressed_size
-
-        # just flush out after every image type, but unlink after writing.
-        # Then, we should be able to interrupt and restart from the last type.
-        os.rename(tmpfile, filepath_gz)
-        write_json(buildmeta_path, buildmeta)
-        os.unlink(filepath)
-        at_least_one = True
-    else:
-        # try to delete the original file if it's somehow still around
-        rm_allow_noent(filepath[:-3])
-
-if at_least_one:
-    print(f"Updated: {buildmeta_path}")
-else:
+    # Note we mutate the build dir in place, similarly to the buildextend
+    # commands.  One cool approach here might be to `cp -al` the whole build
+    # dir, mutate it, then RENAME_EXCHANGE the two... though it doesn't seem
+    # Python exposes it yet so that'd require some hacking around. For now, we
+    # just guarantee that `compress` is idempotent and can resume from
+    # failures.
+
+    at_least_one = False
+
+    for img_format in buildmeta['images']:
+        if img_format in imgs_to_skip:
+            print(f"Skipping {img_format}")
+            continue
+
+        img = buildmeta['images'][img_format]
+        file = img['path']
+        filepath = os.path.join(builddir, file)
+        if not file.endswith('.gz'):
+            tmpfile = os.path.join(tmpdir, (file + '.gz'))
+            # SHA256 + size for uncompressed image was already calculated
+            # during 'build'
+            img['uncompressed-sha256'] = img['sha256']
+            img['uncompressed-size'] = img['size']
+            with open(tmpfile, 'wb') as f:
+                run_verbose(['gzip', '-c', filepath], stdout=f)
+            file_gz = file + '.gz'
+            filepath_gz = filepath + '.gz'
+            compressed_size = os.path.getsize(tmpfile)
+            img['path'] = file_gz
+            img['sha256'] = sha256sum_file(tmpfile)
+            img['size'] = compressed_size
+
+            # Just flush out after every image type, but unlink after writing.
+            # Then, we should be able to interrupt and restart from the last
+            # type.
+            os.rename(tmpfile, filepath_gz)
+            write_json(buildmeta_path, buildmeta)
+            os.unlink(filepath)
+            at_least_one = True
+        else:
+            # try to delete the original file if it's somehow still around
+            rm_allow_noent(filepath[:-3])
+
+    if at_least_one:
+        print(f"Updated: {buildmeta_path}")
+    return at_least_one
+
+
+changed = compress_one_builddir(builds.get_build_dir(build))
+
+if not changed:
     print(f"All builds already compressed")
diff --git a/src/cmd-tag b/src/cmd-tag
index c214ddd8c6..757fbc2bd3 100755
--- a/src/cmd-tag
+++ b/src/cmd-tag
@@ -10,17 +10,11 @@
 #   $ coreos-assembler tag list
 #
 import argparse
-import json
 import os
 import sys
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import fatal, rfc3339_time, write_json
-
-
-# Default location of build metadata; perhaps customizable in
-# the future?
-BUILDS_JSON = "builds/builds.json"
+from cmdlib import fatal, rfc3339_time, Builds  # noqa: E402
 
 
 def main():
@@ -59,27 +53,6 @@ def parse_args():
     return parser.parse_args()
 
 
-def init_build_data():
-    """Initialize build metadata"""
-
-    if os.path.isfile(BUILDS_JSON):
-        with open(BUILDS_JSON) as json_file:
-            build_data = json.load(json_file)
-    else:
-        print("INFO: Did not find existing builds.json; starting with empty "
-              "build list and tags list")
-        build_data = {"builds": [], "tags": []}
-
-    return build_data
-
-
-def finalize_json(build_data):
-    """Wrapper around cmdlib.write_json()"""
-
-    build_data['timestamp'] = rfc3339_time()
-    write_json(BUILDS_JSON, build_data)
-
-
 def available_tags(build_data):
     """Returns the available tag names from build metadata"""
 
@@ -89,7 +62,8 @@ def available_tags(build_data):
 def cmd_list(args):
     """List available tags in build metadata"""
 
-    build_data = init_build_data()
+    builds = Builds()
+    build_data = builds.raw()
     avail_tags = available_tags(build_data)
     if not avail_tags:
         print("No tags found")
@@ -104,16 +78,17 @@ def cmd_list(args):
 def cmd_update(args):
     """Create or update a tag to new build ID"""
 
-    build_data = init_build_data()
+    builds = Builds()
+    build_data = builds.raw()
+
     avail_tags = available_tags(build_data)
-    avail_builds = build_data["builds"]
 
     if "tags" not in build_data:
         build_data["tags"] = []
 
     # if the build doesn't exist, check for the force flag and bail out
     # if it is not used.
-    if args.build not in avail_builds:
+    if not builds.has(args.build):
         if not args.force:
             fatal("Cannot operate on a tag with a build that does not exist")
         print(f"INFO: Operating on a tag ({args.tag}) with a build "
@@ -131,7 +106,7 @@ def cmd_update(args):
                                "target": args.build}
                               if t.get("name") == args.tag
                               else t for t in build_data["tags"]]
-    finalize_json(build_data)
+    builds.flush()
 
 
 def cmd_delete(args):
@@ -139,7 +114,8 @@ def cmd_delete(args):
 
     # To delete a tag, iterate through existing tags list, and
     # drop the entry we want
-    build_data = init_build_data()
+    builds = Builds()
+    build_data = builds.raw()
     avail_tags = available_tags(build_data)
 
     if args.tag not in avail_tags:
@@ -147,9 +123,7 @@ def cmd_delete(args):
 
     build_data["tags"] = [t for t in build_data["tags"]
                           if t["name"] != args.tag]
-
-    # Write out JSON
-    finalize_json(build_data)
+    builds.flush()
 
 
 if __name__ == "__main__":
diff --git a/src/cosalib/build.py b/src/cosalib/build.py
index d2d48064d7..70936431ea 100644
--- a/src/cosalib/build.py
+++ b/src/cosalib/build.py
@@ -9,6 +9,9 @@
 import tempfile
 import shutil
 
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from cmdlib import Builds
+
 # COSA_INPATH is the _in container_ path for the image build source
 COSA_INPATH = "/cosa"
 
@@ -83,16 +86,15 @@ def __init__(self, builds_dir, build="latest", workdir=None):
 
         If workdir is None then no temporary work directory is created.
         """
-        log.info('Evaluating builds.json')
-        builds = load_json('%s/builds.json' % builds_dir)
+        builds = Builds(os.path.dirname(builds_dir))
         if build != "latest":
-            if build not in builds['builds']:
+            if not builds.has(build):
                 raise BuildError("Build was not found in builds.json")
         else:
-            build = builds['builds'][0]
+            build = builds.get_latest()
 
         log.info("Targeting build: %s", build)
-        self._build_dir = os.path.abspath("%s/%s" % (builds_dir, build))
+        self._build_dir = builds.get_build_dir(build)
 
         self._build_json = {
             "commit": None,
diff --git a/src/prune_builds b/src/prune_builds
index 907c09cbb1..c77ad8ccac 100755
--- a/src/prune_builds
+++ b/src/prune_builds
@@ -16,7 +16,7 @@ import collections
 from datetime import timedelta, datetime, timezone
 
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
-from cmdlib import write_json, rfc3339_time
+from cmdlib import Builds
 
 
 def parse_date_string(date_string):
@@ -61,76 +61,77 @@ if args.keep_last_days is not None:
     keep_younger_than = (datetime.now(timezone.utc) -
                          timedelta(days=args.keep_last_days))
 
+builds = Builds(args.workdir)
 
-builds = []
+scanned_builds = []
 builds_dir = os.path.join(args.workdir, "builds")
-builds_json = os.path.join(builds_dir, "builds.json")
-
-if os.path.isfile(builds_json):
-    with open(builds_json) as f:
-        builddata = json.load(f)
-else:
-    builddata = {'builds': []}
 
 # handle --timestamp-only
 if args.timestamp_only:
-    builddata['timestamp'] = rfc3339_time()
-    write_json(builds_json, builddata)
+    builds.bump_timestamp()
     sys.exit(0)
 
 # Handle --insert-only
 if args.insert_only:
-    builddata['builds'].insert(0, args.insert_only)
-    builddata['timestamp'] = rfc3339_time()
-    write_json(builds_json, builddata)
+    builds.insert_build(args.insert_only)
+    builds.flush()
     sys.exit(0)
 
 skip_pruning = (not keep_younger_than and args.keep_last_n == 0)
 
 # collect all builds being pointed to by tags
-tagged_builds = set([tag['target'] for tag in builddata.get('tags', [])])
+tagged_builds = set([tag['target'] for tag in builds.raw().get('tags', [])])
+
+
+def get_timestamp(entry):
+
+    # ignore dirs missing meta.json
+    meta_file = os.path.join(entry.path, 'meta.json')
+    if not os.path.isfile(meta_file):
+        print(f"Ignoring directory {entry.name}")
+        return None
+
+    # collect dirs and timestamps
+    with open(meta_file) as f:
+        j = json.load(f)
+    # Older versions only had ostree-timestamp
+    ts = j.get('coreos-assembler.build-timestamp') or j['ostree-timestamp']
+    return parse_date_string(ts)
+
 
 # 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 three non-dir things we expect there
+            # 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.name}")
+                print(f"Ignoring non-directory {entry.path}")
             continue
 
-        # ignore dirs missing meta.json
-        meta_file = os.path.join(entry.path, 'meta.json')
-        if not os.path.isfile(meta_file):
-            print(f"Ignoring directory {entry.name}")
-            continue
+        ts = get_timestamp(entry)
+        if ts:
+            scanned_builds.append(Build(id=entry.name, timestamp=ts))
 
-        # collect dirs and timestamps
-        with open(meta_file) as f:
-            j = json.load(f)
-        # Older versions only had ostree-timestamp
-        ts = j.get('coreos-assembler.build-timestamp') or j['ostree-timestamp']
-        t = parse_date_string(ts)
-        builds.append(Build(id=entry.name, timestamp=t))
 
 # just get the trivial case out of the way
-if len(builds) == 0:
+if len(scanned_builds) == 0:
     print("No builds to prune!")
     sys.exit(0)
 
 # sort by timestamp, newest first
-builds = sorted(builds, key=lambda x: x.timestamp, reverse=True)
+scanned_builds = sorted(scanned_builds,
+                        key=lambda x: x.timestamp,
+                        reverse=True)
 
 new_builds = []
 builds_to_delete = []
 
 if skip_pruning:
-    new_builds = builds
+    new_builds = scanned_builds
 else:
     if keep_younger_than:
-        for build in builds:
+        for build in scanned_builds:
             if build.id in tagged_builds:
                 print(f"Skipping tagged build {build.id}")
                 new_builds.append(build)
@@ -143,7 +144,7 @@ else:
     else:
         n = args.keep_last_n
         assert(n > 0)
-        for build in builds:
+        for build in scanned_builds:
             # skip tagged builds and don't count them towards the limit
             if build.id in tagged_builds:
                 print(f"Skipping tagged build {build.id}")
@@ -156,9 +157,10 @@ else:
                 new_builds.append(build)
                 n = n - 1
 
-builddata['builds'] = [x.id for x in new_builds]
-builddata['timestamp'] = rfc3339_time()
-write_json(builds_json, builddata)
+builds.raw()['builds'] = []
+for build in reversed(new_builds):
+    builds.insert_build(build.id)
+builds.bump_timestamp()
 
 # if we're not pruning, then we're done!
 if skip_pruning: