Add logdirs table

Author: kawamuray

src/main/java/kmql/Database.java

@@ -5,6 +5,7 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Map.Entry;
@@ -51,6 +52,10 @@ public void prepareTable(String name, AdminClient adminClient) throws Exception
         if (meta.initialized) {
             return;
         }
+        Collection<String> dependencyTables = meta.table.dependencyTables();
+        for (String dependencyTable : dependencyTables) {
+            prepareTable(dependencyTable, adminClient);
+        }
         meta.table.prepare(connection, adminClient);
         meta.initialized = true;
     }

src/main/java/kmql/Table.java

@@ -1,11 +1,18 @@
 package kmql;
 
+import static java.util.Collections.emptyList;
+
 import java.sql.Connection;
+import java.util.Collection;
 
 import org.apache.kafka.clients.admin.AdminClient;
 
 public interface Table {
     String name();
 
+    default Collection<String> dependencyTables() {
+        return emptyList();
+    }
+
     void prepare(Connection connection, AdminClient adminClient) throws Exception;
 }

src/main/java/kmql/TableRegistry.java

@@ -8,6 +8,7 @@
 import java.util.concurrent.ConcurrentMap;
 
 import kmql.table.BrokersTable;
+import kmql.table.LogdirsTable;
 import kmql.table.ReplicasTable;
 
 public class TableRegistry implements Iterable<Map.Entry<String, Table>> {
@@ -16,6 +17,7 @@ public class TableRegistry implements Iterable<Map.Entry<String, Table>> {
     static {
         registerDefault(new ReplicasTable());
         registerDefault(new BrokersTable());
+        registerDefault(new LogdirsTable());
     }
 
     private final ConcurrentMap<String, Table> tables;

src/main/java/kmql/table/LogdirsTable.java

@@ -0,0 +1,80 @@
+package kmql.table;
+
+import static java.util.Collections.singletonList;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.Statement;
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.common.requests.DescribeLogDirsResponse.LogDirInfo;
+import org.apache.kafka.common.requests.DescribeLogDirsResponse.ReplicaInfo;
+
+import kmql.Table;
+
+public class LogdirsTable implements Table {
+    @Override
+    public String name() {
+        return "logdirs";
+    }
+
+    @Override
+    public Collection<String> dependencyTables() {
+        return singletonList("brokers");
+    }
+
+    @Override
+    public void prepare(Connection connection, AdminClient adminClient) throws Exception {
+        try (Statement stmt = connection.createStatement()) {
+            stmt.execute("CREATE TABLE logdirs ("
+                         + "broker_id INT NOT NULL,"
+                         + "path VARCHAR(255) NOT NULL,"
+                         + "topic VARCHAR(255) NOT NULL,"
+                         + "partition INT NOT NULL,"
+                         + "size BIGINT NOT NULL,"
+                         + "offset_lag BIGINT NOT NULL,"
+                         + "is_future BOOLEAN NOT NULL,"
+                         + "PRIMARY KEY (broker_id, path, topic, partition))");
+        }
+
+        List<Integer> brokerIds = new ArrayList<>();
+        try (Statement stmt = connection.createStatement();
+            ResultSet results = stmt.executeQuery("SELECT id FROM brokers")) {
+            while (results.next()) {
+                int brokerId = results.getInt(1);
+                brokerIds.add(brokerId);
+            }
+        }
+
+        Map<Integer, Map<String, LogDirInfo>> logDirs = adminClient.describeLogDirs(brokerIds).all().get();
+        try (PreparedStatement stmt = connection.prepareStatement(
+                "INSERT INTO logdirs (broker_id, path, topic, partition, size, offset_lag, is_future) VALUES (?, ?, ?, ?, ?, ?, ?)")) {
+            for (Entry<Integer, Map<String, LogDirInfo>> entry : logDirs.entrySet()) {
+                int brokerId = entry.getKey();
+                for (Entry<String, LogDirInfo> dirEntry : entry.getValue().entrySet()) {
+                    String path = dirEntry.getKey();
+                    LogDirInfo info = dirEntry.getValue();
+                    for (Entry<TopicPartition, ReplicaInfo> replicaEntry : info.replicaInfos.entrySet()) {
+                        TopicPartition tp = replicaEntry.getKey();
+                        ReplicaInfo replicaInfo = replicaEntry.getValue();
+                        stmt.setInt(1, brokerId);
+                        stmt.setString(2, path);
+                        stmt.setString(3, tp.topic());
+                        stmt.setInt(4, tp.partition());
+                        stmt.setLong(5, replicaInfo.size);
+                        stmt.setLong(6, replicaInfo.offsetLag);
+                        stmt.setBoolean(7, replicaInfo.isFuture);
+                        stmt.executeUpdate();
+                    }
+                }
+            }
+        }
+    }
+}