From a1b4e93a789b537572d2df44f01e61d177bbcd67 Mon Sep 17 00:00:00 2001
From: metamben <103100869+metamben@users.noreply.github.com>
Date: Mon, 15 May 2023 17:48:42 +0300
Subject: [PATCH] Keep original join-aliases in expected columns (#30747)

* Keep original join-aliases in expected columns

Fixes #30648.

The root cause of the bug was that the result_metadata generated for MBQL
queries returned escaped join-aliases, which in BigQueries case means that
spaces are converted underscores (_). These wrong join aliases could not be
resolved in later stages and resulted in incorrect references.

Normally, the escaped aliases are replaced by original ones in the post
processing step, but the post processing is not executed when calculating
expected columns.

This PR makes sure the original aliases are restored in the expected columns.

* Keep source_alias and display_name intact too
---
 src/metabase/query_processor.clj              |  8 +++-
 .../util/add_alias_info_test.clj              | 39 +++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/metabase/query_processor.clj b/src/metabase/query_processor.clj
index 2749cf4cbba..02afbb27c34 100644
--- a/src/metabase/query_processor.clj
+++ b/src/metabase/query_processor.clj
@@ -307,6 +307,10 @@
               (preprocess* query)))]
     (qp query nil nil)))
 
+(defn- restore-join-aliases [preprocessed-query]
+  (let [replacement (-> preprocessed-query :info :alias/escaped->original)]
+    (escape-join-aliases/restore-aliases preprocessed-query replacement)))
+
 (defn query->expected-cols
   "Return the `:cols` you would normally see in MBQL query results by preprocessing the query and calling `annotate` on
   it. This only works for pure MBQL queries, since it does not actually run the queries. Native queries or MBQL
@@ -314,11 +318,11 @@
   [{query-type :type, :as query}]
   (when-not (= (mbql.u/normalize-token query-type) :query)
     (throw (ex-info (tru "Can only determine expected columns for MBQL queries.")
-             {:type qp.error-type/qp})))
+                    {:type qp.error-type/qp})))
   ;; TODO - we should throw an Exception if the query has a native source query or at least warn about it. Need to
   ;; check where this is used.
   (qp.store/with-store
-    (let [preprocessed (preprocess query)]
+    (let [preprocessed (-> query preprocess restore-join-aliases)]
       (driver/with-driver (driver.u/database->driver (:database preprocessed))
         (not-empty (vec (annotate/merged-column-info preprocessed nil)))))))
 
diff --git a/test/metabase/query_processor/util/add_alias_info_test.clj b/test/metabase/query_processor/util/add_alias_info_test.clj
index a37eaf40cd0..28238422004 100644
--- a/test/metabase/query_processor/util/add_alias_info_test.clj
+++ b/test/metabase/query_processor/util/add_alias_info_test.clj
@@ -570,6 +570,45 @@
                             add-alias-info
                             :query)))))))))
 
+(deftest query->expected-cols-test
+  (testing "field_refs in expected columns have the original join aliases (#30648)"
+    (mt/dataset sample-dataset
+      (binding [driver/*driver* ::custom-escape-spaces-to-underscores]
+        (let [query
+              (mt/mbql-query
+               products
+                {:joins
+                 [{:source-query
+                   {:source-table $$orders
+                    :joins
+                    [{:source-table $$people
+                      :alias "People"
+                      :condition [:= $orders.user_id &People.people.id]
+                      :fields [&People.people.address]
+                      :strategy :left-join}]
+                    :fields [$orders.id &People.people.address]}
+                   :alias "Question 54"
+                   :condition [:= $id [:field %orders.id {:join-alias "Question 54"}]]
+                   :fields [[:field %orders.id {:join-alias "Question 54"}]
+                            [:field %people.address {:join-alias "Question 54"}]]
+                   :strategy :left-join}]
+                 :fields
+                 [!default.created_at
+                  [:field %orders.id {:join-alias "Question 54"}]
+                  [:field %people.address {:join-alias "Question 54"}]]})]
+          (is (=? [{:name "CREATED_AT"
+                    :field_ref [:field (mt/id :products :created_at) {:temporal-unit :default}]
+                    :display_name "Created At"}
+                   {:name "ID"
+                    :field_ref [:field (mt/id :orders :id) {:join-alias "Question 54"}]
+                    :display_name "Question 54 → ID"
+                    :source_alias "Question 54"}
+                   {:name "ADDRESS"
+                    :field_ref [:field (mt/id :people :address) {:join-alias "Question 54"}]
+                    :display_name "Question 54 → Address"
+                    :source_alias "Question 54"}]
+                  (qp/query->expected-cols query))))))))
+
 (deftest use-source-unique-aliases-test
   (testing "Make sure uniquified aliases in the source query end up getting used for `::add/source-alias`"
     ;; keep track of the IDs so we don't accidentally fetch the wrong ones after we switch the name of `price`
-- 
GitLab