(PDB-5216) Support a query "origin"

Add support for an optional query "origin", which can be provided via
a query parameters, and log it along with the query id.

Change the formatting of the query id log entries to be a json map,
e.g.

  2021-08-16 11:14:37,723 INFO  [qtp237850705-35] [p.p.query-eng]
  PDBQuery:6dd2a796-a4a3-44be-afb6-2a84459367f6:
  {"ast":["from","catalogs",["=","certname","host-3"]],"origin":"puppet:somewhere"}

to avoid any ambiguity, or the need for restrictions or a bespoke
encoding, with respect to the origin string contents.

Author: rbrw

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

@@ -102,27 +102,32 @@
          :when (:queryable? projection-value)]
      (keyword projection-key))))
 
-(defn maybe-log-sql [logging-options gen-sql]
-  (if-let [log-id (:log-id logging-options)]
-    (let [sql-output (gen-sql)]
-      (log/infof "%s:%s:%s" "PDBQuery" log-id (:results-query sql-output))
-      sql-output)
+(defn maybe-log-sql [options context gen-sql]
+  (if-let [log-id (:log-id context)]
+    (let [{[query & params] :results-query :as sql} (gen-sql)]
+      (log/infof "PDBQuery:%s:%s"
+                 log-id (-> (sorted-map :origin (:origin options)
+                                        :sql query
+                                        :params (vec params))
+                            json/generate-string))
+      sql)
     (gen-sql)))
 
 (defn query->sql
   "Converts a vector-structured `query` to a corresponding SQL query which will
    return nodes matching the `query`."
   ([query entity version options]
    (query->sql query entity version options {}))
-  ([query entity version options logging-options]
+  ([query entity version options context]
    {:pre  [((some-fn nil? sequential?) query)]
     :post [(map? %)
            (jdbc/valid-jdbc-query? (:results-query %))
            (or (not (:include_total options))
                (jdbc/valid-jdbc-query? (:count-query %)))]}
 
    (maybe-log-sql
-    logging-options
+    options
+    context
     (fn []
       (cond
         (= :aggregate-event-counts entity)
@@ -214,8 +219,10 @@
          munge-fn (get-munge-fn entity version options url-prefix)]
 
      (when log-queries
-       ;; log the AST of the incoming query
-       (log/infof "%s:%s:%s" "PDBQuery" log-id query))
+       ;; Log origin and AST of incoming query
+       (log/infof "PDBQuery:%s:%s"
+                  log-id (-> (sorted-map :origin (:origin options) :ast query)
+                             json/generate-string)))
 
      (let [f #(let [{:keys [results-query]}
                     (query->sql remaining-query entity version
@@ -367,8 +374,11 @@
         (user-query->engine-query version query-map query-config warn-experimental)]
 
     (when log-queries
-      ;; log the AST of the incoming query
-      (log/infof "%s:%s:%s" "PDBQuery" log-id (:query query-map)))
+      ;; Log origin and AST of incoming query
+      (let [{:keys [origin query]} query-map]
+        (log/infof "PDBQuery:%s:%s"
+                   log-id (-> (sorted-map :origin origin :ast query)
+                              json/generate-string))))
 
     (try
       (let [munge-fn (get-munge-fn entity version query-options url-prefix)
@@ -407,8 +417,11 @@
         (user-query->engine-query version query-map query-config warn-experimental)]
 
     (when log-queries
-      ;; log the AST of the incoming query
-      (log/infof "%s:%s:%s" "PDBQuery" log-id (:query query-map)))
+      ;; Log origin and AST of incoming query
+      (let [{:keys [origin query]} query-map]
+        (log/infof "PDBQuery:%s:%s"
+                   log-id (-> (sorted-map :origin origin :ast query)
+                              json/generate-string))))
 
     (try
       (jdbc/with-transacted-connection scf-read-db

test/puppetlabs/puppetdb/http/query_logging_test.clj

@@ -54,39 +54,51 @@
     (with-test-db
       (replace-catalog catalog-1)
       (with-http-app* #(assoc % :log-queries true)
-        (doseq [[query ast-exp sql-exp]
+        (doseq [[query exp-ast exp-sql exp-origin]
                 ;; produce-streaming-body
-                [[["/v4" ["from" "nodes"]]
-                  "\"from\" \"nodes\""
-                  "latest_report_noop_pending"]
+                [[["/v4" ["from" "nodes"] {:origin "foo"}]
+                  "\"from\",\"nodes\""
+                  "latest_report_noop_pending"
+                  "foo"]
                  [["/v4" ["from" "facts"]]
-                  "\"from\" \"facts\""
-                  "(jsonb_each((stable||volatile)))"]
+                  "\"from\",\"facts\""
+                  "(jsonb_each((stable||volatile)))"
+                  nil]
                  ;; stream-query-result
-                 [["/v4/catalogs/myhost.localdomain" []]
-                  "\"from\" \"catalogs\""
-                  "row_to_json(edge_data)"]]]
+                 [["/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
-                  uuid (some->> events
-                                (some #(->> % :message (re-find #"^PDBQuery:([^:]+):")))
-                                second)
+                  ;; 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))))]
 
               (is (uuid? (UUID/fromString uuid)))
 
-              (let [asts (filter (qev-matching ast-exp) events)]
-                (when-not (is (= 1 (count asts)))
-                  (println-err "Unexpected log:" events)))
-
-              (let [sqls (filter (qev-matching sql-exp) events)]
-                (when-not (is (= 1 (count sqls)))
-                  (println-err "Unexpected log:" events))))))))))
+              (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))))
+
+              (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
@@ -119,11 +131,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