diff --git a/ChangeLog b/ChangeLog index e9129aec..00346cb1 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,3 +1,7 @@ +2007-02-22 Federico Di Gregorio + + * Added support for per-connection and per-cursor typecasters. + 2007-02-11 Federico Di Gregorio * psycopg/pqpath.c: ported psycopg1 patch from #135. diff --git a/examples/typecast.py b/examples/typecast.py new file mode 100644 index 00000000..2b1d37d4 --- /dev/null +++ b/examples/typecast.py @@ -0,0 +1,67 @@ +# typecast.py - example of per-cursor and per-connection typecasters. +# +# Copyright (C) 2001-2007 Federico Di Gregorio +# +# This program is free software; you can redistribute it and/or modify +# it under the terms of the GNU General Public License as published by the +# Free Software Foundation; either version 2, or (at your option) any later +# version. +# +# This program is distributed in the hope that it will be useful, but +# WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTIBILITY +# or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License +# for more details. + +## put in DSN your DSN string + +DSN = 'dbname=test' + +## don't modify anything below this line (except for experimenting) + +class SimpleQuoter(object): + def sqlquote(x=None): + return "'bar'" + +import sys +import psycopg2 +import psycopg2.extensions + +if len(sys.argv) > 1: + DSN = sys.argv[1] + +print "Opening connection using dns:", DSN +conn = psycopg2.connect(DSN) +print "Encoding for this connection is", conn.encoding + +curs = conn.cursor() +curs.execute("SELECT 'text'::text AS foo") +textoid = curs.description[0][1] +print "Oid for the text datatype is", textoid + +def castA(s, curs): + if s is not None: return "(A) " + s +TYPEA = psycopg2.extensions.new_type((textoid,), "TYPEA", castA) + +def castB(s, curs): + if s is not None: return "(B) " + s +TYPEB = psycopg2.extensions.new_type((textoid,), "TYPEB", castB) + +curs = conn.cursor() +curs.execute("SELECT 'some text.'::text AS foo") +print "Some text from plain connection:", curs.fetchone()[0] + +psycopg2.extensions.register_type(TYPEA, conn) +curs = conn.cursor() +curs.execute("SELECT 'some text.'::text AS foo") +print "Some text from connection with typecaster:", curs.fetchone()[0] + +curs = conn.cursor() +psycopg2.extensions.register_type(TYPEB, curs) +curs.execute("SELECT 'some text.'::text AS foo") +print "Some text from cursor with typecaster:", curs.fetchone()[0] + +curs = conn.cursor() +curs.execute("SELECT 'some text.'::text AS foo") +print "Some text from connection with typecaster again:", curs.fetchone()[0] + + diff --git a/psycopg/connection.h b/psycopg/connection.h index c50ce4e4..fed10c14 100644 --- a/psycopg/connection.h +++ b/psycopg/connection.h @@ -74,7 +74,11 @@ typedef struct { PyObject *exc_IntegrityError; PyObject *exc_DataError; PyObject *exc_NotSupportedError; - + + /* per-connection typecasters */ + PyObject *string_types; /* a set of typecasters for string types */ + PyObject *binary_types; /* a set of typecasters for binary types */ + } connectionObject; /* C-callable functions in connection_int.c and connection_ext.c */ diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c index ec14499d..026d9bf7 100644 --- a/psycopg/connection_type.c +++ b/psycopg/connection_type.c @@ -276,7 +276,11 @@ static struct PyMemberDef connectionObject_members[] = { "The current connection string."}, {"status", T_LONG, offsetof(connectionObject, status), RO, - "The current transaction status."}, + "The current transaction status."}, + {"string_types", T_OBJECT, offsetof(connectionObject, string_types), RO, + "A set of typecasters to convert textual values."}, + {"binary_types", T_OBJECT, offsetof(connectionObject, binary_types), RO, + "A set of typecasters to convert binary values."}, #endif {NULL} }; @@ -301,7 +305,8 @@ connection_setup(connectionObject *self, char *dsn) self->async_cursor = NULL; self->pgconn = NULL; self->mark = 0; - + self->string_types = PyDict_New(); + self->binary_types = PyDict_New(); pthread_mutex_init(&(self->lock), NULL); @@ -339,6 +344,8 @@ connection_dealloc(PyObject* obj) Py_XDECREF(self->notice_list); Py_XDECREF(self->notifies); Py_XDECREF(self->async_cursor); + Py_XDECREF(self->string_types); + Py_XDECREF(self->binary_types); pthread_mutex_destroy(&(self->lock)); diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c index c660d1d0..2953a15c 100644 --- a/psycopg/cursor_type.c +++ b/psycopg/cursor_type.c @@ -1448,7 +1448,9 @@ cursor_dealloc(PyObject* obj) Py_XDECREF(self->tuple_factory); Py_XDECREF(self->tzinfo_factory); Py_XDECREF(self->query); - + Py_XDECREF(self->string_types); + Py_XDECREF(self->binary_types); + IFCLEARPGRES(self->pgres); Dprintf("cursor_dealloc: deleted cursor object at %p, refcnt = %d", diff --git a/psycopg/pqpath.c b/psycopg/pqpath.c index d287c5f1..7e44ead9 100644 --- a/psycopg/pqpath.c +++ b/psycopg/pqpath.c @@ -501,34 +501,44 @@ _pq_fetch_tuples(cursorObject *curs) PyTuple_SET_ITEM(curs->description, i, dtitem); - /* fill the right cast function by accessing the global dictionary of - casting objects. if we got no defined cast use the default - one */ - if (!(cast = PyDict_GetItem(curs->casts, type))) { - Dprintf("_pq_fetch_tuples: cast %d not in per-cursor dict", ftype); - if (!(cast = PyDict_GetItem(psyco_types, type))) { - Dprintf("_pq_fetch_tuples: cast %d not found, using default", - PQftype(curs->pgres,i)); - cast = psyco_default_cast; - } + /* fill the right cast function by accessing three different dictionaries: + - the per-cursor dictionary, if available (can be NULL or None) + - the per-connection dictionary (always exists but can be null) + - the global dictionary (at module level) + if we get no defined cast use the default one */ + + Dprintf("_pq_fetch_tuples: looking for cast %d:", ftype); + if (curs->string_types != NULL && curs->string_types != Py_None) { + cast = PyDict_GetItem(curs->string_types, type); + Dprintf("_pq_fetch_tuples: per-cursor dict: %p", cast); } + if (cast == NULL) { + cast = PyDict_GetItem(curs->conn->string_types, type); + Dprintf("_pq_fetch_tuples: per-connection dict: %p", cast); + } + if (cast == NULL) { + cast = PyDict_GetItem(psyco_types, type); + Dprintf("_pq_fetch_tuples: global dict: %p", cast); + } + if (cast == NULL) cast = psyco_default_cast; + /* else if we got binary tuples and if we got a field that is binary use the default cast FIXME: what the hell am I trying to do here? This just can't work.. */ - else if (pgbintuples && cast == psyco_default_binary_cast) { + if (pgbintuples && cast == psyco_default_binary_cast) { Dprintf("_pq_fetch_tuples: Binary cursor and " "binary field: %i using default cast", PQftype(curs->pgres,i)); cast = psyco_default_cast; } + Dprintf("_pq_fetch_tuples: using cast at %p (%s) for type %d", cast, PyString_AS_STRING(((typecastObject*)cast)->name), PQftype(curs->pgres,i)); Py_INCREF(cast); PyTuple_SET_ITEM(curs->casts, i, cast); - /* 1/ fill the other fields */ PyTuple_SET_ITEM(dtitem, 0, PyString_FromString(PQfname(curs->pgres, i))); diff --git a/psycopg/psycopgmodule.c b/psycopg/psycopgmodule.c index a177639b..5079fc88 100644 --- a/psycopg/psycopgmodule.c +++ b/psycopg/psycopgmodule.c @@ -228,17 +228,40 @@ psyco_connect(PyObject *self, PyObject *args, PyObject *keywds) " the string representation returned by PostgreSQL (`None` if ``NULL``)\n" \ " and ``cur`` is the cursor from which data are read." +static void +_psyco_register_type_set(PyObject **dict, PyObject *type) +{ + if (*dict == NULL) + *dict = PyDict_New(); + typecast_add(type, *dict, 0); +} + static PyObject * psyco_register_type(PyObject *self, PyObject *args) { - PyObject *type; + PyObject *type, *obj; - if (!PyArg_ParseTuple(args, "O!", &typecastType, &type)) { + if (!PyArg_ParseTuple(args, "O!|O", &typecastType, &type, &obj)) { return NULL; } - typecast_add(type, 0); - + if (obj != NULL) { + if (obj->ob_type == &cursorType) { + _psyco_register_type_set(&(((cursorObject*)obj)->string_types), type); + } + else if (obj->ob_type == &connectionType) { + typecast_add(type, ((connectionObject*)obj)->string_types, 0); + } + else { + PyErr_SetString(PyExc_TypeError, + "argument 2 must be a connection, cursor or None"); + return NULL; + } + } + else { + typecast_add(type, NULL, 0); + } + Py_INCREF(Py_None); return Py_None; } diff --git a/psycopg/typecast.c b/psycopg/typecast.c index 625ba1b4..09473021 100644 --- a/psycopg/typecast.c +++ b/psycopg/typecast.c @@ -226,7 +226,7 @@ typecast_init(PyObject *dict) t = (typecastObject *)typecast_from_c(&(typecast_builtins[i]), dict); if (t == NULL) return -1; - if (typecast_add((PyObject *)t, 0) != 0) return -1; + if (typecast_add((PyObject *)t, NULL, 0) != 0) return -1; PyDict_SetItem(dict, t->name, (PyObject *)t); @@ -264,7 +264,7 @@ typecast_init(PyObject *dict) /* typecast_add - add a type object to the dictionary */ int -typecast_add(PyObject *obj, int binary) +typecast_add(PyObject *obj, PyObject *dict, int binary) { PyObject *val; Py_ssize_t len, i; @@ -274,16 +274,14 @@ typecast_add(PyObject *obj, int binary) Dprintf("typecast_add: object at %p, values refcnt = %d", obj, type->values->ob_refcnt); + if (dict == NULL) + dict = (binary ? psyco_binary_types : psyco_types); + len = PyTuple_Size(type->values); for (i = 0; i < len; i++) { val = PyTuple_GetItem(type->values, i); Dprintf("typecast_add: adding val: %ld", PyInt_AsLong(val)); - if (binary) { - PyDict_SetItem(psyco_binary_types, val, obj); - } - else { - PyDict_SetItem(psyco_types, val, obj); - } + PyDict_SetItem(dict, val, obj); } Dprintf("typecast_add: base caster: %p", type->bcast); diff --git a/psycopg/typecast.h b/psycopg/typecast.h index e42dd521..87355b40 100644 --- a/psycopg/typecast.h +++ b/psycopg/typecast.h @@ -69,7 +69,7 @@ extern PyObject *psyco_default_binary_cast; /* used by module.c to init the type system and register types */ extern int typecast_init(PyObject *dict); -extern int typecast_add(PyObject *obj, int binary); +extern int typecast_add(PyObject *obj, PyObject *dict, int binary); /* the C callable typecastObject creator function */ extern PyObject *typecast_from_c(typecastObject_initlist *type, PyObject *d); diff --git a/setup.cfg b/setup.cfg index f7d6d03e..ce6b7fa8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,5 +1,5 @@ [build_ext] -define=PSYCOPG_DEBUG,PSYCOPG_EXTENSIONS,PSYCOPG_DISPLAY_SIZE,PSYCOPG_NEW_BOOLEAN,HAVE_PQFREEMEM,HAVE_PQPROTOCOL3 +define=PSYCOPG_EXTENSIONS,PSYCOPG_DISPLAY_SIZE,PSYCOPG_NEW_BOOLEAN,HAVE_PQFREEMEM,HAVE_PQPROTOCOL3 # PSYCOPG_EXTENSIONS enables extensions to PEP-249 (you really want this) # PSYCOPG_DISPLAY_SIZE enable display size calculation (a little slower) # HAVE_PQFREEMEM should be defined on PostgreSQL >= 7.4