gistfile1.txt

(defn after-batch [{:keys [onyx.core/results onyx.core/task] :as event} lifecycle]
  (let [tree-new (mapv (fn [{:keys [root leaves] :as message}]
                         (assoc message
                                :leaves
                                (mapv (fn [leaf]
                                        (let [input-segment-context (:onyx/context (:message root))
                                              output-segment-context (:onyx/context (:message leaf))
                                              context-path {:path (conj (vec (:path input-segment-context)) task)}]
                                          (assoc-in leaf
                                                    [:message :onyx/context]
                                                    (merge input-segment-context
                                                           output-segment-context
                                                           context-path))))
                                      leaves)))
                       (:tree results))
        segments-new (doall (mapcat :leaves tree-new))]
    {:onyx.core/results (assoc results :tree tree-new :segments segments-new)}))

(def middleware-calls
  {:lifecycle/after-batch after-batch})