bin/prune: Add --only-branch

In 5c940987e768523ef1411b65bcaad09fba6befef / #646 we
added `--retain-branch-depth`; this adds a symmetric
`--only-branch` for the case where a repo owner just
wants to prune a specific branch.

The implementation here is pretty straightforward; we
just walk all refs and inject the equivalent of
`--retain-branch-depth=$ref=-1` if they're *not* in
`--only-branch`.

Closes: https://github.com/ostreedev/ostree/issues/1115

Closes: #1127
Approved by: jlebon

2 files changed, 96 insertions, 2 deletions

diff --git a/src/ostree/ot-builtin-prune.c b/src/ostree/ot-builtin-prune.c
index 2f66334e..6120e0a9 100644
--- a/src/ostree/ot-builtin-prune.c
+++ b/src/ostree/ot-builtin-prune.c
@@ -35,6 +35,7 @@ static gboolean opt_refs_only;
 static char *opt_delete_commit;
 static char *opt_keep_younger_than;
 static char **opt_retain_branch_depth;
+static char **opt_only_branches;
 
 /* ATTENTION:
  * Please remember to update the bash-completion script (bash/ostree) and
@@ -49,6 +50,7 @@ static GOptionEntry options[] = {
   { "keep-younger-than", 0, 0, G_OPTION_ARG_STRING, &opt_keep_younger_than, "Prune all commits older than the specified date", "DATE" },
   { "static-deltas-only", 0, 0, G_OPTION_ARG_NONE, &opt_static_deltas_only, "Change the behavior of delete-commit and keep-younger-than to prune only static deltas" },
   { "retain-branch-depth", 0, 0, G_OPTION_ARG_STRING_ARRAY, &opt_retain_branch_depth, "Additionally retain BRANCH=DEPTH commits", "BRANCH=DEPTH" },
+  { "only-branch", 0, 0, G_OPTION_ARG_STRING_ARRAY, &opt_only_branches, "Only prune BRANCH (may be specified multiple times)", "BRANCH" },
   { NULL }
 };
 
@@ -185,7 +187,7 @@ ostree_builtin_prune (int argc, char **argv, GCancellable *cancellable, GError *
   gint n_objects_total;
   gint n_objects_pruned;
   guint64 objsize_total;
-  if (!(opt_retain_branch_depth || opt_keep_younger_than))
+  if (!(opt_retain_branch_depth || opt_keep_younger_than || opt_only_branches))
     {
       if (!ostree_repo_prune (repo, pruneflags, opt_depth,
                               &n_objects_total, &n_objects_pruned, &objsize_total,
@@ -210,6 +212,7 @@ ostree_builtin_prune (int argc, char **argv, GCancellable *cancellable, GError *
             return glnx_throw (error, "Could not parse '%s'", opt_keep_younger_than);
         }
 
+      /* Process --retain-branch-depth */
       for (char **iter = opt_retain_branch_depth; iter && *iter; iter++)
         {
           /* bd should look like BRANCH=DEPTH where DEPTH is an int */
@@ -239,6 +242,45 @@ ostree_builtin_prune (int argc, char **argv, GCancellable *cancellable, GError *
                                   cancellable, error))
         return FALSE;
 
+      /* Process --only-branch. Note this combines with --retain-branch-depth; one
+       * could do e.g.:
+       *  * --only-branch exampleos/x86_64/foo
+       *  * --only-branch exampleos/x86_64/bar
+       *  * --retain-branch-depth exampleos/x86_64/foo=0
+       *  * --depth 5
+       * to prune exampleos/x86_64/foo to just the latest commit, and
+       * exampleos/x86_64/bar to a depth of 5.
+       */
+      if (opt_only_branches)
+        {
+          /* Turn --only-branch into a set */
+          g_autoptr(GHashTable) only_branches_set = g_hash_table_new (g_str_hash, g_str_equal);
+          for (char **iter = opt_only_branches; iter && *iter; iter++)
+            {
+              const char *ref = *iter;
+              /* Ensure the specified branch exists */
+              if (!ostree_repo_resolve_rev (repo, ref, FALSE, NULL, error))
+                return FALSE;
+              g_hash_table_add (only_branches_set, (char*)ref);
+            }
+
+          /* Iterate over all refs, add equivalent of --retain-branch-depth=$ref=-1
+           * if the ref isn't in --only-branch set and there wasn't already a
+           * --retain-branch-depth specified for it.
+           */
+          GLNX_HASH_TABLE_FOREACH (all_refs, const char *, ref)
+            {
+              if (!g_hash_table_contains (only_branches_set, ref) &&
+                  !g_hash_table_contains (retain_branch_depth, ref))
+                {
+                  g_hash_table_insert (retain_branch_depth, g_strdup (ref), GINT_TO_POINTER ((int)-1));
+                }
+            }
+        }
+
+      /* Traverse each ref, and gather all objects pointed to by it up to a
+       * specific depth (if configured).
+       */
       g_hash_table_iter_init (&hash_iter, all_refs);
       while (g_hash_table_iter_next (&hash_iter, &key, &value))
         {
@@ -275,6 +317,7 @@ ostree_builtin_prune (int argc, char **argv, GCancellable *cancellable, GError *
             return FALSE;
         }
 
+      /* We've gathered the reachable set; start the prune ✀ */
       { OstreeRepoPruneOptions opts = { pruneflags, reachable };
         if (!ostree_repo_prune_from_reachable (repo, &opts,
                                                &n_objects_total,
diff --git a/tests/test-prune.sh b/tests/test-prune.sh
index e41c0fbc..07548822 100755
--- a/tests/test-prune.sh
+++ b/tests/test-prune.sh
@@ -25,7 +25,7 @@ skip_without_user_xattrs
 
 setup_fake_remote_repo1 "archive"
 
-echo '1..5'
+echo '1..12'
 
 cd ${test_tmpdir}
 mkdir repo
@@ -235,3 +235,54 @@ assert_repo_has_n_commits repo 9
 $OSTREE fsck
 
 echo "ok retain branch depth (alone)"
+
+# Test --only-branch with --depth=0; this should be exactly identical to the
+# above with a result of 9.
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=dev --depth=0
+assert_repo_has_n_commits repo 9
+$OSTREE fsck
+echo "ok --only-branch --depth=0"
+
+# Test --only-branch with --depth=1; should just add 1 to the above, for 10.
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=dev --depth=1
+assert_repo_has_n_commits repo 10
+echo "ok --only-branch --depth=1"
+
+# Test --only-branch with all branches
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=dev --only-branch=stable --depth=0
+assert_repo_has_n_commits repo 2
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=dev --only-branch=stable --depth=1
+assert_repo_has_n_commits repo 4
+echo "ok --only-branch (all) --depth=1"
+
+# Test --only-branch and --retain-branch-depth overlap
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=dev --only-branch=stable --depth=0 \
+                     --retain-branch-depth=stable=2
+assert_repo_has_n_commits repo 4
+echo "ok --only-branch and --retain-branch-depth overlap"
+
+# Test --only-branch and --retain-branch-depth together
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=dev --depth=0 --retain-branch-depth=stable=2
+assert_repo_has_n_commits repo 4
+echo "ok --only-branch and --retain-branch-depth together"
+
+# Test --only-branch with --keep-younger-than; this should be identical to the test
+# above for --retain-branch-depth=stable=-1
+reinitialize_datesnap_repo
+${CMD_PREFIX} ostree --repo=repo prune --only-branch=stable --keep-younger-than="1 week ago"
+assert_repo_has_n_commits repo 11
+echo "ok --only-branch --keep-younger-than"
+
+# Test --only-branch with a nonexistent ref
+reinitialize_datesnap_repo
+if ${CMD_PREFIX} ostree --repo=repo prune --only-branch=BACON 2>err.txt; then
+    fatal "we pruned BACON?"
+fi
+assert_file_has_content err.txt "Refspec.*BACON.*not found"
+echo "ok --only-branch=BACON"