From 5235f6399ccdd124af516aa75176d8b0ff880822 Mon Sep 17 00:00:00 2001 From: zhouchenglong Date: Tue, 19 Mar 2024 14:32:16 +0800 Subject: [PATCH] =?UTF-8?q?=E8=A7=A3=E5=86=B3uint=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E6=9F=A5=E8=AF=A2=E6=8A=A5=E9=94=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 连接数据库后动态注册uint类型 无效修改和少提交的修改 注册uint函数封装后再加入到connection_setup函数里 函数名改成pq_get_custom_type_oid; 去掉SHOW前缀 oid改成unsigned int类型 --- psycopg/connection_type.c | 51 +++++++++++++++++++++++++++++++++++++++ psycopg/pqpath.c | 47 +++++++++++++++++++++++++++++++++++- psycopg/typecast.c | 2 +- psycopg/typecast.h | 1 + 4 files changed, 99 insertions(+), 2 deletions(-) diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index 5155c68..33f4adc 100644 --- a/psycopg/connection_type.c +++ b/psycopg/connection_type.c @@ -37,6 +37,8 @@ #include #include +#include "psycopg/typecast.h" +#include "psycopg/typecast_basic.c" extern HIDDEN const char *srv_isolevels[]; extern HIDDEN const char *srv_readonly[]; @@ -1314,6 +1316,54 @@ static struct PyGetSetDef connectionObject_getsets[] = { }; #undef EXCEPTION_GETTER + +/* register the uint typecasters */ +static int +register_type_uint(connectionObject *self, PyThreadState **tstate) +{ + int rv = -1; + typecastObject *obj = NULL; + PyObject *name = NULL, *values = NULL; + Py_ssize_t i, len = 0; + + char* uint_arr[] = {"\'uint1\'", "\'uint2\'", "\'uint4\'", "\'uint8\'"}; + int size = sizeof(uint_arr) / sizeof(uint_arr[0]); + long int* _typecast_INTEGER_types; + _typecast_INTEGER_types = (long int*)malloc((size+1)*sizeof(long int)); + + for (int i=0; i< size; i++) { + unsigned int uint_val = pq_get_custom_type_oid(self, uint_arr[i], tstate); + _typecast_INTEGER_types[i] = (long int)uint_val; + } + typecastObject_initlist _typecast_builtins[] = { + {"INTEGER", _typecast_INTEGER_types, typecast_INTEGER_cast, NULL}, + }; + name = Text_FromUTF8(_typecast_builtins->name); + if (!name) goto end; + + while (_typecast_builtins->values[len] != 0) len++; + + values = PyTuple_New(len); + if (!values) goto end; + + for (i = 0; i < len ; i++) { + PyTuple_SET_ITEM(values, i, PyInt_FromLong(_typecast_builtins->values[i])); + } + obj = (typecastObject *)typecast_new(name, values, NULL, NULL); + if (obj) { + obj->ccast = _typecast_builtins->cast; + obj->pcast = NULL; + } + if (typecast_add((PyObject *)obj, self->string_types, 0) < 0) { goto end; } + rv = 0; + free(_typecast_INTEGER_types); +end: + Py_XDECREF(values); + Py_XDECREF(name); + return rv; +} + + /* initialization and finalization methods */ static int @@ -1364,6 +1414,7 @@ connection_setup(connectionObject *self, const char *dsn, long int async) Py_BEGIN_ALLOW_THREADS; pthread_mutex_lock(&self->lock); sql_compatibility_value = pq_get_guc_locked(self, "sql_compatibility", &_save); + if (register_type_uint(self, &_save)) { goto exit; } pthread_mutex_unlock(&self->lock); Py_END_ALLOW_THREADS; diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c index ec02067..723ee6f 100644 --- a/psycopg/pqpath.c +++ b/psycopg/pqpath.c @@ -43,7 +43,7 @@ #include "psycopg/pgtypes.h" #include "psycopg/error.h" #include "psycopg/column.h" - +#include #include "psycopg/libpq_support.h" #include "libpq-fe.h" @@ -607,6 +607,51 @@ cleanup: return rv; } + + +unsigned int +pq_get_custom_type_oid(connectionObject *conn, const char *param, PyThreadState **tstate) +{ + char query[256]; + int size; + unsigned int rv = 0; + + size = PyOS_snprintf(query, sizeof(query), "select oid from pg_type where typname= %s", param); + if (size < 0 || (size_t)size >= sizeof(query)) { + conn_set_error(conn, "query too large"); + goto cleanup; + } + + if (!psyco_green()) { + conn_set_result(conn, PQexec(conn->pgconn, query)); + } else { + PyEval_RestoreThread(*tstate); + conn_set_result(conn, psyco_exec_green(conn, query)); + *tstate = PyEval_SaveThread(); + } + + if (!conn->pgres) { + Dprintf("pq_get_custom_type_oid: PQexec returned NULL"); + PyEval_RestoreThread(*tstate); + if (!PyErr_Occurred()) { + conn_set_error(conn, PQerrorMessage(conn->pgconn)); + } + *tstate = PyEval_SaveThread(); + goto cleanup; + } + if (PQresultStatus(conn->pgres) != PGRES_TUPLES_OK) { + Dprintf("pq_get_custom_type_oid: result was not TUPLES_OK (%s)", + PQresStatus(PQresultStatus(conn->pgres))); + goto cleanup; + } + + rv = atoi(strdup(PQgetvalue(conn->pgres, 0, 0))); + CLEARPGRES(conn->pgres); + +cleanup: + return rv; +} + /* Set a session parameter. * * The function should be called on a locked connection without diff --git a/psycopg/typecast.c b/psycopg/typecast.c index c1facc4..573084d 100644 --- a/psycopg/typecast.c +++ b/psycopg/typecast.c @@ -459,7 +459,7 @@ PyTypeObject typecastType = { 0, /*tp_new*/ }; -static PyObject * +PyObject * typecast_new(PyObject *name, PyObject *values, PyObject *cast, PyObject *base) { typecastObject *obj; diff --git a/psycopg/typecast.h b/psycopg/typecast.h index 050345f..84af98f 100644 --- a/psycopg/typecast.h +++ b/psycopg/typecast.h @@ -89,3 +89,4 @@ HIDDEN PyObject *typecast_cast( PyObject *self, const char *str, Py_ssize_t len, PyObject *curs); #endif /* !defined(PSYCOPG_TYPECAST_H) */ +PyObject *typecast_new(PyObject *name, PyObject *values, PyObject *cast, PyObject *base); \ No newline at end of file