!22 解决uint数据类型查询报错

Merge pull request !22 from zcl/bug#1

psycopg/connection_type.c

@@ -37,6 +37,8 @@
 
 #include <string.h>
 #include <ctype.h>
+#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;
 

psycopg/pqpath.c

@@ -43,7 +43,7 @@
 #include "psycopg/pgtypes.h"
 #include "psycopg/error.h"
 #include "psycopg/column.h"
-
+#include <stdlib.h>
 #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

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;

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);