From e82ea85b07ffed71f1ab6505149e0d2e683fa542 Mon Sep 17 00:00:00 2001
From: Chris Truter <chris@metabase.com>
Date: Mon, 8 Jul 2024 15:18:07 +0200
Subject: [PATCH 1/2] Support MSSQL square bracket quotes

---
 src/macaw/collect.clj    |  6 ++++--
 src/macaw/core.clj       | 22 ++++++++++++++++++++--
 test/macaw/core_test.clj |  9 +++++----
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/src/macaw/collect.clj b/src/macaw/collect.clj
index 9b76528..d130e2e 100644
--- a/src/macaw/collect.clj
+++ b/src/macaw/collect.clj
@@ -47,12 +47,14 @@
 
 ;;; tables
 
-(def ^:private quotes (map str [\` \"]))
+(def ^:private quotes (map str "`\"["))
+
+(def ^:private closing {"[" "]"})
 
 (defn- quoted? [s]
   (some (fn [q]
           (and (str/starts-with? s q)
-               (str/ends-with? s q)))
+               (str/ends-with? s (closing q q))))
         quotes))
 
 (defn- strip-quotes [s]
diff --git a/src/macaw/core.clj b/src/macaw/core.clj
index 6afdab3..30fa026 100644
--- a/src/macaw/core.clj
+++ b/src/macaw/core.clj
@@ -5,7 +5,9 @@
    [macaw.collect :as collect]
    [macaw.rewrite :as rewrite])
   (:import
-   (net.sf.jsqlparser.parser CCJSqlParserUtil)))
+   (java.util.function Consumer)
+   (net.sf.jsqlparser.parser CCJSqlParser CCJSqlParserUtil)
+   (net.sf.jsqlparser.parser.feature Feature)))
 
 (set! *warn-on-reflection* true)
 
@@ -19,6 +21,22 @@
 (defn- unescape-keywords [sql _keywords]
   (str/replace sql "____escaped____" ""))
 
+(def ^:private features
+  {:backslash-escape-char  Feature/allowBackslashEscapeCharacter
+   :complex-parsing        Feature/allowComplexParsing
+   :postgres-syntax        Feature/allowPostgresSpecificSyntax
+   :square-bracket-quotes  Feature/allowSquareBracketQuotation
+   :unsupported-statements Feature/allowUnsupportedStatements})
+
+(defn- ->Feature ^Feature [k]
+  (get features k))
+
+(defn- ->parser-fn ^Consumer [opts]
+  (reify Consumer
+    (accept [_this parser]
+      (doseq [[f ^boolean v] (:features opts)]
+        (.withFeature ^CCJSqlParser parser (->Feature f) v)))))
+
 (defn parsed-query
   "Main entry point: takes a string query and returns a `Statement` object that can be handled by the other functions."
   [^String query & {:as opts}]
@@ -30,7 +48,7 @@
   (-> query
       (str/replace #"\n{2,}" "\n")
       (escape-keywords (:non-reserved-words opts))
-      (CCJSqlParserUtil/parse)))
+      (CCJSqlParserUtil/parse (->parser-fn opts))))
 
 (defn query->components
   "Given a parsed query (i.e., a [subclass of] `Statement`) return a map with the elements found within it.
diff --git a/test/macaw/core_test.clj b/test/macaw/core_test.clj
index b83e901..87e9af6 100644
--- a/test/macaw/core_test.clj
+++ b/test/macaw/core_test.clj
@@ -19,7 +19,9 @@
     (every? true? xs)
     false))
 
-(def components     (comp m/query->components m/parsed-query))
+(defn components [sql & {:as opts}]
+  (m/query->components (m/parsed-query sql opts) opts))
+
 (def raw-components #(let [xs (empty %)] (into xs (keep :component) %)))
 (def columns        (comp raw-components :columns components))
 (def source-columns (comp :source-columns components))
@@ -88,7 +90,6 @@ from foo")
     (is (= #{{:table "core_user"}}
            (tables "SELECT * FROM (SELECT DISTINCT email FROM core_user) q;")))))
 
-
 (deftest tables-with-complex-aliases-issue-14-test
   (testing "With an alias that is also a table name"
     (is (= #{{:table "user"}
@@ -314,8 +315,8 @@ from foo")
 
 (deftest complicated-mutations-test
   ;; https://github.com/metabase/macaw/issues/18
-  #_  (is (= #{"delete" "insert"}
-             (mutations "WITH outdated_orders AS (
+  #_(is (= #{"delete" "insert"}
+           (mutations "WITH outdated_orders AS (
                        DELETE FROM orders
                        WHERE
                          date <= '2018-01-01'

From 3125e8a3779baf60a29ae50193dff49b08c9e352 Mon Sep 17 00:00:00 2001
From: Chris Truter <chris@metabase.com>
Date: Mon, 8 Jul 2024 17:19:23 +0200
Subject: [PATCH 2/2] Add test

---
 test/macaw/core_test.clj | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/test/macaw/core_test.clj b/test/macaw/core_test.clj
index 87e9af6..2b75eb1 100644
--- a/test/macaw/core_test.clj
+++ b/test/macaw/core_test.clj
@@ -604,6 +604,16 @@ from foo")
                              :columns {{:table "x" :column "final"} "y"}}
                             {:non-reserved-words [:final]})))))
 
+(deftest square-bracket-test
+  (testing "We can opt into allowing square brackets to quote things"
+    (is (=? {:tables  #{{:schema "s" :table "t"}}
+             :columns #{{:schema "s" :table "t" :column "f"}}}
+            (update-vals
+             (components "SELECT [f] FROM [s].[t]"
+                         {:features              {:square-bracket-quotes true}
+                          :preserve-identifiers? false})
+             raw-components)))))
+
 (comment
  (require 'hashp.core)
  (require 'virgil)