From 6f4d286ff4484e2d4e49580ae9a49be6328a0b52 Mon Sep 17 00:00:00 2001
From: Christian Kellner <christian@kellner.me>
Date: Mon, 27 Jan 2020 16:21:47 +0100
Subject: [PATCH] osbuild: support for checkpoints during build

Add a new `--checkpoint` option, which can be provided multiple
times, that indicate after which stages a the current stage of
the tree should be committed to the object store; the tree id
will be the treesum of the tree at that point and a reference
is created with the id of the stage at the point.
The argument to `--checkpoint` is the id of the stage. If not
all the given checkpoints can be found the execution will be
aborted.
---
 osbuild/__main__.py | 29 +++++++++++++++++++++++++++++
 osbuild/pipeline.py |  3 +++
 2 files changed, 32 insertions(+)

diff --git a/osbuild/__main__.py b/osbuild/__main__.py
index 70161075..9f087e9a 100755
--- a/osbuild/__main__.py
+++ b/osbuild/__main__.py
@@ -10,6 +10,25 @@ BOLD = "\033[1m"
 RED = "\033[31m"
 
 
+def mark_checkpoints(pipeline, checkpoints):
+    points = set(checkpoints)
+    def mark_stage(stage):
+        for c in points:
+            if stage.id == c:
+                stage.checkpoint = True
+                points.remove(c)
+                return
+
+    def mark_pipeline(pl):
+        for stage in pl.stages:
+            mark_stage(stage)
+        if pl.build:
+            mark_pipeline(pl.build)
+
+    mark_pipeline(pipeline)
+    return points
+
+
 def main():
     parser = argparse.ArgumentParser(description="Build operating system images")
     parser.add_argument("pipeline_path", metavar="PIPELINE",
@@ -25,6 +44,8 @@ def main():
                         help="json file containing a dictionary of secrets that are passed to sources")
     parser.add_argument("-l", "--libdir", metavar="DIRECTORY", type=os.path.abspath,
                         help="the directory containing stages, assemblers, and the osbuild library")
+    parser.add_argument("--checkpoint", metavar="CHECKPOINT", action="append", type=str, default=None,
+                        help="stage to commit to the object store during build (can be passed multiple times)")
     parser.add_argument("--json", action="store_true",
                         help="output results in JSON format")
     args = parser.parse_args()
@@ -51,6 +72,14 @@ def main():
         with open(args.secrets) as f:
             secrets = json.load(f)
 
+    if args.checkpoint:
+        missed = mark_checkpoints(pipeline, args.checkpoint)
+        if missed:
+            for checkpoint in missed:
+                print(f"Checkpoint {BOLD}{checkpoint}{RESET} not found!")
+            print(f"{RESET}{BOLD}{RED}Failed{RESET}")
+            return 1
+
     try:
         r = pipeline.run(
             args.store,
diff --git a/osbuild/pipeline.py b/osbuild/pipeline.py
index a0411fca..8984a0a7 100644
--- a/osbuild/pipeline.py
+++ b/osbuild/pipeline.py
@@ -51,6 +51,7 @@ class Stage:
         self.build = build
         self.base = base
         self.options = options
+        self.checkpoint = False
 
     @property
     def id(self):
@@ -271,6 +272,8 @@ class Pipeline:
                                               var=store,
                                               source_options=source_options,
                                               secrets=secrets)
+                                if stage.checkpoint:
+                                    object_store.snapshot(tree, stage.id)
                                 results["stages"].append(r.as_dict())
                     except BuildError as err:
                         results["stages"].append(err.as_dict())