summaryrefslogtreecommitdiff
path: root/lib
diff options
context:
space:
mode:
Diffstat (limited to 'lib')
-rw-r--r--lib/output/pq/pqConn.cpp36
-rw-r--r--lib/output/pq/pqConn.h4
2 files changed, 40 insertions, 0 deletions
diff --git a/lib/output/pq/pqConn.cpp b/lib/output/pq/pqConn.cpp
index 76a3a75..53fd3db 100644
--- a/lib/output/pq/pqConn.cpp
+++ b/lib/output/pq/pqConn.cpp
@@ -1,6 +1,7 @@
#include "pqConn.h"
#include "pqBindings.h"
#include "pqStmt.h"
+#include <compileTimeFormatter.h>
#include <dbConn.h>
#include <helpers.h>
#include <libpq-fe.h>
@@ -62,6 +63,41 @@ namespace MyGrate::Output::Pq {
query("ROLLBACK");
}
+ FILE *
+ PqConn::beginBulkUpload(const char * schema, const char * table)
+ {
+ ResPtr res {PQexec(conn.get(), scprintf<"COPY %?.%? FROM STDIN">(schema, table).c_str()), &PQclear};
+ verify<PqErr>(PQresultStatus(res.get()) == PGRES_COPY_IN, "begin copy", res.get());
+ return fopencookie(this, "w",
+ {nullptr,
+ [](void * cookie, const char * buf, size_t size) -> ssize_t {
+ auto pq = static_cast<PqConn *>(cookie);
+ int rc;
+ while (!(rc = PQputCopyData(pq->conn.get(), buf, size))) {
+ sleep(1);
+ }
+ verify<PqErr>(rc == 1, "copy data", pq->conn.get());
+ return size;
+ },
+ nullptr,
+ [](void * cookie) -> int {
+ static_cast<PqConn *>(cookie)->endBulkUpload(nullptr);
+ return 0;
+ }});
+ }
+
+ void
+ PqConn::endBulkUpload(const char * errormsg)
+ {
+ int rc;
+ while (!(rc = PQputCopyEnd(this->conn.get(), errormsg))) {
+ sleep(1);
+ }
+ verify<PqErr>(rc == 1, "copy end rc", conn.get());
+ ResPtr res {PQgetResult(conn.get()), &PQclear};
+ verify<PqErr>(PQresultStatus(res.get()) == PGRES_COMMAND_OK, "end copy", res.get());
+ }
+
void
PqConn::notice_processor(void * p, const char * n)
{
diff --git a/lib/output/pq/pqConn.h b/lib/output/pq/pqConn.h
index 3b27f97..6ee6dc0 100644
--- a/lib/output/pq/pqConn.h
+++ b/lib/output/pq/pqConn.h
@@ -2,6 +2,7 @@
#define MYGRATE_OUTPUT_PQ_PQCONN_H
#include <cstddef>
+#include <cstdio>
#include <dbConn.h>
#include <functional>
#include <initializer_list>
@@ -36,6 +37,9 @@ namespace MyGrate::Output::Pq {
void commitTx() override;
void rollbackTx() override;
+ FILE * beginBulkUpload(const char * schema, const char * table);
+ void endBulkUpload(const char * errormsg);
+
const std::string connstr;
private: