Set bytea PreparedStatement parameters via .setBytes and .setArray

In 42.7.7 (pgjdbc/pgjdbc#3757) pgjdbc broke
parameters set via

  (.setObject stmt i (doto (PGobject.)
                       (.setType "bytea")
                       (.setValue "\x42abcdef)))

which is the way pdb has handled them, whether explicitly, or
indirectly via clojure.java.jdbc and next.jdbc, which both default
to (.setObject ...) for parameter types without a more specific
protocol extension.

While this is intended to work (see the issue thread), the pgjdbc
documentation specifies `.setBytes` for bytea parameters, and
`.setArray` for arrays of bytea:

  https://jdbc.postgresql.org/documentation/binary-data/
  https://jdbc.postgresql.org/documentation/server-prepare/#arrays

Add new PDBBytea and VecPDBBytea classes for the bytea and bytea[]
parameters and extend the relevant clojure.java.jdbc and next.jdbc
protocols to call the appropriate PreparedStatement method for each,
rather than the `.setObject` default. Note that we cannot use
defrecords because (by default) honeysql would treat them
as unrecognized extensions.

[[email protected]: remove unused import, fix typo, add
hashCode method to PDBBytea]

Reviewed-by: Austin Blatt <[email protected]>

Author: rlbdv

project.clj

@@ -78,17 +78,23 @@
          "puppetserver"
          "vendor"]))
 
+(def pdb-aot-classes
+  ;; Compile classes first for now:
+  ;; https://codeberg.org/leiningen/leiningen/issues/99
+  '[puppetlabs.puppetdb.jdbc.PDBBytea puppetlabs.puppetdb.jdbc.VecPDBBytea])
+
 (def pdb-aot-namespaces
-  (apply vector
-         #"puppetlabs\.puppetdb\..*"
-         (->> "resources/puppetlabs/puppetdb/bootstrap.cfg"
-              clojure.java.io/reader
-              line-seq
-              (map clojure.string/trim)
-              (remove #(re-matches #"#.*" %))  ;; # comments
-              (remove #(re-matches #"puppetlabs\.puppetdb\.." %))
-              (map #(clojure.string/replace % #"(.*)/[^/]+$" "$1"))
-              (map symbol))))
+  (into []
+        (concat pdb-aot-classes
+                [#"puppetlabs\.puppetdb\..*"]
+                (->> "resources/puppetlabs/puppetdb/bootstrap.cfg"
+                     clojure.java.io/reader
+                     line-seq
+                     (map clojure.string/trim)
+                     (remove #(re-matches #"#.*" %)) ;; # comments
+                     (remove #(re-matches #"puppetlabs\.puppetdb\.." %))
+                     (map #(clojure.string/replace % #"(.*)/[^/]+$" "$1"))
+                     (map symbol)))))
 
 ;; Avoid startup reflection warnings due to
 ;; https://clojure.atlassian.net/browse/CLJ-2066
@@ -216,6 +222,8 @@
   ;; "test"].  See the :testutils profile below.
   :classifiers  {:test :testutils}
 
+  :aot ~pdb-aot-classes
+
   :profiles {:defaults {:resource-paths ["test-resources"]
                         :dependencies ~pdb-dev-deps
                         :injections [(do

src/puppetlabs/puppetdb/jdbc/PDBBytea.clj

@@ -0,0 +1,31 @@
+(ns puppetlabs.puppetdb.jdbc.PDBBytea
+  "Carrier for bytea parameters, to support clojure.java.jdbc and next.jdbc
+  protocol extensions. Essentially just a typed wrapper around a byte[]."
+  (:import
+   (java.util Arrays))
+  (:gen-class
+   ;;:extends org.postgresql.util.PGobject
+   :state ^"[B" data
+   :init init
+   :constructors {["[B"] []}))
+
+(def ^:private warn-on-reflection-orig *warn-on-reflection*)
+(set! *warn-on-reflection* true)
+
+(defn -init [data] [[] data])
+
+;; Implemented for now because edge-replacement-differential test, for
+;; example, uses munge-hash-for-storage to create the test values that
+;; it compares for equality -- and who knows what else. Apparently
+;; PGobject bytea values, which PDBBytea replaces, compare content for
+;; equality, and since we use bytea for hashes, best to be
+;; conservative.
+(defn -equals [^puppetlabs.puppetdb.jdbc.PDBBytea this x]
+  (and (instance? puppetlabs.puppetdb.jdbc.PDBBytea x)
+       (Arrays/equals ^"[B" (.data this)
+                      ^"[B" (.data ^puppetlabs.puppetdb.jdbc.PDBBytea x))))
+
+(defn -hashCode [^puppetlabs.puppetdb.jdbc.PDBBytea this]
+  (Arrays/hashCode ^"[B" (.data this)))
+
+(set! *warn-on-reflection* warn-on-reflection-orig)

src/puppetlabs/puppetdb/jdbc/VecPDBBytea.clj

@@ -0,0 +1,16 @@
+(ns puppetlabs.puppetdb.jdbc.VecPDBBytea
+  "Carrier for bytea[] parameters, to support clojure.java.jdbc and next.jdbc
+  protocol extensions. Essentially just a typed wrapper around a byte[][]."
+  ;; It doesn't look like java.sql.Arrays, which is what this
+  ;; replaces, define content-based equals(), so we don't need one.
+  (:gen-class
+   :state ^"[[B" data
+   :init init
+   :constructors {["[[B"] []}))
+
+(def ^:private warn-on-reflection-orig *warn-on-reflection*)
+(set! *warn-on-reflection* true)
+
+(defn -init [data] [[] data])
+
+(set! *warn-on-reflection* warn-on-reflection-orig)

src/puppetlabs/puppetdb/scf/migrate.clj

@@ -71,7 +71,7 @@
    (java.time LocalDateTime LocalDate ZonedDateTime Instant)
    (java.time.temporal ChronoUnit)
    (java.time.format DateTimeFormatter)
-   (org.postgresql.util PGobject)))
+   (puppetlabs.puppetdb.jdbc PDBBytea)))
 
 (defn init-through-2-3-8
   []
@@ -1432,7 +1432,7 @@
             from (select unnest(?) as id, unnest(?) as hash) in_data
             where fact_values.id = in_data.id"
           [(sutils/array-to-param "bigint" Long ids)
-           (sutils/array-to-param "bytea" PGobject hashes)])))))
+           (sutils/array-to-param "bytea" PDBBytea hashes)])))))
 
   (log/info (trs "[7/8] Indexing fact_values table..."))
   (jdbc/do-commands

src/puppetlabs/puppetdb/scf/storage.clj

@@ -58,7 +58,7 @@
    (java.time ZoneId ZonedDateTime)
    (java.util Arrays)
    (org.joda.time Period)
-   (org.postgresql.util PGobject)))
+   (puppetlabs.puppetdb.jdbc PDBBytea)))
 
 ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
 ;;; Schemas
@@ -475,7 +475,7 @@
   (if (seq resource-hashes)
     (let [resource-array (->> (vec resource-hashes)
                               (map sutils/munge-hash-for-storage)
-                              (sutils/array-to-param "bytea" org.postgresql.util.PGobject))
+                              (sutils/array-to-param "bytea" PDBBytea))
           query (format "SELECT DISTINCT %s as resource FROM resource_params_cache WHERE resource=ANY(?)"
                         (sutils/sql-hash-as-str "resource"))
           sql-params [query resource-array]]
@@ -1039,7 +1039,7 @@
 (defn find-package-hashes [package-hashes]
   (jdbc/call-with-query-rows [(format "SELECT id, %s as package_hash FROM packages WHERE hash = ANY(?)"
                                       (sutils/sql-hash-as-str "hash"))
-                              (sutils/array-to-param "bytea" PGobject
+                              (sutils/array-to-param "bytea" PDBBytea
                                                      (map sutils/munge-hash-for-storage package-hashes))]
                              (fn [rows]
                                (reduce (fn [acc {:keys [id package_hash]}]

src/puppetlabs/puppetdb/scf/storage_utils.clj

@@ -4,26 +4,36 @@
    [clojure.java.jdbc :as sql]
    [clojure.string :as str]
    [clojure.tools.logging :as log]
+   [next.jdbc.prepare :refer [SettableParameter]]
    [puppetlabs.i18n.core :refer [trs tru]]
    [puppetlabs.kitchensink.core :as kitchensink]
    [puppetlabs.puppetdb.cheshire :as json]
    [puppetlabs.puppetdb.honeysql]
    [puppetlabs.puppetdb.jdbc :as jdbc]
    [puppetlabs.puppetdb.query.common :refer [bad-query-ex]]
    [puppetlabs.puppetdb.schema :as pls]
+   [puppetlabs.puppetdb.utils :refer [byte-array-class]]
    [schema.core :as s])
   (:import
-   [java.sql Connection]
-   [java.util UUID]
-   [org.postgresql.util PGobject]))
+   (java.sql Connection PreparedStatement)
+   (java.util HexFormat UUID)
+   (org.postgresql.util PGobject)
+   (puppetlabs.puppetdb.jdbc PDBBytea VecPDBBytea)))
+
+(def ^:private warn-on-reflection-orig *warn-on-reflection*)
+(set! *warn-on-reflection* true)
 
 ;; SCHEMA
 
 (defn array-to-param
   [col-type java-type values]
-  (.createArrayOf ^Connection (:connection (jdbc/db))
-                  col-type
-                  (into-array java-type values)))
+  ;; We create the VecPDBBytea so we can redirect the parameter to
+  ;; setArray in the set-parameter specializations.
+  (if (= java-type PDBBytea)
+    ;; FIXME: change byte-array-class to byte/1 once we require clojure 1.12+
+    (VecPDBBytea. (into-array byte-array-class (map #(.data ^PDBBytea %) values)))
+    (.createArrayOf ^Connection (:connection (jdbc/db))
+                    col-type (into-array java-type values))))
 
 (def pg-extension-map
   "Maps to the table definition in postgres, but only includes some of the
@@ -427,9 +437,8 @@
   (sql/with-db-connection [_conn db]
     (sql/execute! db ["vacuum analyze"] {:transaction? false})))
 
-(defn parse-db-hash
-  [^PGobject db-hash]
-  (str/replace (.getValue db-hash) "\\x" ""))
+(defn db-hash->hex [^PDBBytea pdbb]
+  (.formatHex (HexFormat/of) (.data pdbb)))
 
 (defn parse-db-uuid
   [^UUID db-uuid]
@@ -454,9 +463,43 @@
 (defn bytea-escape [s]
   (format "\\x%s" s))
 
-(defn munge-hash-for-storage
-  [hash]
-  (str->pgobject "bytea" (bytea-escape hash)))
+(extend-protocol sql/ISQLParameter
+  ;; This became necessary with more recent pgjdbc
+  ;; versions (42.7.7 (at least) and newer), because it stopped
+  ;; handling a "bytea" PGobject; it started crashing in
+  ;; PGBytea/toPGLiteral, which doesn't have a clause to handle the
+  ;; escaped String that it ends up with, presumably from the PGobject
+  ;; value. The pgjdbc docs specify .setBytes, so do
+  ;; that:
+  ;;   https://jdbc.postgresql.org/documentation/binary-data/
+  ;;   https://clojure-doc.org/articles/ecosystem/java_jdbc/using_sql/
+  ;;   https://clojure.github.io/java.jdbc/#clojure.java.jdbc/ISQLParameter
+  PDBBytea
+  (sql/set-parameter [v ^PreparedStatement s i]
+    (.setBytes s i (.data v)))
+  VecPDBBytea
+  (sql/set-parameter [v ^PreparedStatement s i]
+    (.setArray s i (.createArrayOf ^Connection (:connection (jdbc/db))
+                                   "bytea" (.data v)))))
+
+(extend-protocol SettableParameter
+  PDBBytea
+  (sql/set-parameter [v ^PreparedStatement s i] (.setBytes s i (.data v)))
+  VecPDBBytea
+  (sql/set-parameter [v ^PreparedStatement s i]
+    (.setArray s i (.createArrayOf ^Connection (:connection (jdbc/db))
+                                   "bytea" (.data v)))))
+
+(defn munge-hash-for-storage [hash]
+  ;; Represent hashes as PDBBytea so we can add them to
+  ;; PreparedStatments via setBytes via the clojure/next jdbc
+  ;; protocols above. It'd be more efficient if eventually pdb just
+  ;; used something like PDBBytea or byte[] everywhere, though the
+  ;; latter would allow mutation and might raise "ownership"
+  ;; questions. The *most* efficient hash that's not just a byte[]
+  ;; might be a class with two long members and a short to store the
+  ;; 20 bytes (i.e. so that the data would be inline in the class).
+  (PDBBytea. (.parseHex (HexFormat/of) hash)))
 
 (defn munge-json-for-storage
   "Prepare a clojure object for storage depending on db type."
@@ -497,3 +540,5 @@
   (log/info (trs "Analyzing small tables"))
   (apply jdbc/do-commands-outside-txn
          (map #(str "analyze " %) small-tables)))
+
+(set! *warn-on-reflection* warn-on-reflection-orig)

test/puppetlabs/puppetdb/scf/storage_test.clj

@@ -290,7 +290,7 @@
                           :stable {:hostname "myhost" :kernel "Linux" :uptime_seconds 3600}
                           :volatile {:domain "mynewdomain.com" :fqdn "myhost.mynewdomain.com"}}
                          (-> m
-                             (update :hash sutils/parse-db-hash)
+                             (update :hash sutils/db-hash->hex)
                              (update :paths_hash bytes->hex)
                              (update :stable_hash bytes->hex)
                              (update :stable sutils/parse-db-json)