metabase.query-processor.middleware.add-source-metadata

(ns metabase.query-processor.middleware.add-source-metadata
  (:require
   [clojure.walk :as walk]
   [metabase.legacy-mbql.schema :as mbql.s]
   [metabase.lib.metadata :as lib.metadata]
   [metabase.lib.util.match :as lib.util.match]
   [metabase.query-processor.interface :as qp.i]
   [metabase.query-processor.store :as qp.store]
   [metabase.request.core :as request]
   [metabase.util.log :as log]
   [metabase.util.malli :as mu]))

(defn- has-same-fields-as-nested-source?
  [{nested-source-query    :source-query
    nested-source-metadata :source-metadata
    breakouts              :breakout
    aggregations           :aggregation
    fields                 :fields}]
  (when nested-source-query
    (and (every? empty? [breakouts aggregations])
         (or (empty? fields)
             (and (= (count fields) (count nested-source-metadata))
                  (every? #(lib.util.match/match-one % [:field (_ :guard string?) _])
                          fields))))))

(mu/defn- native-source-query->metadata :- [:maybe [:sequential mbql.s/SourceQueryMetadata]]
  "Given a `source-query`, return the source metadata that should be added at the parent level (i.e., at the same
  level where this `source-query` was present.) This metadata is used by other middleware to determine what Fields to
  expect from the source query."
  [{nested-source-metadata :source-metadata, :as source-query} :- mbql.s/SourceQuery]
  ;; If the source query has a nested source with metadata and does not change the fields that come back, return
  ;; metadata as-is
  (if (has-same-fields-as-nested-source? source-query)
    nested-source-metadata
    ;; Otherwise we cannot determine the metadata automatically; usually, this is because the source query itself has
    ;; a native source query
    (do
      (when-not qp.i/*disable-qp-logging*
        (log/warn "Cannot infer `:source-metadata` for source query with native source query without source metadata."
                  {:source-query source-query}))
      nil)))

(mu/defn mbql-source-query->metadata :- [:maybe [:sequential mbql.s/SourceQueryMetadata]]
  "Preprocess a `source-query` so we can determine the result columns."
  [source-query :- mbql.s/MBQLQuery]
  (try
    (let [cols (request/as-admin
                 ((requiring-resolve 'metabase.query-processor.preprocess/query->expected-cols)
                  {:database (:id (lib.metadata/database (qp.store/metadata-provider)))
                   :type     :query
                   ;; don't add remapped columns to the source metadata for the source query, otherwise we're going
                   ;; to end up adding it again when the middleware runs at the top level
                   :query    (assoc-in source-query [:middleware :disable-remaps?] true)}))]
      (for [col cols
            :when (not (:remapped_from col))]
        (select-keys col [:name :id :table_id :display_name :base_type :effective_type :coercion_strategy
                          :semantic_type :unit :fingerprint :settings :source_alias :field_ref :nfc_path :parent_id])))
    (catch Throwable e
      (log/errorf e "Error determining expected columns for query: %s" (ex-message e))
      nil)))

(mu/defn- add-source-metadata :- [:map
                                  [:source-metadata
                                   {:optional true}
                                   [:maybe [:sequential mbql.s/SourceQueryMetadata]]]]
  [{{native-source-query? :native, :as source-query} :source-query, :as inner-query} :- :map]
  (let [metadata ((if native-source-query?
                    native-source-query->metadata
                    mbql-source-query->metadata) source-query)]
    (cond-> inner-query
      (seq metadata) (assoc :source-metadata metadata))))

(defn- legacy-source-metadata?
  [source-metadata]
  (and (seq source-metadata)
       (every? nil? (map :field_ref source-metadata))))

(defn- should-add-source-metadata?
  [{{native-source-query?              :native
     source-query-has-source-metadata? :source-metadata
     :as                               source-query} :source-query
    :keys                                            [source-metadata]}]
  (and source-query
       (or (not source-metadata)
           (legacy-source-metadata? source-metadata))
       (or (not native-source-query?)
           source-query-has-source-metadata?)))

(defn- maybe-add-source-metadata [x]
  (if (and (map? x) (should-add-source-metadata? x))
    (add-source-metadata x)
    x))

(defn- add-source-metadata-at-all-levels [inner-query]
  (walk/postwalk maybe-add-source-metadata inner-query))

(defn add-source-metadata-for-source-queries
  [{query-type :type, :as query}]
  (if-not (= query-type :query)
    query
    (update query :query add-source-metadata-at-all-levels)))