Merge pull request #1941 from joto/pgsql-binary

lonvia · web-flow · commit ae807b0c8c47 · 2023-03-02T09:52:17.000+01:00
Add functionality to send data to PostgreSQL in binary
diff --git 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
@@ -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();
