print explain callback

tokuhirom · Apr 5, 2016 · e5c8437 · e5c8437
1 parent efd8e9b
commit e5c8437
Show file tree

Hide file tree

Showing 2 changed files with 71 additions and 30 deletions.
diff --git a/jdbc-querylog/src/main/java/me/geso/jdbcquerylog/PreparedStatementLogger.java b/jdbc-querylog/src/main/java/me/geso/jdbcquerylog/PreparedStatementLogger.java
@@ -1,13 +1,9 @@
 package me.geso.jdbcquerylog;
 
-import de.vandermeer.asciitable.v2.RenderedTable;
-import de.vandermeer.asciitable.v2.V2_AsciiTable;
-import de.vandermeer.asciitable.v2.render.V2_AsciiTableRenderer;
-import de.vandermeer.asciitable.v2.render.WidthLongestLine;
-import de.vandermeer.asciitable.v2.themes.V2_E_TableThemes;
 import me.geso.jdbctracer.PreparedStatementListener;
 
 import java.sql.*;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -27,48 +23,49 @@ public void trace(Connection connection, long elapsed, String query, List<Object
         if (QueryLogDriver.isCompact()) {
             query = compact(query);
         }
-        System.err.println(query);
+        QueryLogDriver.getPrintQuery().accept(connection, query);
+
         if (QueryLogDriver.isExplain()) {
             try (PreparedStatement preparedStatement = connection.prepareStatement("EXPLAIN " + query)) {
                 try (ResultSet resultSet = preparedStatement.executeQuery()) {
-                    V2_AsciiTable at = new V2_AsciiTable();
                     ResultSetMetaData metaData = resultSet.getMetaData();
                     int columnCount = metaData.getColumnCount();
+
+                    Object[] header = IntStream.rangeClosed(1, columnCount)
+                            .mapToObj(i -> {
+                                try {
+                                    return metaData.getColumnName(i);
+                                } catch (SQLException e) {
+                                    throw new RuntimeException(e);
+                                }
+                            }).toArray();
+
+                    List<Object[]> rows = new ArrayList<>();
                     while (resultSet.next()) {
-                        at.addRule();
-                        at.addRow(IntStream.rangeClosed(1, columnCount)
-                                .mapToObj(i -> {
-                                    try {
-                                        return metaData.getColumnName(i);
-                                    } catch (SQLException e) {
-                                        throw new RuntimeException(e);
-                                    }
-                                }).toArray());
-                        at.addRule();
-                        at.addRow(IntStream.rangeClosed(1, columnCount)
+                        Object[] row = IntStream.rangeClosed(1, columnCount)
                                 .mapToObj(i -> {
                                     try {
                                         return resultSet.getString(i);
                                     } catch (SQLException e) {
                                         throw new RuntimeException(e);
                                     }
-                                }).toArray());
-                        at.addRule();
-                    }
+                                }).toArray();
+                        rows.add(row);
 
-                    V2_AsciiTableRenderer rend = new V2_AsciiTableRenderer();
-                    rend.setTheme(V2_E_TableThemes.UTF_LIGHT.get());
-                    rend.setWidth(new WidthLongestLine());
-                    RenderedTable render = rend.render(at);
-                    for (String line : render.toString().split("\n")) {
-                        System.err.println(line);
                     }
+
+                    QueryLogDriver.getPrintExplain().accept(
+                            connection,
+                            query,
+                            header,
+                            rows
+                    );
                 }
             }
         }
     }
 
-    public String compact(String query) {
+    private String compact(String query) {
         return query.replaceAll("\\n", " ");
     }
 
@@ -90,6 +87,6 @@ private String bind(String query, List<Object> binds) {
             matcher.appendTail(sb);
             return sb.toString();
         }
-        return query.toString();
+        return query;
     }
 }
diff --git a/jdbc-querylog/src/main/java/me/geso/jdbcquerylog/QueryLogDriver.java b/jdbc-querylog/src/main/java/me/geso/jdbcquerylog/QueryLogDriver.java
@@ -1,18 +1,41 @@
 package me.geso.jdbcquerylog;
 
+import de.vandermeer.asciitable.v2.RenderedTable;
+import de.vandermeer.asciitable.v2.V2_AsciiTable;
+import de.vandermeer.asciitable.v2.render.V2_AsciiTableRenderer;
+import de.vandermeer.asciitable.v2.render.WidthLongestLine;
+import de.vandermeer.asciitable.v2.themes.V2_E_TableThemes;
 import me.geso.jdbctracer.TracerConnection;
 
 import java.sql.*;
 import java.util.Enumeration;
+import java.util.List;
 import java.util.Objects;
 import java.util.Properties;
-import java.util.function.Consumer;
+import java.util.function.BiConsumer;
 import java.util.logging.Logger;
 
 public class QueryLogDriver implements Driver {
     private static boolean enabled = true;
     private static boolean explain = false;
     private static boolean compact = true;
+    private static BiConsumer<Connection, String> printQuery = (connection, query) -> {
+        System.err.println(query);
+    };
+    private static PrintExplainCallback printExplain = (connection, query, header, rows) -> {
+        V2_AsciiTable at = new V2_AsciiTable();
+        at.addRule();
+        at.addRow(header);
+        at.addRule();
+        rows.forEach(at::addRow);
+        at.addRule();
+
+        V2_AsciiTableRenderer rend = new V2_AsciiTableRenderer();
+        rend.setTheme(V2_E_TableThemes.UTF_LIGHT.get());
+        rend.setWidth(new WidthLongestLine());
+        RenderedTable render = rend.render(at);
+        System.err.println(render.toString());
+    };
 
     static {
         try {
@@ -46,6 +69,27 @@ public static void setCompact(boolean compact) {
         QueryLogDriver.compact = compact;
     }
 
+    public static void setPrintQuery(BiConsumer<Connection, String> printQuery) {
+        QueryLogDriver.printQuery = printQuery;
+    }
+
+    public static BiConsumer<Connection, String> getPrintQuery() {
+        return QueryLogDriver.printQuery;
+    }
+
+    public static PrintExplainCallback getPrintExplain() {
+        return printExplain;
+    }
+
+    public static void setPrintExplain(PrintExplainCallback printExplain) {
+        QueryLogDriver.printExplain = printExplain;
+    }
+
+    @FunctionalInterface
+    public interface PrintExplainCallback {
+        void accept(Connection connection, String query, Object[] header, List<Object[]> rows);
+    }
+
     @Override
     public Connection connect(String url, Properties info) throws SQLException {
         String underlingUri = parseURL(url);