Use EnumMap for callbacks (#9)

metabase · Feb 28, 2024 · aba1c72 · aba1c72
1 parent cad4956
commit aba1c72
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 24 deletions.
diff --git a/java/com/metabase/macaw/AstWalker.java b/java/com/metabase/macaw/AstWalker.java
@@ -5,10 +5,9 @@
 import clojure.lang.IFn;
 import clojure.lang.Keyword;
 
-import java.util.HashMap;
+import java.util.EnumMap;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 
 import net.sf.jsqlparser.expression.AllValue;
 import net.sf.jsqlparser.expression.AnalyticExpression;
@@ -169,6 +168,9 @@
 import net.sf.jsqlparser.statement.update.Update;
 import net.sf.jsqlparser.statement.upsert.Upsert;
 
+import static com.metabase.macaw.AstWalker.CallbackKey.COLUMN;
+import static com.metabase.macaw.AstWalker.CallbackKey.TABLE;
+
 /**
  * Walks the AST, using JSqlParser's `visit()` methods. Each `visit()` method additionally calls an applicable callback
  * method provided in the `callbacks` map. Supported callbacks have a corresponding key string (see below).
@@ -207,43 +209,43 @@
 public class AstWalker<Acc> implements SelectVisitor, FromItemVisitor, ExpressionVisitor,
        SelectItemVisitor, StatementVisitor {
 
-    public static final String COLUMN_STR = "column";
-    public static final String TABLE_STR = "table";
-    public static final Set<String> SUPPORTED_CALLBACK_KEYS = Set.of(COLUMN_STR, TABLE_STR);
+    public enum CallbackKey {
+        COLUMN,
+        TABLE;
+
+        public String toString() {
+            return name().toLowerCase();
+        }
+    }
 
     private static final String NOT_SUPPORTED_YET = "Not supported yet.";
 
     private Acc acc;
-    private final Map<String, IFn> callbacks;
+    private final EnumMap<CallbackKey, IFn> callbacks;
 
     /**
      * Construct a new walker with the given `callbacks`. The `callbacks` should be a (Clojure) map like so:
      *
      * <pre><code>
-     *   (ASTWalker. {:column (fn [col] (do-something-with-the-found col))
-     *                :table  (fn [table] (...))})
+     *   (ASTWalker. {AstWalker$CallbackKey/COLUMN (fn [acc col] (do-something-with-the-found-column acc col))
+     *                AstWalker$CallbackKey/TABLE  (fn [acc table] (...))})
      * </code></pre>
      *
-     * The appropriate callback fn will be invoked for every matching element found. The list of supported keys can be found in [[SUPPORTED_CALLBACK_KEYS]].
+     * The appropriate callback fn will be invoked for every matching element found. Valid keys are of course defined by
+     * the nested [[CallbackKey]] enum.
      * <p>
      * Silently rejects invalid keys.
      */
-    public AstWalker(Map<Keyword, IFn> callbacksWithKeywordKeys, Acc val) {
+    public AstWalker(Map<CallbackKey, IFn> rawCallbacks, Acc val) {
         this.acc = val;
-        this.callbacks = new HashMap<>(SUPPORTED_CALLBACK_KEYS.size());
-        for(Map.Entry<Keyword, IFn> entry : callbacksWithKeywordKeys.entrySet()) {
-            String keyName = entry.getKey().getName();
-            if (SUPPORTED_CALLBACK_KEYS.contains(keyName)) {
-                this.callbacks.put(keyName, entry.getValue());
-            }
-        }
+        this.callbacks = new EnumMap<>(rawCallbacks);
     }
 
     /**
      * Safely invoke the given callback by name.
      */
-    public void invokeCallback(String callbackName, Object visitedItem) {
-        IFn callback = this.callbacks.get(callbackName);
+    public void invokeCallback(CallbackKey key, Object visitedItem) {
+        IFn callback = this.callbacks.get(key);
         if (callback != null) {
             //noinspection unchecked
             acc = (Acc) callback.invoke(acc, visitedItem);
@@ -340,7 +342,7 @@ public void visit(PlainSelect plainSelect) {
 
     @Override
     public void visit(Table table) {
-        invokeCallback(TABLE_STR, table);
+        invokeCallback(TABLE, table);
     }
 
     @Override
@@ -368,7 +370,7 @@ public void visit(OverlapsCondition overlapsCondition) {
 
     @Override
     public void visit(Column tableColumn) {
-        invokeCallback(COLUMN_STR, tableColumn);
+        invokeCallback(COLUMN, tableColumn);
         if (tableColumn.getTable() != null
                 && tableColumn.getTable().getName() != null) {
             visit(tableColumn.getTable());

diff --git a/src/macaw/core.clj b/src/macaw/core.clj
@@ -1,7 +1,8 @@
 (ns macaw.core
   (:import
    (com.metabase.macaw
-    AstWalker)
+    AstWalker
+    AstWalker$CallbackKey)
    (net.sf.jsqlparser.parser
     CCJSqlParserUtil)
    (net.sf.jsqlparser.schema
@@ -12,6 +13,12 @@
 
 (set! *warn-on-reflection* true)
 
+(def callback-keys
+  "keyword->key map for the AST-folding callbacks."
+  ;; TODO: Move this to a Malli schema to simplify the indirection
+  {:column AstWalker$CallbackKey/COLUMN
+   :table  AstWalker$CallbackKey/TABLE})
+
 (defn- walk-query [parsed-query callbacks init-val]
   (.walk (AstWalker. callbacks init-val) parsed-query))
 
@@ -22,8 +29,8 @@
   the query'; this function doesn't do additional inference work to find out a table's schema.)"
   [^Statement parsed-query]
   (walk-query parsed-query
-              {:column #(update %1 :columns conj (.getColumnName ^Column %2))
-               :table  #(update %1 :tables conj (.getName ^Table %2))}
+              {(:column callback-keys) #(update %1 :columns conj (.getColumnName ^Column %2))
+               (:table callback-keys)  #(update %1 :tables conj (.getName ^Table %2))}
               {:columns #{}
                :tables  #{}}))