From 7b76d947198fbb025a201a77f76cc4462d557d14 Mon Sep 17 00:00:00 2001 From: Yorhel Date: Sat, 8 Feb 2025 17:24:41 +0100 Subject: [PATCH] pg: Add dynamic type loading & support enum types Least efficient way to support enums, really. *shrug* --- c/libpq.h | 1 + c/pgconn.c | 75 +++++++++++++++++++++++++++++++++++++++------ c/pgtypes.c | 31 +++++++++++-------- t/pgtypes-dynamic.t | 20 ++++++++++++ t/pgtypes.t | 16 ++++++---- 5 files changed, 114 insertions(+), 29 deletions(-) create mode 100644 t/pgtypes-dynamic.t diff --git a/c/libpq.h b/c/libpq.h index c62b3f5..e3cf312 100644 --- a/c/libpq.h +++ b/c/libpq.h @@ -43,6 +43,7 @@ typedef enum { PQTRANS_IDLE, PQTRANS_ACTIVE, PQTRANS_INTRANS, PQTRANS_INERROR, P X(PQdescribePrepared, PGresult *, PGconn *, const char *) \ X(PQerrorMessage, char *, const PGconn *) \ X(PQexec, PGresult *, PGconn *, const char *) \ + X(PQexecParams, PGresult *, PGconn *, const char *, int, const Oid *, const char * const *, const int *, const int *, int) \ X(PQexecPrepared, PGresult *, PGconn *, const char *, int, const char * const *, const int *, const int *, int) \ X(PQfinish, void, PGconn *) \ X(PQfmod, int, const PGresult *, int) \ diff --git a/c/pgconn.c b/c/pgconn.c index e6fd3b5..c12c465 100644 --- a/c/pgconn.c +++ b/c/pgconn.c @@ -10,6 +10,8 @@ typedef struct { UV cookie_counter; UV cookie; /* currently active transaction object; 0 = none active */ int stflags; + int ntypes; + fupg_type *types; fustr buf; /* Scratch space for query params */ } fupg_conn; @@ -154,6 +156,8 @@ static SV *fupg_connect(pTHX_ const char *str) { c->conn = conn; c->prep_counter = c->cookie_counter = c->cookie = 0; c->stflags = 0; + c->ntypes = 0; + c->types = NULL; fustr_init(&c->buf, NULL, SIZE_MAX); return fu_selfobj(c, "FU::PG::conn"); } @@ -176,6 +180,8 @@ static void fupg_conn_disconnect(fupg_conn *c) { static void fupg_conn_destroy(fupg_conn *c) { PQfinish(c->conn); + if (c->buf.sv) SvREFCNT_dec(c->buf.sv); + safefree(c->types); safefree(c); } @@ -252,6 +258,54 @@ static void fupg_txn_destroy(pTHX_ fupg_txn *t) { safefree(t); } +/* XXX: It feels a bit wasteful to load *all* types; even on an empty database + * that's ~55k of data, but it's easier and (potentially) faster than fetching + * each type seperately as we encounter them. + */ +static void fupg_refresh_types(pTHX_ fupg_conn *c) { + safefree(c->types); + c->types = 0; + c->ntypes = 0; + + const char *sql = "SELECT oid, typname, typtype FROM pg_type ORDER BY oid"; + PGresult *r = PQexecParams(c->conn, sql, 0, NULL, NULL, NULL, NULL, 1); + if (!r) fupg_conn_croak(c, "exec"); + if (PQresultStatus(r) != PGRES_TUPLES_OK) fupg_result_croak(r, "exec", sql); + + c->ntypes = PQntuples(r); + c->types = calloc(c->ntypes, sizeof(*c->types)); + int i; + for (i=0; intypes; i++) { + fupg_type *t = c->types + i; + t->oid = __builtin_bswap32(*((Oid *)PQgetvalue(r, i, 0))); + snprintf(t->name, sizeof(t->name), "%s", PQgetvalue(r, i, 1)); + char typ = *PQgetvalue(r, i, 2); + + /* enum, can use text send/recv */ + if (typ == 'e') { + t->send = fupg_send_text; + t->recv = fupg_recv_text; + continue; + } + /* TODO: Array types, records, custom overrides, by-name lookup for dynamic-oid types */ + const fupg_type *builtin = fupg_builtin_byoid(t->oid); + if (builtin) { + t->send = builtin->send; + t->recv = builtin->recv; + } + } + PQclear(r); +} + +static const fupg_type *fupg_lookup_type(pTHX_ fupg_conn *c, int *refresh_done, Oid oid) { + const fupg_type *t = NULL; + if (c->types && (t = fupg_type_byoid(c->types, c->ntypes, oid))) return t; + if ((t = fupg_builtin_byoid(oid))) return t; + if (*refresh_done) return NULL; + *refresh_done = 1; + fupg_refresh_types(c); + return fupg_type_byoid(c->types, c->ntypes, oid); +} @@ -358,7 +412,7 @@ static void fupg_st_check_dupcols(pTHX_ PGresult *r) { SvREFCNT_dec((SV *)hv); } -static void fupg_params_setup(pTHX_ fupg_st *st) { +static void fupg_params_setup(pTHX_ fupg_st *st, int *refresh_done) { int i; st->param_values = safecalloc(st->nbind, sizeof(*st->param_values)); if (st->stflags & FUPG_TEXT_PARAMS) { @@ -379,9 +433,9 @@ static void fupg_params_setup(pTHX_ fupg_st *st) { } fupg_send send; send.oid = PQparamtype(st->describe, i); - const fupg_core_type *t = fupg_core_type_byoid(send.oid); - if (!t) - fu_confess("Unable to use type oid %u as bind parameter", send.oid); + const fupg_type *t = fupg_lookup_type(aTHX_ st->conn, refresh_done, send.oid); + if (!t) fu_confess("No type found with oid %u", send.oid); + if (!t->send) fu_confess("Unable to use type '%s' (oid %u) as bind parameter", t->name, t->oid); send.name = t->name; send.fn = t->send; off = fustr_len(buf); @@ -401,7 +455,7 @@ static void fupg_params_setup(pTHX_ fupg_st *st) { } } -static void fupg_results_setup(pTHX_ fupg_st *st) { +static void fupg_results_setup(pTHX_ fupg_st *st, int *refresh_done) { int i; st->recv = safecalloc(st->nfields, sizeof(*st->recv)); if (st->stflags & FUPG_TEXT_RESULTS) { @@ -413,8 +467,9 @@ static void fupg_results_setup(pTHX_ fupg_st *st) { for (i=0; infields; i++) { fupg_recv *r = st->recv + i; r->oid = PQftype(st->result, i); - const fupg_core_type *t = fupg_core_type_byoid(r->oid); - if (!t) fu_confess("Unable to receive query results of type oid %u", r->oid); + const fupg_type *t = fupg_lookup_type(aTHX_ st->conn, refresh_done, r->oid); + if (!t) fu_confess("No type found with oid %u", r->oid); + if (!t->recv) fu_confess("Unable to receive data of type '%s' (oid %u)", t->name, t->oid); r->name = t->name; r->fn = t->recv; } @@ -430,7 +485,8 @@ static void fupg_st_execute(pTHX_ fupg_st *st) { fupg_st_prepare(aTHX_ st); if (PQnparams(st->describe) != st->nbind) fu_confess("Statement expects %d bind parameters but %d were given", PQnparams(st->describe), st->nbind); - fupg_params_setup(aTHX_ st); + int refresh_done = 0; + fupg_params_setup(aTHX_ st, &refresh_done); /* I'm not super fond of this approach. Storing the full query results in a * PGresult involves unnecessary parsing, memory allocation and copying. @@ -458,7 +514,7 @@ static void fupg_st_execute(pTHX_ fupg_st *st) { st->result = r; st->nfields = PQnfields(r); - fupg_results_setup(aTHX_ st); + fupg_results_setup(aTHX_ st, &refresh_done); } static SV *fupg_st_getval(pTHX_ fupg_st *st, int row, int col) { @@ -547,4 +603,3 @@ static void fupg_st_destroy(fupg_st *st) { /* TODO: $st->alla, allh, flat, kvv, kva, kvh */ /* TODO: Prepared statement caching */ -/* TODO: Custom type handling */ diff --git a/c/pgtypes.c b/c/pgtypes.c index 284b8c2..ac5b2d0 100644 --- a/c/pgtypes.c +++ b/c/pgtypes.c @@ -24,10 +24,10 @@ struct fupg_recv { typedef struct { Oid oid; - char name[16]; /* Postgres has a 64 byte limit on names, but this is sufficient for the core types listed here */ + char name[64]; fupg_send_fn send; fupg_recv_fn recv; -} fupg_core_type; +} fupg_type; @@ -132,9 +132,10 @@ SENDFN(char) { fustr_write(out, buf, len); } -/* Works for many text-based column types. - * Assumes client_encoding=utf8, will create a mess otherwise */ +/* Works for many text-based column types, including receiving any value in the text format */ RECVFN(text) { + if (!is_c9strict_utf8_string((const U8*)buf, len)) + fu_confess("Received invalid UTF-8 for type '%s' (oid %u)", ctx->name, ctx->oid); return newSVpvn_utf8(buf, len, 1); } @@ -248,7 +249,7 @@ SENDFN(jsonpath) { Ordered by oid to support binary search. (name is only used when formatting error messages, for now) */ -#define CORETYPES \ +#define BUILTINS \ B( 16, "bool", bool )\ B( 17, "bytea", bytea )\ B( 18, "char", char )\ @@ -320,24 +321,28 @@ SENDFN(jsonpath) { /* 5038 pg_snapshot */\ /* 5069 xid8 */ -static const fupg_core_type fupg_core_types[] = { +static const fupg_type fupg_builtin[] = { #define B(oid, name, fun) { oid, name"\0", fupg_send_##fun, fupg_recv_##fun }, - CORETYPES + BUILTINS #undef B }; -#undef CORETYPES +#undef BUILTINS -#define FUPG_CORE_TYPES (sizeof(fupg_core_types) / sizeof(fupg_core_type)) +#define FUPG_BUILTIN (sizeof(fupg_builtin) / sizeof(fupg_type)) -static const fupg_core_type *fupg_core_type_byoid(Oid oid) { - int i, b = 0, e = FUPG_CORE_TYPES-1; +static const fupg_type *fupg_type_byoid(const fupg_type *list, int len, Oid oid) { + int i, b = 0, e = len-1; while (b <= e) { i = b + (e - b)/2; - if (fupg_core_types[i].oid == oid) return fupg_core_types+i; - if (fupg_core_types[i].oid < oid) b = i+1; + if (list[i].oid == oid) return list+i; + if (list[i].oid < oid) b = i+1; else e = i-1; } return NULL; } + +static const fupg_type *fupg_builtin_byoid(Oid oid) { + return fupg_type_byoid(fupg_builtin, FUPG_BUILTIN, oid); +} diff --git a/t/pgtypes-dynamic.t b/t/pgtypes-dynamic.t new file mode 100644 index 0000000..fddcda8 --- /dev/null +++ b/t/pgtypes-dynamic.t @@ -0,0 +1,20 @@ +use v5.36; +use Test::More; + +plan skip_all => $@ if !eval { require FU::PG; } && $@ =~ /Unable to load libpq/; +die $@ if $@; +plan skip_all => 'Please set FU_TEST_DB to a PostgreSQL connection string to run these tests' if !$ENV{FU_TEST_DB}; + +my $conn = FU::PG->connect($ENV{FU_TEST_DB}); + +ok !eval { $conn->q('SELECT $1::aclitem', '')->exec; 1 }; +like $@, qr/Unable to use type/; + +{ + my $txn = $conn->txn; + $txn->exec("CREATE TYPE fupg_test_enum AS ENUM('a', 'b', 'ccccccccccccccccccc')"); + is $txn->q("SELECT 'a'::fupg_test_enum")->val, 'a'; + is $txn->q('SELECT $1::fupg_test_enum', 'ccccccccccccccccccc')->val, 'ccccccccccccccccccc'; +} + +done_testing; diff --git a/t/pgtypes.t b/t/pgtypes.t index 39dafc8..50a1cd7 100644 --- a/t/pgtypes.t +++ b/t/pgtypes.t @@ -17,23 +17,27 @@ sub v($type, $p_in, @args) { my $s_in = @args > 1 && defined $args[1] ? $args[1] : $p_in; my $s_out = @args > 2 && defined $args[2] ? $args[2] : $s_in; + my $test = "$type $s_in" =~ s/\n/\\n/rg; + utf8::encode($test); { my $res = $conn->q("SELECT \$1::$type", $s_in)->text_params->val; - ok is_bool($res), "$type $s_in is bool" if $type eq 'bool'; - ok created_as_number($res), "$type $s_in is number" if $type =~ /^int/; - is_deeply $res, $p_out, "$type $s_in text->bin" =~ s/\n/\\n/rg; + ok is_bool($res), "$test is bool" if $type eq 'bool'; + ok created_as_number($res), "$test is number" if $type =~ /^(int|float)/; + is_deeply $res, $p_out, "$test text->bin"; } { my $res = $conn->q("SELECT \$1::$type", $p_in)->text_results->val; - is $res, $s_out, "$type $s_out bin->text" =~ s/\n/\\n/rg; + is $res, $s_out, "$test bin->text"; } { my $res = $conn->q("SELECT \$1::$type", $p_in)->val; - is_deeply $res, $p_out, "$type $s_in bin->bin" =~ s/\n/\\n/rg; + is_deeply $res, $p_out, "$test bin->bin"; } } sub f($type, $p_in) { - ok !eval { $conn->q("SELECT \$1::$type", $p_in)->val; 1 }, "$type $p_in fail"; + my $test = "$type $p_in" =~ s/\n/\\n/rg; + utf8::encode($test); + ok !eval { $conn->q("SELECT \$1::$type", $p_in)->val; 1 }, "$test fail"; } v bool => true, undef, 1, 't';