diff --git a/src/metabase/driver/generic_sql/query_processor/annotate.clj b/src/metabase/driver/generic_sql/query_processor/annotate.clj
index df66f9228a5051789af68c82422ddab7f33dae08..708d14fb99539fda6a9cc4a37d3ee69b8b538b87 100644
--- a/src/metabase/driver/generic_sql/query_processor/annotate.clj
+++ b/src/metabase/driver/generic_sql/query_processor/annotate.clj
@@ -17,14 +17,13 @@
* `:columns` ordered sequence of column names
* `:cols` ordered sequence of information about each column, such as `:base_type` and `:special_type`"
[query results]
- (let [column-names (get-column-names query results)]
- {:status :completed
- :row_count (count results)
- :data {:rows (->> results
- (map #(map % ; pull out the values in each result in the same order we got from get-column-names
- (map keyword column-names))))
- :columns (map uncastify column-names)
- :cols (get-column-info query column-names)}}))
+ (let [column-names (get-column-names query results)
+ column-name-kws (map keyword column-names)]
+ {:rows (->> results
+ (map (fn [row]
+ (map row column-name-kws))))
+ :columns (map uncastify column-names)
+ :cols (get-column-info query column-names)}))
(defn- order-columns
[query castified-field-names]
diff --git a/src/metabase/driver/mongo/query_processor.clj b/src/metabase/driver/mongo/query_processor.clj
index 65970cae37dbf65571d5ca8585d623a5c0b49b7b..024335a00a300381423b3595c410d773f84754e5 100644
--- a/src/metabase/driver/mongo/query_processor.clj
+++ b/src/metabase/driver/mongo/query_processor.clj
@@ -108,9 +108,9 @@
(defn aggregate
"Generate a Monger `aggregate` form."
[& forms]
- `(mc/aggregate *mongo-connection* ~*collection-name* [~@(when *constraints*
- [{$match *constraints*}])
- ~@(filter identity forms)]))
+ `(mc/aggregate ^DBApiLayer *mongo-connection* ~*collection-name* [~@(when *constraints*
+ [{$match *constraints*}])
+ ~@(filter identity forms)]))
(defn field-id->$string
"Given a FIELD-ID, return a `$`-qualified field name for use in a Mongo aggregate query, e.g. `\"$user_id\"`."
@@ -119,12 +119,12 @@
(defaggregation ["rows"]
- `(doall (with-collection *mongo-connection* ~*collection-name*
+ `(doall (with-collection ^DBApiLayer *mongo-connection* ~*collection-name*
~@(when *constraints* [`(find ~*constraints*)])
~@(mapcat apply-clause *query*))))
(defaggregation ["count"]
- `[{:count (mc/count *mongo-connection* ~*collection-name*
+ `[{:count (mc/count ^DBApiLayer *mongo-connection* ~*collection-name*
~*constraints*)}])
(defaggregation ["avg" field-id]
@@ -241,12 +241,10 @@
(let [field-name->field (sel :many :field->obj [Field :name] :table_id source_table)
column-keys (qp/order-columns {:query query} (keys (first results)))
column-names (map name column-keys)]
- {:row_count (count results)
- :status :completed
- :data {:columns column-names
- :cols (qp/get-column-info {:query query} column-names)
- :rows (map #(map % column-keys)
- results)}}))
+ {:columns column-names
+ :cols (qp/get-column-info {:query query} column-names)
+ :rows (map #(map % column-keys)
+ results)}))
;; ## CLAUSE APPLICATION 2.0
diff --git a/src/metabase/driver/query_processor.clj b/src/metabase/driver/query_processor.clj
index a234eccff6d620b3c0ccab90c5a69e3186e7052b..47c37615dd14f5eca4247bbb4ca5eee8be807ccd 100644
--- a/src/metabase/driver/query_processor.clj
+++ b/src/metabase/driver/query_processor.clj
@@ -16,15 +16,13 @@
(def ^:const empty-response
"An empty response dictionary to return when there's no query to run."
- {:status :completed
- :row_count 0
- :data {:rows [], :columns [], :cols []}})
+ {:rows [], :columns [], :cols []})
;; # DYNAMIC VARS
(def ^:dynamic *query*
- "The structured query we're currently processing, before any preprocessing occurs (i.e. the `:query` part of the API call body)"
+ "The query we're currently processing (i.e., the body of the query API call)."
nil)
(def ^:dynamic *disable-qp-logging*
@@ -130,7 +128,7 @@
(defn post-process-cumulative-sum
"Cumulative sum the values of the aggregate `Field` in RESULTS."
{:arglists '([query results])}
- [{cum-sum-field :cum_sum, :as query} {{rows :rows, cols :cols, :as data} :data, :as results}]
+ [{cum-sum-field :cum_sum, :as query} {rows :rows, cols :cols, :as results}]
(if-not cum-sum-field results
(let [ ;; Determine the index of the field we need to cumulative sum
cum-sum-field-index (->> cols
@@ -170,7 +168,22 @@
reverse
(filter (complement same-breakout-field-values-as-previous-row?))
reverse
- (assoc-in results [:data :rows])))))
+ (assoc results :rows)))))
+
+;; ### ADD-ROW-COUNT-AND-STATUS
+
+(defn add-row-count-and-status
+ "Wrap the results of a successfully processed query in the format expected by the frontend (add `row_count` and `status`)."
+ [results]
+ {:pre [(map? results)
+ (sequential? (:columns results))
+ (sequential? (:cols results))
+ (sequential? (:rows results))]}
+ {:row_count (count (:rows results))
+ :status :completed
+ :data results})
+
+;; ### POST-PROCESS
(defn post-process
"Apply post-processing steps to the RESULTS of a QUERY, such as applying cumulative sum."
@@ -179,7 +192,8 @@
:native results
:query (let [query (:query query)]
(->> results
- (post-process-cumulative-sum query)))))
+ (post-process-cumulative-sum query)
+ add-row-count-and-status))))
;; # COMMON ANNOTATION FNS
diff --git a/test/metabase/driver/query_processor_test.clj b/test/metabase/driver/query_processor_test.clj
index 3b25447c1177ebf5c9ce82d4d168c6f4ebc931bd..c19e4fa9222bf24a19c0b7445ff582d520d7d62e 100644
--- a/test/metabase/driver/query_processor_test.clj
+++ b/test/metabase/driver/query_processor_test.clj
@@ -648,3 +648,16 @@
{:source_table (id :users)
:breakout [(id :users :name)]
:aggregation ["cum_sum" (id :users :id)]})
+
+;; ### Cumulative sum w/ a different breakout field that requires grouping
+(qp-expect-with-all-drivers
+ {:columns (->columns "price" "id")
+ :cols [(venue-col :price)
+ (venue-col :id)]
+ :rows [[1 1211]
+ [2 4066]
+ [3 4681]
+ [4 5050]]}
+ {:source_table (id :venues)
+ :breakout [(id :venues :price)]
+ :aggregation ["cum_sum" (id :venues :id)]})