rewrite CassandraData.parse to use Files.walkFileTree instead of Files.find

Author: smiklosovic

src/main/java/com/instaclustr/esop/impl/CassandraData.java

@@ -1,8 +1,9 @@
 package com.instaclustr.esop.impl;
 
 import static java.lang.String.format;
-import static java.util.stream.Collectors.groupingBy;
+import static java.util.stream.Collectors.toList;
 
+import java.io.IOException;
 import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -17,6 +18,7 @@
 import java.util.Map.Entry;
 import java.util.Optional;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import com.instaclustr.esop.impl.RenamedEntities.Renamed;
 
@@ -158,27 +160,102 @@ public RenamedEntities getRenamedEntities() {
         return renamedEntities;
     }
 
+    public static class SnapshotsLister extends SimpleFileVisitor<Path> {
+
+        private boolean isDropped = false;
+
+        @Override
+        public FileVisitResult preVisitDirectory(final Path dir, final BasicFileAttributes attrs) throws IOException {
+            if (dir.getParent().getFileName().toString().equals("snapshots")) {
+                if (dir.getFileName().toString().startsWith("dropped-")) {
+                    isDropped = true;
+                    return FileVisitResult.TERMINATE;
+                } else {
+                    return FileVisitResult.SKIP_SUBTREE;
+                }
+            } else {
+                return FileVisitResult.CONTINUE;
+            }
+        }
+
+        public boolean isDropped() {
+            return isDropped;
+        }
+    }
+
+    public static class KeyspaceTableLister extends SimpleFileVisitor<Path> {
+
+        private final Path cassandraDir;
+        private final Map<Path, List<Path>> dataDirs = new HashMap<>();
+
+        public KeyspaceTableLister(final Path cassandraDir) {
+            this.cassandraDir = cassandraDir;
+        }
+
+        @Override
+        public FileVisitResult preVisitDirectory(final Path dir, final BasicFileAttributes attrs) throws IOException {
+            // we hit keyspace
+            if (dir.getParent().equals(cassandraDir)) {
+                dataDirs.putIfAbsent(dir, new ArrayList<>());
+                return FileVisitResult.CONTINUE;
+                // we hit table
+            } else if (dir.getParent().getParent().equals(cassandraDir)) {
+                // detect if it is a dropped table
+                Path snapshotsDir = dir.resolve("snapshots");
+                if (Files.exists(snapshotsDir)) {
+                    SnapshotsLister snapshotsLister = new SnapshotsLister();
+                    Files.walkFileTree(snapshotsDir, snapshotsLister);
+                    if (!snapshotsLister.isDropped()) {
+                        dataDirs.get(dir.getParent()).add(dir);
+                    }
+                } else {
+                    dataDirs.get(dir.getParent()).add(dir);
+                }
+
+                return FileVisitResult.SKIP_SUBTREE;
+            } else {
+                return FileVisitResult.CONTINUE;
+            }
+        }
+
+        /**
+         * Remove keyspaces which have 0 tables, it means that each table has a snapshot with "dropped-" snapshot name
+         */
+        public void removeDroppedKeyspaces() {
+            final List<Path> droppedKeyspaces = dataDirs
+                .entrySet()
+                .stream()
+                .filter(entry -> entry.getValue().isEmpty())
+                .map(Entry::getKey)
+                .collect(toList());
+
+            for (final Path droppedKeyspace : droppedKeyspaces) {
+                dataDirs.remove(droppedKeyspace);
+            }
+        }
+
+        public Map<Path, List<Path>> getDataDirs() {
+            return dataDirs;
+        }
+
+        @Override
+        public String toString() {
+            return dataDirs.toString();
+        }
+    }
+
     public static CassandraData parse(final Path cassandraDir) throws Exception {
 
         if (!Files.exists(cassandraDir)) {
             return CassandraData.empty();
         }
 
-        final Map<Path, List<Path>> dataDirs = Files.find(cassandraDir,
-                                                          2,
-                                                          (path, basicFileAttributes) -> basicFileAttributes.isDirectory() &&
-                                                              !path.getParent().equals(cassandraDir) &&
-                                                              !path.equals(cassandraDir))
-            // take only these into consideration which do not have "snapshots/dropped-"
-            .filter(table -> {
-                try {
-                    return Files.find(table, 2, (p, b) -> b.isDirectory() && p.toString().contains("snapshots/dropped-")).count() == 0;
-                } catch (final Exception ex) {
-                    return false;
-                }
+        final KeyspaceTableLister lister = new KeyspaceTableLister(cassandraDir);
+
+        Files.walkFileTree(cassandraDir, lister);
+        lister.removeDroppedKeyspaces();
 
-            })
-            .collect(groupingBy(Path::getParent));
+        final Map<Path, List<Path>> dataDirs = lister.getDataDirs();
 
         final Map<String, Map<String, String>> tableIdsMap = new HashMap<>();
 

src/test/java/com/instaclustr/esop/backup/KeyspaceTableListerTest.java

@@ -0,0 +1,68 @@
+package com.instaclustr.esop.backup;
+
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+
+import com.instaclustr.esop.impl.CassandraData.KeyspaceTableLister;
+import com.instaclustr.io.FileUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class KeyspaceTableListerTest {
+
+    private static final Logger logger = LoggerFactory.getLogger(KeyspaceTableLister.class);
+
+    @Test
+    public void testKeyspaceTableLister() throws Exception {
+        Path cassandrDir = Files.createTempDirectory("keyspace-table-lister");
+
+        // in total 20000 files in 1000 tables in 100 keyspaces
+
+        for (int i = 0; i < 90; i++) {
+            Files.createDirectory(cassandrDir.resolve("ks" + i));
+            for (int j = 0; j < 10; j++) {
+                Files.createDirectory(cassandrDir.resolve("ks" + i + "/tb-" + j));
+                for (int k = 0; k < 20; k++) {
+                    Files.createFile(cassandrDir.resolve("ks" + i + "/tb-" + j + "/file" + k));
+                }
+            }
+        }
+
+        for (int i = 0; i < 5; i++) {
+            for (int j = 0; j < 5; j++) {
+                Files.createDirectory(cassandrDir.resolve("ks" + i + "/tb-" + j + "/snapshots"));
+                Files.createDirectory(cassandrDir.resolve("ks" + i + "/tb-" + j + "/snapshots/dropped-3424324"));
+            }
+        }
+
+        // keyspaces from no. 90 - 99 will have all tables dropped which renders them dropped too
+
+        for (int i = 90; i < 99; i++) {
+            Files.createDirectory(cassandrDir.resolve("ks" + i));
+            for (int j = 0; j < 5; j++) {
+                Files.createDirectory(cassandrDir.resolve("ks" + i + "/tb-" + j));
+                Files.createDirectory(cassandrDir.resolve("ks" + i + "/tb-" + j + "/snapshots"));
+                Files.createDirectory(cassandrDir.resolve("ks" + i + "/tb-" + j + "/snapshots/dropped-3424324"));
+            }
+        }
+
+        KeyspaceTableLister lister = new KeyspaceTableLister(cassandrDir);
+        Files.walkFileTree(cassandrDir, lister);
+
+        lister.removeDroppedKeyspaces();
+
+        Map<Path, List<Path>> dataDirs = lister.getDataDirs();
+        Assert.assertFalse(dataDirs.isEmpty());
+
+        Assert.assertEquals(90, dataDirs.size());
+
+        // 25 tables are dropped and 10 keyspaces, each having 10 tables, have all tables dropped: 1000 - 100 - 25 = 875
+        Assert.assertEquals(875, dataDirs.values().stream().map(List::size).reduce(0, Integer::sum).intValue());
+
+        FileUtils.deleteDirectory(cassandrDir);
+    }
+}