From 2d8f6549808e8b855d784ea80380b10f77b70b19 Mon Sep 17 00:00:00 2001
From: Jochen Topf <jochen@topf.org>
Date: Mon, 27 Feb 2023 10:50:17 +0100
Subject: [PATCH] Add functionality to send data to PostgreSQL in binary

Prepared statements can now send parameters in binary as well as in the
default text format. To do this, wrap the binary data in an instance of
the binary_param class, which is just a wrapper of std::string_view used
to signal to the exec_prepared() function that this parameter should be
sent in binary.

You can mix and match any number of parameters in text or binary format.

It isn't possible any more to use exec_prepared() without any
parameters, but that wasn't use (and doesn't make much sense anyway), so
no loss there.
---
 src/pgsql.hpp        | 35 +++++++++++++++++++++++++++++++----
 tests/test-pgsql.cpp | 41 +++++++++++++++++++++++++++++------------
 2 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/src/pgsql.hpp b/src/pgsql.hpp
index 782d79ae8..3736a115b 100644
--- a/src/pgsql.hpp
+++ b/src/pgsql.hpp
@@ -126,6 +126,20 @@ class pg_result_t
     std::unique_ptr<PGresult, pg_result_deleter_t> m_result;
 };
 
+/**
+ * Wrapper class for query parameters that should be sent to the database
+ * as binary parameter.
+ */
+class binary_param : public std::string_view
+{
+public:
+    using std::string_view::string_view;
+
+    binary_param(std::string const &str)
+    : std::string_view(str.data(), str.size())
+    {}
+};
+
 /**
  * PostgreSQL connection.
  *
@@ -229,6 +243,8 @@ class pg_conn_t
             return 0;
         } else if constexpr (std::is_same_v<T, std::string>) {
             return 0;
+        } else if constexpr (std::is_same_v<T, binary_param>) {
+            return 0;
         }
         return 1;
     }
@@ -240,12 +256,18 @@ class pg_conn_t
      * strings.
      */
     template <typename T>
-    static char const *to_str(std::vector<std::string> *data, T const &param)
+    static char const *to_str(std::vector<std::string> *data, int *length,
+                              int *bin, T const &param)
     {
         if constexpr (std::is_same_v<T, char const *>) {
             return param;
         } else if constexpr (std::is_same_v<T, std::string>) {
+            *length = param.size();
             return param.c_str();
+        } else if constexpr (std::is_same_v<T, binary_param>) {
+            *length = param.size();
+            *bin = 1;
+            return param.data();
         }
         return data->emplace_back(fmt::to_string(param)).c_str();
     }
@@ -275,16 +297,21 @@ class pg_conn_t
         std::vector<std::string> exec_params;
         exec_params.reserve(total_buffers_needed);
 
+        std::array<int, sizeof...(params)> lengths = {0};
+        std::array<int, sizeof...(params)> bins = {0};
+
         // This array holds the pointers to all parameter strings, either
         // to the original string parameters or to the recently converted
         // in the exec_params vector.
+        std::size_t n = 0;
+        std::size_t m = 0;
         std::array<char const *, sizeof...(params)> param_ptrs = {
-            to_str<std::decay_t<TArgs>>(&exec_params,
+            to_str<std::decay_t<TArgs>>(&exec_params, &lengths[n++], &bins[m++],
                                         std::forward<TArgs>(params))...};
 
         return exec_prepared_internal(stmt, sizeof...(params),
-                                      param_ptrs.data(), nullptr, nullptr,
-                                      result_as_binary ? 1 : 0);
+                                      param_ptrs.data(), lengths.data(),
+                                      bins.data(), result_as_binary ? 1 : 0);
     }
 
     struct pg_conn_deleter_t
diff --git a/tests/test-pgsql.cpp b/tests/test-pgsql.cpp
index f394fc224..60783c6dd 100644
--- a/tests/test-pgsql.cpp
+++ b/tests/test-pgsql.cpp
@@ -56,18 +56,6 @@ TEST_CASE("exec with invalid SQL should fail")
     REQUIRE_THROWS(conn.exec("XYZ"));
 }
 
-TEST_CASE("exec_prepared without parameters should work")
-{
-    auto const conn = db.db().connect();
-    conn.exec("PREPARE test AS SELECT 42");
-
-    auto const result = conn.exec_prepared("test");
-    REQUIRE(result.status() == PGRES_TUPLES_OK);
-    REQUIRE(result.num_fields() == 1);
-    REQUIRE(result.num_tuples() == 1);
-    REQUIRE(result.get(0, 0) == "42");
-}
-
 TEST_CASE("exec_prepared with single string parameters should work")
 {
     auto const conn = db.db().connect();
@@ -108,6 +96,35 @@ TEST_CASE("exec_prepared with non-string parameters should work")
     REQUIRE(result.get(0, 0) == "6");
 }
 
+TEST_CASE("exec_prepared with binary parameter should work")
+{
+    auto const conn = db.db().connect();
+    conn.exec("PREPARE test(bytea) AS SELECT length($1)");
+
+    binary_param const p{"foo \x01 bar"};
+    auto const result = conn.exec_prepared("test", p);
+    REQUIRE(result.status() == PGRES_TUPLES_OK);
+    REQUIRE(result.num_fields() == 1);
+    REQUIRE(result.num_tuples() == 1);
+    REQUIRE(result.get(0, 0) == "9");
+}
+
+TEST_CASE("exec_prepared with mixed parameter types should work")
+{
+    auto const conn = db.db().connect();
+    conn.exec("PREPARE test(text, bytea, int) AS"
+              " SELECT length($1) + length($2) + $3");
+
+    std::string const p1{"foo bar"};
+    binary_param const p2{"foo \x01 bar"};
+    int const p3 = 17;
+    auto const result = conn.exec_prepared("test", p1, p2, p3);
+    REQUIRE(result.status() == PGRES_TUPLES_OK);
+    REQUIRE(result.num_fields() == 1);
+    REQUIRE(result.num_tuples() == 1);
+    REQUIRE(result.get(0, 0) == "33"); // 7 + 9 + 17
+}
+
 TEST_CASE("create table and insert something")
 {
     auto const conn = db.db().connect();