Implement FileFinder::prune

2 files changed, 47 insertions, 0 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinder.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinder.java
index a41281d0479..f9a369aab27 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinder.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinder.java
@@ -14,6 +14,8 @@ import java.nio.file.attribute.BasicFileAttributes;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Set;
@@ -35,6 +37,7 @@ public class FileFinder {
     private static final Logger logger = Logger.getLogger(FileFinder.class.getName());
 
     private final Path basePath;
+    private final Set<Path> pruned = new HashSet<>();
     private Predicate<FileAttributes> matcher;
     private int maxDepth = Integer.MAX_VALUE;
 
@@ -72,6 +75,27 @@ public class FileFinder {
     }
 
     /**
+     * Path for which whole directory tree will be skipped, including the path itself.
+     * The path must be under {@code basePath} or be relative to {@code basePath}.
+     */
+    public FileFinder prune(Path path) {
+        if (!path.isAbsolute())
+            path = basePath.resolve(path);
+
+        if (!path.startsWith(basePath))
+            throw new IllegalArgumentException("Prune path " + path + " is not under base path " + basePath);
+
+        this.pruned.add(path);
+        return this;
+    }
+
+    /** Convenience method for pruning multiple paths, see {@link #prune(Path)}. */
+    public FileFinder prune(Collection<Path> paths) {
+        paths.forEach(this::prune);
+        return this;
+    }
+
+    /**
      * Maximum depth (relative to basePath) where contents should be matched with the given filters.
      * Default is unlimited.
      */
@@ -145,6 +169,8 @@ public class FileFinder {
 
                 @Override
                 public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) {
+                    if (pruned.contains(dir)) return FileVisitResult.SKIP_SUBTREE;
+
                     currentLevel++;
 
                     FileAttributes attributes = new FileAttributes(dir, attrs);
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinderTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinderTest.java
index 8fff2141252..26897839b32 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinderTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/file/FileFinderTest.java
@@ -61,6 +61,27 @@ public class FileFinderTest {
         }
 
         @Test
+        public void all_files_recursive_with_prune_relative() {
+            assertFileHelper(FileFinder.files(testRoot()).prune(fileSystem.getPath("test")),
+
+                    of("file-1.json", "test.json", "test.txt"),
+                    of("test", "test/file.txt", "test/data.json", "test/subdir-1", "test/subdir-1/test", "test/subdir-2"));
+        }
+
+        @Test
+        public void all_files_recursive_with_prune_absolute() {
+            assertFileHelper(FileFinder.files(testRoot()).prune(testRoot().resolve("test/subdir-1")),
+
+                    of("file-1.json", "test.json", "test.txt", "test/file.txt", "test/data.json"),
+                    of("test", "test/subdir-1", "test/subdir-1/test", "test/subdir-2"));
+        }
+
+        @Test(expected = IllegalArgumentException.class)
+        public void throws_if_prune_path_not_under_base_path() {
+            FileFinder.files(Paths.get("/some/path")).prune(Paths.get("/other/path"));
+        }
+
+        @Test
         public void with_file_filter_recursive() {
             assertFileHelper(FileFinder.files(testRoot())
                             .match(FileFinder.nameEndsWith(".json")),