Connection.hpp

Go to the documentation of this file.

#pragma once
 
 
#if __has_include(<libpq-fe.h>)
#  include <libpq-fe.h>
#else
#  include <postgresql/libpq-fe.h>
#endif
#include <stdexcept>
#include <string>
#include <functional>
#include <arpa/inet.h>
#include <cstring>
 
namespace PluginSQL {
 
class Connection {
  PGconn* conn = nullptr;
  std::vector<char> rowBuf;
 
public:
  Connection(const std::string& user,
             const std::string& pass,
             const std::string& host,
             int port,
             const std::string& dbname)
  {
    std::string conninfo =
      "user="     + user +
      " password="+ pass +
      " host="    + host +
      " port="    + std::to_string(port) +
      " dbname="  + dbname;
    conn = PQconnectdb(conninfo.c_str());
    if (PQstatus(conn) != CONNECTION_OK)
      throw std::runtime_error(PQerrorMessage(conn));
  }
 
  ~Connection() {
    if (conn) PQfinish(conn);
  }
 
  void Execute(const std::string &sql) {
    PGresult* res = PQexec(conn, sql.c_str());
    if (PQresultStatus(res) != PGRES_COMMAND_OK) {
      std::string err = PQerrorMessage(conn);
      PQclear(res);
      throw std::runtime_error(err);
    }
    PQclear(res);
  }
 
  // Lance COPY BINARY sur la table en mode binaire
  void CopyBinary(
    const std::string &table,
    const std::string &columns,
    const std::function<void()> &writeRows
  ) {
    std::string stmt =
      "COPY " + table + " " + columns + " FROM STDIN WITH (FORMAT binary)";
    PGresult* res = PQexec(conn, stmt.c_str());
    if (PQresultStatus(res) != PGRES_COPY_IN) {
      std::string err = PQerrorMessage(conn);
      PQclear(res);
      throw std::runtime_error("COPY start failed: " + err);
    }
    PQclear(res);
 
    writeRows();  // header, lignes, trailer
 
    if (PQputCopyEnd(conn, nullptr) != 1)
      throw std::runtime_error("COPY end failed");
 
    res = PQgetResult(conn);
    if (PQresultStatus(res) != PGRES_COMMAND_OK) {
      std::string err = PQerrorMessage(conn);
      PQclear(res);
      throw std::runtime_error("COPY finalization failed: " + err);
    }
    PQclear(res);
  }
 
  // Écrit l'en-tête global du protocole COPY BINARY
  inline void headerCopy(size_t N, size_t rowSize) {
    //static const char sig[] = "PGCOPY\n\377\r\n\0";
    prepareBuffer(N, rowSize);
 
    static const unsigned char sig[11] = { 'P','G','C','O','P','Y','\n', 0xFF, '\r', '\n', '\0' };
    appendRaw(sig, sizeof(sig));
 
    int32_t zero = htonl(0);
    appendRaw(&zero, sizeof(zero));
    appendRaw(&zero, sizeof(zero));
  }
 
  // Écrit le trailer (fin de flux COPY)
  inline void footerCopy() {
    int16_t tr = htons(-1);
    appendRaw(&tr, sizeof(tr)); 
    PQputCopyData(conn, rowBuf.data(), (int)rowBuf.size());
    rowBuf.clear();
  }
 
  void prepareBuffer(size_t N, size_t rowSize) {
    size_t headerSize  = 19; // 11 for signature + 4 for flags + 4 for zero
    size_t footerSize  = 2;   // 2 for trailer
    size_t totalSize = headerSize + N * rowSize + footerSize;
    rowBuf.reserve(totalSize);
  }
 
  void beginRow(int16_t nCols) {
    int16_t be = htons(nCols);
    appendRaw(&be, sizeof(be));
  }
 
  // b) append d’un bloc brut
  void appendRaw(const void* data, size_t len) {
    auto p = reinterpret_cast<const char*>(data);
    rowBuf.insert(rowBuf.end(), p, p + len);
  }
 
  // c) append d’un champ (length + data)
  void appendField(const void* data, size_t len) {
    int32_t len_be = htonl((int32_t)len);
    appendRaw(&len_be, sizeof(len_be));
    appendRaw(data, len);
  }
 
inline void appendBool(bool b) {
  int32_t len_be = htonl(1);
  appendRaw(&len_be, sizeof(len_be));
  unsigned char v = b ? 1 : 0;
  appendRaw(&v, sizeof(v));
}
  inline void appendInt(int v) {
    int32_t v_be = htonl(static_cast<int32_t>(v));
    appendField(&v_be, sizeof(v_be));
  }
 
  inline void appendDouble(double d) {
    uint64_t u;
    std::memcpy(&u, &d, sizeof(u));
    u = __builtin_bswap64(u);
    appendField(&u, sizeof(u));
  }
};
 
} // namespace PluginSQL