From 13cae0ebb158834b2604714aee7ba1de3c48ddee Mon Sep 17 00:00:00 2001
From: Alexander Solovyov <alexander@solovyov.net>
Date: Fri, 5 Apr 2024 16:53:40 +0300
Subject: [PATCH] account for group by columns

---
 .dir-locals.el                         | 12 ++++++++----
 java/com/metabase/macaw/AstWalker.java | 20 +++++++++++++++-----
 test/macaw/core_test.clj               |  5 ++++-
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/.dir-locals.el b/.dir-locals.el
index 8363bc4..09580a5 100644
--- a/.dir-locals.el
+++ b/.dir-locals.el
@@ -1,11 +1,15 @@
 ((nil . ((indent-tabs-mode . nil)       ; always use spaces for tabs
          (require-final-newline . t)))  ; add final newline on save
  (clojure-mode . ((cider-preferred-build-tool . clojure-cli)
-                  (cider-clojure-cli-aliases . "dev")
+                  (cider-clojure-cli-aliases . "dev:user")
                   (cljr-favor-prefix-notation . nil)
-                  (fill-column . 120)
-                  (column-enforce-column . 120)
-                  (clojure-docstring-fill-column . 120)
+                  ;; prefer keeping source width about ~118, GitHub seems to cut
+                  ;; off stuff at either 119 or 120 and it's nicer to look at
+                  ;; code in GH when you don't have to scroll back and forth
+                  (fill-column . 118)
+                  (whitespace-line-column . 118)
+                  (column-enforce-column . 118)
+                  (clojure-docstring-fill-column . 118)
                   (eval . (put-clojure-indent 'with-meta '(:form)))
                   (eval . (put-clojure-indent 'with-bindings* '(:form)))))
  (markdown-mode . ((fill-column . 80)
diff --git a/java/com/metabase/macaw/AstWalker.java b/java/com/metabase/macaw/AstWalker.java
index 8ae207c..5052665 100644
--- a/java/com/metabase/macaw/AstWalker.java
+++ b/java/com/metabase/macaw/AstWalker.java
@@ -146,6 +146,8 @@
 import net.sf.jsqlparser.statement.refresh.RefreshMaterializedViewStatement;
 import net.sf.jsqlparser.statement.select.AllColumns;
 import net.sf.jsqlparser.statement.select.AllTableColumns;
+import net.sf.jsqlparser.statement.select.GroupByElement;
+import net.sf.jsqlparser.statement.select.GroupByVisitor;
 import net.sf.jsqlparser.statement.select.FromItemVisitor;
 import net.sf.jsqlparser.statement.select.Join;
 import net.sf.jsqlparser.statement.select.LateralSubSelect;
@@ -209,7 +211,7 @@
  * map and functions to implement the necessary behavior; no `reify` necesary.
  */
 public class AstWalker<Acc> implements SelectVisitor, FromItemVisitor, ExpressionVisitor,
-       SelectItemVisitor, StatementVisitor {
+       SelectItemVisitor, StatementVisitor, GroupByVisitor {
 
     public enum CallbackKey {
         ALL_COLUMNS,
@@ -343,6 +345,10 @@ public void visit(PlainSelect plainSelect) {
         if (plainSelect.getOracleHierarchical() != null) {
             plainSelect.getOracleHierarchical().accept(this);
         }
+
+        if (plainSelect.getGroupBy() != null) {
+            plainSelect.getGroupBy().accept(this);
+        }
     }
 
     @Override
@@ -799,7 +805,6 @@ public void visit(UserVariable var) {
     @Override
     public void visit(NumericBind bind) {
 
-
     }
 
     @Override
@@ -967,7 +972,6 @@ public void visit(RowGetExpression rowGetExpression) {
     @Override
     public void visit(HexValue hexValue) {
 
-
     }
 
     @Override
@@ -1012,13 +1016,11 @@ public void visit(TimeKeyExpression timeKeyExpression) {
     @Override
     public void visit(DateTimeLiteralExpression literal) {
 
-
     }
 
     @Override
     public void visit(Commit commit) {
 
-
     }
 
     @Override
@@ -1044,6 +1046,14 @@ public void visit(ParenthesedFromItem parenthesis) {
         visitJoins(parenthesis.getJoins());
     }
 
+    @Override
+    public void visit(GroupByElement element) {
+        element.getGroupByExpressionList().accept(this);
+        for (ExpressionList exprList : element.getGroupingSets()) {
+            exprList.accept(this);
+        }
+    }
+
     /**
      * visit join block
      *
diff --git a/test/macaw/core_test.clj b/test/macaw/core_test.clj
index 249731f..71add36 100644
--- a/test/macaw/core_test.clj
+++ b/test/macaw/core_test.clj
@@ -36,7 +36,10 @@
 (deftest query->columns-test
   (testing "Simple queries"
     (is (= #{"foo" "bar" "id" "quux_id"}
-           (columns "select foo, bar from baz inner join quux on quux.id = baz.quux_id")))))
+           (columns "select foo, bar from baz inner join quux on quux.id = baz.quux_id"))))
+  (testing "'group by' columns present"
+    (is (= #{"id" "user_id"}
+           (columns "select id from orders group by user_id")))))
 
 (deftest alias-inclusion-test
   (testing "Aliases are not included"