Fix nested query issue with SparkSQL driver [ci drivers]

Previously the driver was only looking one and two levels deep for the
source table. In the case of nested and binned queries, this could be
buried deeper in the MBQL query. This commit switches form an `or`
clause and a couple of `get-in`s to a `postwalk` that will find the
table however it's nested in the query.

src/metabase/driver/sparksql.clj

@@ -15,6 +15,7 @@
              [generic-sql :as sql]
              [hive-like :as hive-like]]
             [metabase.driver.generic-sql.query-processor :as sqlqp]
+            [metabase.models.table :refer [Table]]
             [metabase.query-processor.util :as qputil]
             [metabase.util.honeysql-extensions :as hx]
             [puppetlabs.i18n.core :refer [trs]])
@@ -32,9 +33,13 @@
 
 (def ^:private source-table-alias "t1")
 
+(defn- find-source-table [query]
+  (first (qputil/postwalk-collect #(instance? (type Table) %)
+                                  identity
+                                  query)))
+
 (defn- resolve-table-alias [{:keys [schema-name table-name special-type field-name] :as field}]
-  (let [source-table (or (get-in sqlqp/*query* [:query :source-table])
-                         (get-in sqlqp/*query* [:query :source-query :source-table]))]
+  (let [source-table (find-source-table sqlqp/*query*)]
     (if (and (= schema-name (:schema source-table))
              (= table-name (:name source-table)))
       (-> (assoc field :schema-name nil)
@@ -49,7 +54,7 @@
 
 (defmethod  sqlqp/->honeysql [SparkSQLDriver Field]
   [driver field-before-aliasing]
-  (let [{:keys [schema-name table-name special-type field-name]} (resolve-table-alias field-before-aliasing)
+  (let [{:keys [schema-name table-name special-type field-name] :as foo} (resolve-table-alias field-before-aliasing)
         field (keyword (hx/qualify-and-escape-dots schema-name table-name field-name))]
     (cond
       (isa? special-type :type/UNIXTimestampSeconds)      (sql/unix-timestamp->timestamp driver field :seconds)

src/metabase/query_processor/util.clj

@@ -4,7 +4,9 @@
              [codecs :as codecs]
              [hash :as hash]]
             [cheshire.core :as json]
-            [clojure.string :as str]
+            [clojure
+             [string :as str]
+             [walk :as walk]]
             [metabase.util :as u]
             [metabase.util.schema :as su]
             [schema.core :as s]))
@@ -138,3 +140,30 @@
     (when (string? source-table)
       (when-let [[_ card-id-str] (re-matches #"^card__(\d+$)" source-table)]
         (Integer/parseInt card-id-str)))))
+
+;;; ---------------------------------------- General Tree Manipulation Helpers ---------------------------------------
+
+(defn postwalk-pred
+  "Transform `form` by applying `f` to each node where `pred` returns true"
+  [pred f form]
+  (walk/postwalk (fn [node]
+                   (if (pred node)
+                     (f node)
+                     node))
+                 form))
+
+(defn postwalk-collect
+  "Invoke `collect-fn` on each node satisfying `pred`. If `collect-fn` returns a value, accumulate that and return the
+  results.
+
+  Note: This would be much better as a zipper. It could have the same API, would be faster and would avoid side
+  affects."
+  [pred collect-fn form]
+  (let [results (atom [])]
+    (postwalk-pred pred
+                   (fn [node]
+                     (when-let [result (collect-fn node)]
+                       (swap! results conj result))
+                     node)
+                   form)
+    @results))

test/metabase/query_processor/util_test.clj

@@ -167,3 +167,43 @@
 (expect
   nil
   (qputil/dissoc-normalized nil :num-toucans))
+
+(defrecord ^:private TestRecord1 [x])
+(defrecord ^:private TestRecord2 [x])
+
+(def ^:private test-tree
+  {:a {:aa (TestRecord1. 1)
+       :ab (TestRecord2. 1)}
+   :b (TestRecord1. 1)
+   :c (TestRecord2. 1)
+   :d [1 2 3 4]})
+
+;; Test that we can change only the items matching the `instance?` predicate
+(expect
+  (-> test-tree
+      (update-in [:a :aa :x] inc)
+      (update-in [:b :x] inc))
+  (qputil/postwalk-pred #(instance? TestRecord1 %)
+                        #(update % :x inc)
+                        test-tree))
+
+;; If nothing matches, the original tree should be returned
+(expect
+  test-tree
+  (qputil/postwalk-pred set?
+                        #(set (map inc %))
+                        test-tree))
+
+;; We should be able to collect items matching the predicate
+(expect
+  [(TestRecord1. 1) (TestRecord1. 1)]
+  (qputil/postwalk-collect #(instance? TestRecord1 %)
+                           identity
+                           test-tree))
+
+;; Not finding any of the items should just return an empty seq
+(expect
+  []
+  (qputil/postwalk-collect set?
+                           identity
+                           test-tree))