Pg: Add COPY support

This commit is contained in:
Yorhel 2025-03-10 12:24:52 +01:00
parent dc752e2a23
commit d9d2ad0434
5 changed files with 301 additions and 29 deletions

View file

@ -36,6 +36,7 @@ typedef enum { PQTRANS_IDLE, PQTRANS_ACTIVE, PQTRANS_INTRANS, PQTRANS_INERROR, P
#define PG_DIAG_SOURCE_FUNCTION 'R'
#define PG_FUNCS \
X(PQbinaryTuples, int, const PGresult *) \
X(PQclear, void, PGresult *) \
X(PQclosePrepared, PGresult *, PGconn *, const char *) \
X(PQcmdTuples, char *, PGresult *) \
@ -51,9 +52,10 @@ typedef enum { PQTRANS_IDLE, PQTRANS_ACTIVE, PQTRANS_INTRANS, PQTRANS_INERROR, P
X(PQfname, char *, const PGresult *, int) \
X(PQfreemem, void, void *) \
X(PQftype, Oid, const PGresult *, int) \
X(PQgetCopyData, int, PGconn *, char **, int) \
X(PQgetResult, PGresult *, PGconn *) \
X(PQgetisnull, int, const PGresult *, int, int) \
X(PQgetlength, int, const PGresult *, int, int) \
X(PQgetResult, PGresult *, PGconn *) \
X(PQgetvalue, char *, const PGresult *, int, int) \
X(PQlibVersion, int, void) \
X(PQnfields, int, const PGresult *) \
@ -61,6 +63,8 @@ typedef enum { PQTRANS_IDLE, PQTRANS_ACTIVE, PQTRANS_INTRANS, PQTRANS_INERROR, P
X(PQntuples, int, const PGresult *) \
X(PQparamtype, Oid, const PGresult *, int) \
X(PQpipelineSync, int, PGconn *) \
X(PQputCopyData, int, PGconn *, const char *, int) \
X(PQputCopyEnd, int, PGconn *, const char *) \
X(PQresStatus, char *, ExecStatusType) \
X(PQresultErrorField, char *, const PGresult *, int) \
X(PQresultErrorMessage, char *, const PGresult *) \

View file

@ -503,3 +503,82 @@ static SV *fupg_st_kvh(pTHX_ fupg_st *st) {
}
return sv;
}
/* COPY support */
typedef struct {
SV *self;
fupg_conn *conn;
char in;
char bin;
char rddone;
char closed;
} fupg_copy;
static SV *fupg_copy_exec(pTHX_ fupg_conn *c, const char *sql) {
PGresult *r = PQexec(c->conn, sql);
if (!r) fupg_conn_croak(c, "exec");
int s = PQresultStatus(r);
switch (s) {
case PGRES_COPY_OUT:
case PGRES_COPY_IN:
break;
default: fupg_result_croak(r, "exec", sql);
}
fupg_copy *copy = safecalloc(1, sizeof(fupg_copy));
copy->conn = c;
SvREFCNT_inc(c->self);
copy->bin = !!PQbinaryTuples(r);
copy->in = s == PGRES_COPY_IN;
PQclear(r);
return fu_selfobj(copy, "FU::Pg::copy");
}
static void fupg_copy_write(pTHX_ fupg_copy *c, SV *data) {
STRLEN len;
const char *buf = c->bin ? SvPVbyte(data, len) : SvPVutf8(data, len);
if (PQputCopyData(c->conn->conn, buf, len) < 0) fupg_conn_croak(c->conn, "copy");
}
static SV *fupg_copy_read(pTHX_ fupg_copy *c, int discard) {
char *buf = NULL;
int len = PQgetCopyData(c->conn->conn, &buf, 0);
if (len == -1) {
c->rddone = 1;
return &PL_sv_undef;
} else if (len < 0) {
if (discard) c->rddone = 1;
else fupg_conn_croak(c->conn, "copy");
}
SV *r = discard ? &PL_sv_undef : newSVpvn_flags(buf, len, SVs_TEMP | (c->bin ? 0 : SVf_UTF8));
PQfreemem(buf);
return r;
}
static void fupg_copy_close(pTHX_ fupg_copy *c, int ignerror) {
if (c->closed) return;
c->closed = 1; /* Mark as closed even on error, a second attempt won't help anyway */
if (c->in && PQputCopyEnd(c->conn->conn, NULL) < 0 && !ignerror)
fupg_conn_croak(c->conn, "copyEnd");
while (!c->in && !c->rddone) fupg_copy_read(aTHX_ c, 1);
PGresult *r = PQgetResult(c->conn->conn);
if (!ignerror && !r) fupg_conn_croak(c->conn, "copyEnd");
if (!ignerror && PQresultStatus(r) != PGRES_COMMAND_OK) fupg_result_croak(r, "copy", "");
PQclear(r);
while ((r = PQgetResult(c->conn->conn))) PQclear(r);
}
static void fupg_copy_destroy(pTHX_ fupg_copy *c) {
fupg_copy_close(aTHX_ c, 1);
SvREFCNT_dec(c->conn->self);
safefree(c);
}