Skip to content

Commit 0a36253

Browse files
rbrwrbrw
committed
(PDB-5216) Add pdb-query-id/pdb-query-origin to query threads' MDC
Provide the query origin and uuid in the log MDC (Mapped Diagnostic Context) for all threads handling the query.
1 parent 6d9b46c commit 0a36253

File tree

2 files changed

+125
-76
lines changed

2 files changed

+125
-76
lines changed

‎src/puppetlabs/puppetdb/query_eng.clj‎

Lines changed: 86 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
[puppetlabs.puppetdb.query-eng.default-reports :as dr]
2323
[puppetlabs.puppetdb.scf.storage-utils :as sutils]
2424
[puppetlabs.puppetdb.schema :as pls]
25-
[puppetlabs.puppetdb.utils :as utils]
25+
[puppetlabs.puppetdb.utils :as utils :refer [with-log-mdc]]
2626
[puppetlabs.puppetdb.utils.string-formatter :as formatter]
2727
[puppetlabs.puppetdb.query-eng.engine :as eng]
2828
[ring.util.io :as rio]
@@ -212,27 +212,30 @@
212212
;; For now, generate the ids here; perhaps later, higher up
213213
(assert (not (:query-id context)))
214214
(let [query-id (str (java.util.UUID/randomUUID))
215-
context (assoc context :query-id query-id)]
216-
(let [{:keys [scf-read-db url-prefix warn-experimental pretty-print log-queries]
217-
:or {warn-experimental true}} context
218-
{:keys [remaining-query entity]} (eng/parse-query-context version query warn-experimental)
219-
munge-fn (get-munge-fn entity version options url-prefix)]
215+
context (assoc context :query-id query-id)
216+
origin (:origin options)]
217+
(with-log-mdc ["pdb-query-id" query-id
218+
"pdb-query-origin" origin]
219+
(let [{:keys [scf-read-db url-prefix warn-experimental pretty-print log-queries]
220+
:or {warn-experimental true}} context
221+
{:keys [remaining-query entity]} (eng/parse-query-context version query warn-experimental)
222+
munge-fn (get-munge-fn entity version options url-prefix)]
220223

221-
(when log-queries
222-
;; Log origin and AST of incoming query
223-
(log/infof "PDBQuery:%s:%s"
224-
query-id (-> (sorted-map :origin (:origin options) :ast query)
225-
json/generate-string)))
224+
(when log-queries
225+
;; Log origin and AST of incoming query
226+
(log/infof "PDBQuery:%s:%s"
227+
query-id (-> (sorted-map :origin origin :ast query)
228+
json/generate-string)))
226229

227-
(let [f #(let [{:keys [results-query]}
228-
(query->sql remaining-query entity version
229-
options
230-
(select-keys context [:log-queries :query-id]))]
231-
(jdbc/call-with-array-converted-query-rows results-query
232-
(comp row-fn munge-fn)))]
233-
(if use-preferred-streaming-method?
234-
(jdbc/with-db-connection scf-read-db (jdbc/with-db-transaction [] (f)))
235-
(jdbc/with-transacted-connection scf-read-db (f))))))))
230+
(let [f #(let [{:keys [results-query]}
231+
(query->sql remaining-query entity version
232+
options
233+
(select-keys context [:log-queries :query-id]))]
234+
(jdbc/call-with-array-converted-query-rows results-query
235+
(comp row-fn munge-fn)))]
236+
(if use-preferred-streaming-method?
237+
(jdbc/with-db-connection scf-read-db (jdbc/with-db-transaction [] (f)))
238+
(jdbc/with-transacted-connection scf-read-db (f)))))))))
236239

237240
;; Do we still need this, i.e. do we need the pass-through, and the
238241
;; strict selectivity in the caller below?
@@ -325,43 +328,45 @@
325328
stream (generated-stream
326329
;; Runs in a future
327330
(fn [out]
328-
(with-open! [out (io/writer out :encoding "UTF-8")]
329-
(try
330-
(jdbc/with-db-connection db
331-
(jdbc/with-db-transaction []
332-
(let [{:keys [results-query count-query]}
333-
(query->sql query entity version query-options context)
334-
st (when count-query
335-
{:count (jdbc/get-result-count count-query)})
336-
stream-row (fn [row]
337-
(let [r (munge-fn row)]
338-
(when-not (instance? PGobject r)
339-
(first r))
340-
(when-not (realized? status)
341-
(deliver status st))
342-
(try
343-
(http/stream-json r out pretty-print)
344-
(catch IOException ex
345-
(log/debug ex (trs "Unable to stream response: {0}"
346-
(.getMessage ex)))
347-
(throw quiet-exit)))))]
348-
(jdbc/call-with-array-converted-query-rows results-query
349-
stream-row)
350-
(when-not (realized? status)
351-
(deliver status st)))))
352-
(catch Exception ex
353-
;; If it's an exit, we've already handled it.
354-
(when-not (identical? quiet-exit ex)
331+
(with-log-mdc ["pdb-query-id" query-id
332+
"pdb-query-origin" (:origin query-options)]
333+
(with-open! [out (io/writer out :encoding "UTF-8")]
334+
(try
335+
(jdbc/with-db-connection db
336+
(jdbc/with-db-transaction []
337+
(let [{:keys [results-query count-query]}
338+
(query->sql query entity version query-options context)
339+
st (when count-query
340+
{:count (jdbc/get-result-count count-query)})
341+
stream-row (fn [row]
342+
(let [r (munge-fn row)]
343+
(when-not (instance? PGobject r)
344+
(first r))
345+
(when-not (realized? status)
346+
(deliver status st))
347+
(try
348+
(http/stream-json r out pretty-print)
349+
(catch IOException ex
350+
(log/debug ex (trs "Unable to stream response: {0}"
351+
(.getMessage ex)))
352+
(throw quiet-exit)))))]
353+
(jdbc/call-with-array-converted-query-rows results-query
354+
stream-row)
355+
(when-not (realized? status)
356+
(deliver status st)))))
357+
(catch Exception ex
358+
;; If it's an exit, we've already handled it.
359+
(when-not (identical? quiet-exit ex)
360+
(if (realized? status)
361+
(throw ex)
362+
(deliver status {:error ex}))))
363+
(catch Throwable ex
355364
(if (realized? status)
356-
(throw ex)
357-
(deliver status {:error ex}))))
358-
(catch Throwable ex
359-
(if (realized? status)
360-
(do
361-
(log/error ex (trs "Query streaming failed: {0} {1}"
362-
query query-options))
363-
(throw ex))
364-
(deliver status {:error ex})))))))]
365+
(do
366+
(log/error ex (trs "Query streaming failed: {0} {1}"
367+
query query-options))
368+
(throw ex))
369+
(deliver status {:error ex}))))))))]
365370
{:status status
366371
:stream stream}))
367372

@@ -417,11 +422,12 @@
417422
:or {warn-experimental true}} context
418423
query-config (select-keys context [:node-purge-ttl :add-agent-report-filter])
419424
{:keys [query remaining-query entity query-options]}
420-
(user-query->engine-query version query-map query-config warn-experimental)]
425+
(user-query->engine-query version query-map query-config warn-experimental)
426+
origin (:origin query-map)]
421427

422428
(when log-queries
423429
;; Log origin and AST of incoming query
424-
(let [{:keys [origin query]} query-map]
430+
(let [query (:query query-map)]
425431
(log/infof "PDBQuery:%s:%s"
426432
query-id (-> (sorted-map :origin origin :ast query)
427433
json/generate-string))))
@@ -437,19 +443,21 @@
437443
query-error (promise)
438444
resp (http/streamed-response
439445
buffer
440-
(try (jdbc/with-transacted-connection scf-read-db
441-
(jdbc/call-with-array-converted-query-rows
442-
results-query
443-
(comp #(http/stream-json % buffer pretty-print)
444-
#(do (when-not (instance? PGobject %)
445-
(first %))
446-
(deliver query-error nil) %)
447-
(munge-fn-hook munge-fn))))
448-
;; catch throwable to make sure any trouble is forwarded to the
449-
;; query-error promise below. If something throws and is not passed
450-
;; along the deref will block indefinitely.
451-
(catch Throwable e
452-
(deliver query-error e))))]
446+
(with-log-mdc ["pdb-query-id" query-id
447+
"pdb-query-origin" origin]
448+
(try (jdbc/with-transacted-connection scf-read-db
449+
(jdbc/call-with-array-converted-query-rows
450+
results-query
451+
(comp #(http/stream-json % buffer pretty-print)
452+
#(do (when-not (instance? PGobject %)
453+
(first %))
454+
(deliver query-error nil) %)
455+
(munge-fn-hook munge-fn))))
456+
;; catch throwable to make sure any trouble is forwarded to the
457+
;; query-error promise below. If something throws and is not passed
458+
;; along the deref will block indefinitely.
459+
(catch Throwable e
460+
(deliver query-error e)))))]
453461
(if @query-error
454462
(throw @query-error)
455463
(cond-> (http/json-response* resp)
@@ -483,10 +491,13 @@
483491
context :- query-context-schema]
484492
;; For now, generate the ids here; perhaps later, higher up
485493
(assert (not (:query-id context)))
486-
(let [context (assoc context :query-id (str (java.util.UUID/randomUUID)))]
487-
(if use-preferred-streaming-method?
488-
(preferred-produce-streaming-body version query-map context)
489-
(deprecated-produce-streaming-body version query-map context))))
494+
(let [qid (str (java.util.UUID/randomUUID))
495+
context (assoc context :query-id qid)]
496+
(with-log-mdc ["pdb-query-id" qid
497+
"pdb-query-origin" (:origin query-map)]
498+
(if use-preferred-streaming-method?
499+
(preferred-produce-streaming-body version query-map context)
500+
(deprecated-produce-streaming-body version query-map context)))))
490501

491502
(pls/defn-validated object-exists? :- s/Bool
492503
"Returns true if an object exists."

‎test/puppetlabs/puppetdb/http/query_logging_test.clj‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,21 @@
77
[puppetlabs.kitchensink.core :as kitchensink]
88
[puppetlabs.puppetdb.cheshire :as json]
99
[puppetlabs.puppetdb.cli.services :as svcs]
10+
[puppetlabs.puppetdb.query-eng :as qeng]
1011
[puppetlabs.puppetdb.testutils.catalogs :refer [replace-catalog]]
1112
[puppetlabs.puppetdb.testutils.db :refer [with-test-db *db*]]
1213
[puppetlabs.puppetdb.testutils.http
1314
:refer [call-with-http-app query-response with-http-app*]]
15+
[puppetlabs.puppetdb.testutils.log :refer [notable-pdb-event?]]
1416
[puppetlabs.puppetdb.testutils.services
1517
:refer [call-with-puppetdb-instance create-temp-config *server*]]
18+
[puppetlabs.puppetdb.time :as time]
1619
[puppetlabs.puppetdb.utils :refer [println-err]]
1720
[puppetlabs.trapperkeeper.app :refer [get-service]]
1821
[puppetlabs.trapperkeeper.testutils.logging :as tk-log
19-
:refer [with-logged-event-maps]])
22+
:refer [with-log-suppressed-unless-notable
23+
with-logged-event-maps
24+
with-logging-to-atom]])
2025
(:import
2126
(java.util UUID)))
2227

@@ -152,3 +157,36 @@
152157
(svcs/query pdb-service :v4 ["from" "facts"] nil identity)
153158
(svcs/query pdb-service :v4 ["from" "nodes"] nil identity)
154159
(is (empty? (prep-logs logs))))))))))
160+
161+
(deftest queries-have-expected-log-mdc
162+
;; For now, we assume that all log messages generated by the
163+
;; query-eng ns during this period will be from threads handling the
164+
;; query.
165+
(tk-log/with-log-level "puppetlabs.puppetdb.query-eng" :debug
166+
(with-log-suppressed-unless-notable notable-pdb-event?
167+
(with-test-db
168+
(let [context {:scf-read-db *db*
169+
:url-prefix "/pdb"
170+
:log-queries true
171+
:add-agent-report-filter true
172+
:node-purge-ttl (time/parse-period "14d")}]
173+
(replace-catalog catalog-1)
174+
(let [events (atom [])]
175+
(with-logging-to-atom "puppetlabs.puppetdb.query-eng" events
176+
(qeng/stream-query-result :v4
177+
["from" "catalogs"]
178+
{:origin "foo"}
179+
context)
180+
(doseq [event @events
181+
:let [mdc (.getMDCPropertyMap event)]]
182+
(is (= "foo" (get mdc "pdb-query-origin")))
183+
(is (uuid? (UUID/fromString (get mdc "pdb-query-id")))))))
184+
(let [events (atom [])]
185+
(with-logging-to-atom "puppetlabs.puppetdb.query-eng" events
186+
(qeng/produce-streaming-body :v4
187+
{:query ["from" "nodes"] :origin "foo"}
188+
context)
189+
(doseq [event @events
190+
:let [mdc (.getMDCPropertyMap event)]]
191+
(is (= "foo" (get mdc "pdb-query-origin")))
192+
(is (uuid? (UUID/fromString (get mdc "pdb-query-id"))))))))))))

0 commit comments

Comments
 (0)