Merge pull request #3554 from rbrw/pdb-5126-add-query-origin

(PDB-5126) Add support for query "origin"

Author: austb

documentation/api/query/v4/index.markdown

@@ -37,6 +37,12 @@ general info about queries, see [our guide to query structure.][query]
 * `ast_only`: optional. A boolean value. When true, the query response will be the supplied 
 `query` in AST, either exactly as supplied or translated from PQL. False by default.
 
+* `origin`: optional. A string describing the source of the query.  It
+can be anything, and will be reported in the log when PuppetDB is
+configured to log queries.  Note that Puppet intends to use origin
+names beginning with `puppet:` for its own queries, so it is
+recommended that other clients choose something else.
+
 ### Response format
 
 The response will be in `application/json`, and will contain a list of JSON

src/puppetlabs/puppetdb/http/handlers.clj

@@ -68,7 +68,7 @@
 (def global-engine-params
   "Parameters that should always be forwarded from the incoming query to
   the engine."
-  [:optimize_drop_unused_joins :include_facts_expiration :explain])
+  [:optimize_drop_unused_joins :include_facts_expiration :explain :origin])
 
 (defn status-response
   "Executes `query` and if a result is found, calls `found-fn` with
@@ -148,7 +148,8 @@
 (def global-params {:optional ["optimize_drop_unused_joins"
                                "include_facts_expiration"
                                "include_package_inventory"
-                               "explain"]})
+                               "explain"
+                               "origin"]})
 (def paging-params {:optional paging/query-params})
 (def pretty-params {:optional ["pretty"]})
 (def typical-params (merge-param-specs global-params

src/puppetlabs/puppetdb/http/query.clj

@@ -46,6 +46,7 @@
    (s/optional-key :include_package_inventory) (s/maybe s/Bool)
    (s/optional-key :order_by) (s/maybe [[(s/one s/Keyword "field")
                                          (s/one (s/enum :ascending :descending) "order")]])
+   (s/optional-key :origin) (s/maybe s/Str)
    (s/optional-key :distinct_resources) (s/maybe s/Bool)
    (s/optional-key :distinct_start_time) s/Any
    (s/optional-key :distinct_end_time) s/Any

src/puppetlabs/puppetdb/query_eng.clj

@@ -417,7 +417,7 @@
                :subquery? false}))
 
 (def fact-paths-query
-  "Query for the resource-params query, mostly used as a subquery"
+  "Query for the fact-paths query, mostly used as a subquery"
   (map->Query {::which-query :fact-paths
                :can-drop-unused-joins? true
                :projections {"type" {:type :string

src/puppetlabs/puppetdb/query_eng/engine.clj

@@ -17,7 +17,8 @@
    [java.net MalformedURLException URISyntaxException URL]
    [java.nio ByteBuffer CharBuffer]
    [java.nio.charset Charset CoderResult StandardCharsets]
-   (java.util.concurrent ScheduledThreadPoolExecutor TimeUnit)))
+   (java.util.concurrent ScheduledThreadPoolExecutor TimeUnit)
+   (org.apache.log4j MDC)))
 
 (defmacro with-captured-throw [& body]
   `(try [(do ~@body)] (catch Throwable ex# ex#)))
@@ -34,6 +35,36 @@
     (apply print args)
     (flush)))
 
+(defmacro with-log-mdc
+  "Establishes the MDC contexts given by the alternating-kvs key value
+  pairs during the execution of the body, and ensures that the
+  original values (if any) are always restored before returning."
+  [alternating-kvs & body]
+  ;; For now, assume this isn't used in performance-critical code,
+  ;; i.e. within tight loops.
+  (when-not (even? (count alternating-kvs))
+    (throw (RuntimeException. "Odd number of MDC key value pairs")))
+  (when-not (every? string? (take-nth 2 alternating-kvs))
+    (throw (RuntimeException. "MDC keys are not all strings")))
+  (loop [[k v & alternating-kvs] alternating-kvs
+         expansion `(do ~@body)]
+    (if-not k
+      expansion
+      (recur alternating-kvs
+             ;; We know k is a string, so it's fine to repeat ~k
+             `(let [v# ~v]
+                (if (nil? v#)
+                  ~expansion
+                  (let [orig# (MDC/get ~k)]
+                    (try
+                      (MDC/put ~k v#)
+                      ~expansion
+                      (finally
+                        ;; After you put a nil value MDC/getContext crashes
+                        (if (nil? orig#)
+                          (MDC/remove ~k)
+                          (MDC/put ~k orig#)))))))))))
+
 (defn flush-and-exit
   "Attempts to flush *out* and *err*, reporting any failures to *err*,
   if possible, and then invokes (System/exit status)."

src/puppetlabs/puppetdb/utils.clj

@@ -1,16 +1,29 @@
 (ns puppetlabs.puppetdb.http.query-logging-test
-  (:require  [clojure.test :refer :all]
-             [clojure.string :as str]
-             [puppetlabs.kitchensink.core :as kitchensink]
-             [puppetlabs.puppetdb.testutils.http :refer [query-response]]
-             [puppetlabs.puppetdb.testutils.db :refer [with-test-db *db*]]
-             [puppetlabs.puppetdb.testutils.http :refer [call-with-http-app]]
-             [puppetlabs.puppetdb.testutils.services :refer [call-with-puppetdb-instance
-                                                             create-temp-config
-                                                             *server*]]
-             [puppetlabs.puppetdb.cli.services :as svcs]
-             [puppetlabs.trapperkeeper.app :refer [get-service]]
-             [puppetlabs.trapperkeeper.testutils.logging :as tk-log]))
+  (:require
+   [clojure.java.io :as io]
+   [clojure.string :as str]
+   [clojure.test :refer :all]
+   [clojure.walk :refer [keywordize-keys]]
+   [puppetlabs.kitchensink.core :as kitchensink]
+   [puppetlabs.puppetdb.cheshire :as json]
+   [puppetlabs.puppetdb.cli.services :as svcs]
+   [puppetlabs.puppetdb.query-eng :as qeng]
+   [puppetlabs.puppetdb.testutils.catalogs :refer [replace-catalog]]
+   [puppetlabs.puppetdb.testutils.db :refer [with-test-db *db*]]
+   [puppetlabs.puppetdb.testutils.http
+    :refer [call-with-http-app query-response with-http-app*]]
+   [puppetlabs.puppetdb.testutils.log :refer [notable-pdb-event?]]
+   [puppetlabs.puppetdb.testutils.services
+    :refer [call-with-puppetdb-instance create-temp-config *server*]]
+   [puppetlabs.puppetdb.time :as time]
+   [puppetlabs.puppetdb.utils :refer [println-err]]
+   [puppetlabs.trapperkeeper.app :refer [get-service]]
+   [puppetlabs.trapperkeeper.testutils.logging :as tk-log
+    :refer [with-log-suppressed-unless-notable
+            with-logged-event-maps
+            with-logging-to-atom]])
+  (:import
+   (java.util UUID)))
 
 (defn logs-include?
   "Returns true if only one instance of unique-msg is found in the log."
@@ -37,29 +50,60 @@
 (defn prep-logs [logs]
   (->> @logs (map :message) keep-only-pdbquery-logs))
 
-(deftest setting-log-queries-triggers-ast-sql-logging
-  (tk-log/with-logged-event-maps logs
-    (tk-log/with-log-level "puppetlabs.puppetdb.query-eng" :debug
-      (with-test-db
-        (call-with-http-app
-         (fn []
-           ;; make a couple http queries to trigger a debug AST and SQL log message for each
-           (is (= 200 (:status (query-response :get "/v4" ["from" "nodes"]))))
-           (is (= 200 (:status (query-response :get "/v4" ["from" "facts"]))))
+(def catalog-1
+  (-> "puppetlabs/puppetdb/cli/export/tiny-catalog.json"
+      io/resource slurp json/parse-string keywordize-keys))
+
+(deftest queries-are-logged-when-log-queries-is-true
+  (tk-log/with-log-level "puppetlabs.puppetdb.query-eng" :debug
+    (with-test-db
+      (replace-catalog catalog-1)
+      (with-http-app* #(assoc % :log-queries true)
+        (doseq [[query exp-ast exp-sql exp-origin]
+                ;; produce-streaming-body
+                [[["/v4" ["from" "nodes"] {:origin "foo"}]
+                  "\"from\",\"nodes\""
+                  "latest_report_noop_pending"
+                  "foo"]
+                 [["/v4" ["from" "facts"]]
+                  "\"from\",\"facts\""
+                  "(jsonb_each((stable||volatile)))"
+                  nil]
+                 ;; stream-query-result
+                 [["/v4/catalogs/myhost.localdomain" [] {:origin "bar"}]
+                  "\"from\",\"catalogs\""
+                  "row_to_json(edge_data)"
+                  "bar"]]]
+          (with-logged-event-maps events
+            (is (= 200 (:status (apply query-response :get query))))
+
+            (let [events @events
+                  ;; Returns [everything uuid query-info]
+                  parse-event-msg #(->> % :message (re-find #"^PDBQuery:([^:]+):(.*)"))
+                  parse-event-info #(-> % parse-event-msg (nth 2) json/parse-string)
+                  uuid (some->> events (some parse-event-msg) second)
+                  qev-matching (fn [expected]
+                                 (fn [{:keys [message] :as event}]
+                                   (and (str/starts-with? message (str "PDBQuery:" uuid ":"))
+                                        (str/includes? message expected))))]
 
-           (let [logs (prep-logs logs)]
-             (testing "uuids match for the AST and SQL logged per query"
-               (is (= [2 2] (vals (count-logs-uuids logs)))))
+              (is (uuid? (UUID/fromString uuid)))
 
-             (testing "AST/SQL is logged for both queries above"
-               ;; match the AST/SQL logs for nodes query
-               (is (logs-include? logs "\"from\" \"nodes\""))
-               (is (logs-include? logs "latest_report_noop_pending"))
+              (let [[ev & evs] (filter (qev-matching exp-ast) events)]
+                (is (not (seq evs)))
+                (when (seq evs)
+                  (println-err "Unexpected log:" events))
+                (let [{:strs [ast origin] :as info} (parse-event-info ev)]
+                  (is ast)
+                  (is (= exp-origin origin))))
 
-               ;; match the AST/SQL logs for facts query
-               (is (logs-include? logs "\"from\" \"facts\""))
-               (is (logs-include? logs "(jsonb_each((stable||volatile)))")))))
-         #(assoc % :log-queries true))))))
+              (let [[ev & evs] (filter (qev-matching exp-sql) events)]
+                (is (not (seq evs)))
+                (when (seq evs)
+                  (println-err "Unexpected log:" events))
+                (let [{:strs [sql origin] :as info} (parse-event-info ev)]
+                  (is sql)
+                  (is (= exp-origin origin)))))))))))
 
 (deftest no-queries-are-logged-when-log-queires-is-false
   (tk-log/with-logged-event-maps logs
@@ -92,11 +136,11 @@
 
                (testing "AST/SQL is logged for both queries above"
                  ;; match the AST/SQL logs for nodes query
-                 (is (logs-include? logs "\"from\" \"nodes\""))
+                 (is (logs-include? logs "\"from\",\"nodes\""))
                  (is (logs-include? logs "latest_report_noop_pending"))
 
                  ;; match the AST/SQL logs for facts query
-                 (is (logs-include? logs "\"from\" \"facts\""))
+                 (is (logs-include? logs "\"from\",\"facts\""))
                  (is (logs-include? logs "(jsonb_each((stable||volatile)))")))))))))))
 
 (deftest no-PuppetDBServer-tk-service-queries-are-logged-when-log-queries-is-false
@@ -113,3 +157,36 @@
              (svcs/query pdb-service :v4 ["from" "facts"] nil identity)
              (svcs/query pdb-service :v4 ["from" "nodes"] nil identity)
              (is (empty? (prep-logs logs))))))))))
+
+(deftest queries-have-expected-log-mdc
+  ;; For now, we assume that all log messages generated by the
+  ;; query-eng ns during this period will be from threads handling the
+  ;; query.
+  (tk-log/with-log-level "puppetlabs.puppetdb.query-eng" :debug
+    (with-log-suppressed-unless-notable notable-pdb-event?
+      (with-test-db
+        (let [context {:scf-read-db *db*
+                       :url-prefix "/pdb"
+                       :log-queries true
+                       :add-agent-report-filter true
+                       :node-purge-ttl (time/parse-period "14d")}]
+          (replace-catalog catalog-1)
+          (let [events (atom [])]
+            (with-logging-to-atom "puppetlabs.puppetdb.query-eng" events
+              (qeng/stream-query-result :v4
+                                        ["from" "catalogs"]
+                                        {:origin "foo"}
+                                        context)
+              (doseq [event @events
+                      :let [mdc (.getMDCPropertyMap event)]]
+                (is (= "foo" (get mdc "pdb-query-origin")))
+                (is (uuid? (UUID/fromString (get mdc "pdb-query-id")))))))
+          (let [events (atom [])]
+            (with-logging-to-atom "puppetlabs.puppetdb.query-eng" events
+              (qeng/produce-streaming-body :v4
+                                           {:query ["from" "nodes"] :origin "foo"}
+                                           context)
+              (doseq [event @events
+                      :let [mdc (.getMDCPropertyMap event)]]
+                (is (= "foo" (get mdc "pdb-query-origin")))
+                (is (uuid? (UUID/fromString (get mdc "pdb-query-id"))))))))))))

test/puppetlabs/puppetdb/http/query_logging_test.clj

@@ -2,7 +2,9 @@
   (:require [puppetlabs.puppetdb.utils :refer :all]
             [clojure.test :refer :all]
             [puppetlabs.puppetdb.testutils :as tu]
-            [clojure.walk :as walk]))
+            [clojure.walk :as walk])
+  (:import
+   (org.apache.log4j MDC)))
 
 (deftest test-println-err
   (is (= "foo\n"
@@ -17,6 +19,23 @@
   (is (= "foo bar"
          (tu/with-err-str (print-err "foo" "bar")))))
 
+(deftest with-log-mdc-behavior
+  (with-log-mdc ["foo" "bar"]
+    (is (= "bar" (MDC/get "foo"))))
+  (with-log-mdc ["foo" nil]
+    (is (= {} (MDC/getContext))))
+  (with-log-mdc ["foo" nil "bar" 1]
+    (is (= {"bar" "1"} (MDC/getContext))))
+  (with-log-mdc ["foo" 1 "bar" nil]
+    (is (= {"foo" "1"} (MDC/getContext))))
+  (with-log-mdc ["foo" "x" "bar" "y"]
+    (is (= {"foo" "x" "bar" "y"} (MDC/getContext))))
+  (with-log-mdc ["foo" "x" "bar" nil "baz" "z"]
+    (is (= {"foo" "x" "baz" "z"} (MDC/getContext)))
+    (with-log-mdc ["foo" "x" "bar" "y" "baz" "z"]
+      (is (= {"foo" "x" "bar" "y" "baz" "z"} (MDC/getContext))))
+    (is (= {"foo" "x" "baz" "z"} (MDC/getContext)))))
+
 (deftest test-assoc-when
   (is (= {:a 1 :b 2}
          (assoc-when {:a 1 :b 2} :b 100)))