From 9491bf03fa6c33e59ae92479d0fb54a9f13d4c92 Mon Sep 17 00:00:00 2001
From: Alexander Solovyov <alexander@solovyov.net>
Date: Thu, 9 May 2024 18:37:51 +0300
Subject: [PATCH] Infer a table for column if there is only one table mentioned
 (#29)

---
 build.clj                              |  2 +-
 deps.edn                               |  3 ++-
 java/com/metabase/macaw/AstWalker.java |  2 ++
 src/macaw/core.clj                     |  7 +++++--
 test/macaw/core_test.clj               | 18 +++++++++---------
 5 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/build.clj b/build.clj
index 2460f31..64b4621 100644
--- a/build.clj
+++ b/build.clj
@@ -69,7 +69,7 @@
   (b/javac (merge default-options
                   opts
                   {:src-dirs   ["java"]
-                   :javac-opts ["--release" "11"]})))
+                   :javac-opts ["--release" "11" "-Xlint:-options" #_"-Xlint:unchecked"]})))
 
 (defn jar [opts]
   (println "\nStarting to build a JAR...")
diff --git a/deps.edn b/deps.edn
index d8a74a5..082984c 100644
--- a/deps.edn
+++ b/deps.edn
@@ -13,7 +13,8 @@
  :aliases
  {:dev
   {:extra-deps
-   {io.github.metabase/hawk {:sha "539eefaa31a43d52d7c9b5731f471bb6742e7131"}}
+   {io.github.metabase/hawk {:sha "539eefaa31a43d52d7c9b5731f471bb6742e7131"}
+    virgil/virgil           {:mvn/version "0.3.0"}}
 
    :extra-paths
    ["test"]
diff --git a/java/com/metabase/macaw/AstWalker.java b/java/com/metabase/macaw/AstWalker.java
index 5fa6d9b..a12b1ff 100644
--- a/java/com/metabase/macaw/AstWalker.java
+++ b/java/com/metabase/macaw/AstWalker.java
@@ -268,6 +268,7 @@ public AstWalker(Map<CallbackKey, IFn> rawCallbacks, Acc val) {
     /**
      * Safely invoke the given callback by name.
      */
+    @SuppressWarnings("unchecked")
     public void invokeCallback(CallbackKey key, Object visitedItem) {
         IFn callback = this.callbacks.get(key);
         if (callback != null) {
@@ -882,6 +883,7 @@ public void visit(Delete delete) {
         popContext(); // DELETE
     }
 
+    @SuppressWarnings("deprecation")
     @Override
     public void visit(Update update) {
         pushContext(UPDATE);
diff --git a/src/macaw/core.clj b/src/macaw/core.clj
index 225fbf3..3dc8d41 100644
--- a/src/macaw/core.clj
+++ b/src/macaw/core.clj
@@ -43,10 +43,13 @@
 (defn- make-column [alias-map table-map ^Column c]
   (merge
     {:column (.getColumnName c)}
-    (when-let [t (.getTable c)]
+    (if-let [t (.getTable c)]
       (or
         (get alias-map (.getName t))
-        (:component (get table-map (.getName t)))))))
+        (:component (get table-map (.getName t))))
+      ;; if we see only a single table, we can safely say it's the table of that column
+      (when (= (count table-map) 1)
+        (:component (val (first table-map)))))))
 
 (defn- alias-mapping
   [^Table table]
diff --git a/test/macaw/core_test.clj b/test/macaw/core_test.clj
index ac12e81..1e23cb1 100644
--- a/test/macaw/core_test.clj
+++ b/test/macaw/core_test.clj
@@ -69,8 +69,8 @@
              {:column "quux_id" :table "baz"}}
            (columns "SELECT foo, bar FROM baz INNER JOIN quux ON quux.id = baz.quux_id"))))
   (testing "'group by' columns present"
-    (is (= #{{:column "id"}
-             {:column "user_id"}}
+    (is (= #{{:column "id" :table "orders"}
+             {:column "user_id" :table "orders"}}
            (columns "SELECT id FROM orders GROUP BY user_id"))))
   (testing "table alias present"
     (is (= #{{:column "id" :table "orders" :schema "public"}}
@@ -165,9 +165,9 @@
 (deftest context-test
   (testing "Sub-select with outer wildcard"
     (is (= {:columns
-            #{{:component {:column "total"}, :context ["SELECT" "SUB_SELECT" "FROM" "SELECT"]}
-              {:component {:column "id"},    :context ["SELECT" "SUB_SELECT" "FROM" "SELECT"]}
-              {:component {:column "total"}, :context ["WHERE" "JOIN" "FROM" "SELECT"]}},
+            #{{:component {:column "total" :table "orders"}, :context ["SELECT" "SUB_SELECT" "FROM" "SELECT"]}
+              {:component {:column "id"    :table "orders"}, :context ["SELECT" "SUB_SELECT" "FROM" "SELECT"]}
+              {:component {:column "total" :table "orders"}, :context ["WHERE" "JOIN" "FROM" "SELECT"]}},
             :has-wildcard?     #{{:component true, :context ["SELECT"]}},
             :mutation-commands #{},
             :tables            #{{:component {:table "orders"}, :context ["FROM" "SELECT" "SUB_SELECT" "FROM" "SELECT"]}},
@@ -175,16 +175,16 @@
            (components "SELECT * FROM (SELECT id, total FROM orders) WHERE total > 10"))))
   (testing "Sub-select with inner wildcard"
     (is (= {:columns
-            #{{:component {:column "id"},    :context ["SELECT"]}
-              {:component {:column "total"}, :context ["SELECT"]}
-              {:component {:column "total"}, :context ["WHERE" "JOIN" "FROM" "SELECT"]}},
+            #{{:component {:column "id"    :table "orders"}, :context ["SELECT"]}
+              {:component {:column "total" :table "orders"}, :context ["SELECT"]}
+              {:component {:column "total" :table "orders"}, :context ["WHERE" "JOIN" "FROM" "SELECT"]}},
             :has-wildcard?     #{{:component true, :context ["SELECT" "SUB_SELECT" "FROM" "SELECT"]}},
             :mutation-commands #{},
             :tables            #{{:component {:table "orders"}, :context ["FROM" "SELECT" "SUB_SELECT" "FROM" "SELECT"]}},
             :table-wildcards   #{}}
            (components "SELECT id, total FROM (SELECT * FROM orders) WHERE total > 10"))))
   (testing "Sub-select with dual wildcards"
-    (is (= {:columns           #{{:component {:column "total"}, :context ["WHERE" "JOIN" "FROM" "SELECT"]}},
+    (is (= {:columns           #{{:component {:column "total" :table "orders"}, :context ["WHERE" "JOIN" "FROM" "SELECT"]}},
             :has-wildcard?
             #{{:component true, :context ["SELECT" "SUB_SELECT" "FROM" "SELECT"]}
               {:component true, :context ["SELECT"]}},