trunk/psycopg2/psycopg/adapter_qstring.c

   1 /* adapter_qstring.c - QuotedString objects
   2  *
   3  * Copyright (C) 2003-2004 Federico Di Gregorio <fog@debian.org>
   4  *
   5  * This file is part of psycopg.
   6  *
   7  * This program is free software; you can redistribute it and/or
   8  * modify it under the terms of the GNU General Public License
   9  * as published by the Free Software Foundation; either version 2,
  10  * or (at your option) any later version.
  11  *
  12  * This program is distributed in the hope that it will be useful,
  13  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15  * GNU General Public License for more details.
  16  *
  17  * You should have received a copy of the GNU General Public License
  18  * along with this program; if not, write to the Free Software
  19  * Foundation, 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
  20  */
  21
  22 #include <Python.h>
  23 #include <structmember.h>
  24 #include <stringobject.h>
  25
  26 #include <libpq-fe.h>
  27 #include <string.h>
  28
  29 #define PSYCOPG_MODULE
  30 #include "psycopg/config.h"
  31 #include "psycopg/python.h"
  32 #include "psycopg/psycopg.h"
  33 #include "psycopg/connection.h"
  34 #include "psycopg/adapter_qstring.h"
  35 #include "psycopg/microprotocols_proto.h"
  36
  37
  38 /** the quoting code */
  39
  40 #ifndef PSYCOPG_OWN_QUOTING
  41 static size_t
  42 qstring_escape(char *to, char *from, size_t len, PGconn *conn)
  43 {
  44 #if PG_MAJOR_VERSION > 8 || \
  45  (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION > 1) || \
  46  (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION == 1 && PG_PATCH_VERSION >= 4)
  47     int err;
  48     if (conn)
  49         return PQescapeStringConn(conn, to, from, len, &err);
  50     else
  51 #endif
  52         return PQescapeString(to, from, len);
  53 }
  54 #else
  55 static size_t
  56 qstring_escape(char *to, char *from, size_t len, PGconn *conn)
  57 {
  58     int i, j;
  59
  60     for (i=0, j=0; i<len; i++) {
  61         switch(from[i]) {
  62
  63         case '\'':
  64             to[j++] = '\'';
  65             to[j++] = '\'';
  66             break;
  67
  68         case '\\':
  69             to[j++] = '\\';
  70             to[j++] = '\\';
  71             break;
  72
  73         case '\0':
  74             /* do nothing, embedded \0 are discarded */
  75             break;
  76
  77         default:
  78             to[j++] = from[i];
  79         }
  80     }
  81     to[j] = '\0';
  82
  83     Dprintf("qstring_quote: to = %s", to);
  84     return strlen(to);
  85 }
  86 #endif
  87
  88 /* qstring_quote - do the quote process on plain and unicode strings */
  89
  90 static PyObject *
  91 qstring_quote(qstringObject *self)
  92 {
  93     PyObject *str;
  94     char *s, *buffer;
  95     int len;
  96
  97     /* if the wrapped object is an unicode object we can encode it to match
  98        self->encoding but if the encoding is not specified we don't know what
  99        to do and we raise an exception */
 100
 101     /* TODO: we need a real translation table from postgres encoding names to
 102            python ones here */
 103
 104     if (PyUnicode_Check(self->wrapped) && self->encoding) {
 105         PyObject *enc = PyDict_GetItemString(psycoEncodings, self->encoding);
 106         /* note that pgenc is a borrowed reference */
 107
 108         if (enc) {
 109             char *s = PyString_AsString(enc);
 110             Dprintf("qstring_quote: encoding unicode object to %s", s);
 111             str = PyUnicode_AsEncodedString(self->wrapped, s, NULL);
 112             Dprintf("qstring_quote: got encoded object at %p", str);
 113             if (str == NULL) return NULL;
 114         }
 115         else {
 116             /* can't find the right encoder, raise exception */
 117             PyErr_Format(InterfaceError,
 118                          "can't encode unicode string to %s", self->encoding);
 119             return NULL;
 120         }
 121     }
 122
 123     /* if the wrapped object is a simple string, we don't know how to
 124        (re)encode it, so we pass it as-is */
 125     else if (PyString_Check(self->wrapped)) {
 126         str = self->wrapped;
 127         /* INCREF to make it ref-wise identical to unicode one */
 128         Py_INCREF(str);
 129     }
 130
 131     /* if the wrapped object is not a string, this is an error */
 132     else {
 133         PyErr_SetString(PyExc_TypeError,
 134                         "can't quote non-string object (or missing encoding)");
 135         return NULL;
 136     }
 137
 138     /* encode the string into buffer */
 139     PyString_AsStringAndSize(str, &s, &len);
 140
 141     buffer = (char *)PyMem_Malloc((len*2+3) * sizeof(char));
 142     if (buffer == NULL) {
 143         Py_DECREF(str);
 144         PyErr_NoMemory();
 145         return NULL;
 146     }
 147
 148     Py_BEGIN_ALLOW_THREADS;
 149     len = qstring_escape(buffer+1, s, len,
 150         self->conn ? ((connectionObject*)self->conn)->pgconn : NULL);
 151     buffer[0] = '\'' ; buffer[len+1] = '\'';
 152     Py_END_ALLOW_THREADS;
 153
 154     self->buffer = PyString_FromStringAndSize(buffer, len+2);
 155     PyMem_Free(buffer);
 156     Py_DECREF(str);
 157
 158     return self->buffer;
 159 }
 160
 161 /* qstring_str, qstring_getquoted - return result of quoting */
 162
 163 static PyObject *
 164 qstring_str(qstringObject *self)
 165 {
 166     if (self->buffer == NULL) {
 167         qstring_quote(self);
 168     }
 169     Py_XINCREF(self->buffer);
 170     return self->buffer;
 171 }
 172
 173 PyObject *
 174 qstring_getquoted(qstringObject *self, PyObject *args)
 175 {
 176     if (!PyArg_ParseTuple(args, "")) return NULL;
 177     return qstring_str(self);
 178 }
 179
 180 PyObject *
 181 qstring_prepare(qstringObject *self, PyObject *args)
 182 {
 183     connectionObject *conn;
 184
 185     if (!PyArg_ParseTuple(args, "O", &conn))
 186         return NULL;
 187
 188     /* we bother copying the encoding only if the wrapped string is unicode,
 189        we don't need the encoding if that's not the case */
 190     if (PyUnicode_Check(self->wrapped)) {
 191         if (self->encoding) free(self->encoding);
 192         self->encoding = strdup(conn->encoding);
 193         Dprintf("qstring_prepare: set encoding to %s", conn->encoding);
 194     }
 195
 196     Py_XDECREF(self->conn);
 197     if (conn) {
 198         self->conn = (PyObject*)conn;
 199         Py_INCREF(self->conn);
 200     }
 201
 202     Py_INCREF(Py_None);
 203     return Py_None;
 204 }
 205
 206 PyObject *
 207 qstring_conform(qstringObject *self, PyObject *args)
 208 {
 209     PyObject *res, *proto;
 210
 211     if (!PyArg_ParseTuple(args, "O", &proto)) return NULL;
 212
 213     if (proto == (PyObject*)&isqlquoteType)
 214         res = (PyObject*)self;
 215     else
 216         res = Py_None;
 217
 218     Py_INCREF(res);
 219     return res;
 220 }
 221
 222 /** the QuotedString object **/
 223
 224 /* object member list */
 225
 226 static struct PyMemberDef qstringObject_members[] = {
 227     {"adapted", T_OBJECT, offsetof(qstringObject, wrapped), RO},
 228     {"buffer", T_OBJECT, offsetof(qstringObject, buffer), RO},
 229     {"encoding", T_STRING, offsetof(qstringObject, encoding), RO},
 230     {NULL}
 231 };
 232
 233 /* object method table */
 234
 235 static PyMethodDef qstringObject_methods[] = {
 236     {"getquoted", (PyCFunction)qstring_getquoted, METH_VARARGS,
 237      "getquoted() -> wrapped object value as SQL-quoted string"},
 238     {"prepare", (PyCFunction)qstring_prepare, METH_VARARGS,
 239      "prepare(conn) -> set encoding to conn->encoding and store conn"},
 240     {"__conform__", (PyCFunction)qstring_conform, METH_VARARGS, NULL},
 241     {NULL}  /* Sentinel */
 242 };
 243
 244 /* initialization and finalization methods */
 245
 246 static int
 247 qstring_setup(qstringObject *self, PyObject *str, char *enc)
 248 {
 249     Dprintf("qstring_setup: init qstring object at %p, refcnt = %d",
 250             self, ((PyObject *)self)->ob_refcnt);
 251
 252     self->buffer = NULL;
 253     self->conn = NULL;
 254
 255     /* FIXME: remove this orrible strdup */
 256     if (enc) self->encoding = strdup(enc);
 257
 258     self->wrapped = str;
 259     Py_INCREF(self->wrapped);
 260
 261     Dprintf("qstring_setup: good qstring object at %p, refcnt = %d",
 262             self, ((PyObject *)self)->ob_refcnt);
 263     return 0;
 264 }
 265
 266 static void
 267 qstring_dealloc(PyObject* obj)
 268 {
 269     qstringObject *self = (qstringObject *)obj;
 270
 271     Py_XDECREF(self->wrapped);
 272     Py_XDECREF(self->buffer);
 273     Py_XDECREF(self->conn);
 274
 275     if (self->encoding) free(self->encoding);
 276
 277     Dprintf("qstring_dealloc: deleted qstring object at %p, refcnt = %d",
 278             obj, obj->ob_refcnt);
 279
 280     obj->ob_type->tp_free(obj);
 281 }
 282
 283 static int
 284 qstring_init(PyObject *obj, PyObject *args, PyObject *kwds)
 285 {
 286     PyObject *str;
 287     char *enc = "latin-1"; /* default encoding as in Python */
 288
 289     if (!PyArg_ParseTuple(args, "O|s", &str, &enc))
 290         return -1;
 291
 292     return qstring_setup((qstringObject *)obj, str, enc);
 293 }
 294
 295 static PyObject *
 296 qstring_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 297 {
 298     return type->tp_alloc(type, 0);
 299 }
 300
 301 static void
 302 qstring_del(PyObject* self)
 303 {
 304     PyObject_Del(self);
 305 }
 306
 307 static PyObject *
 308 qstring_repr(qstringObject *self)
 309 {
 310     return PyString_FromFormat("<psycopg2._psycopg.QuotedString object at %p>",
 311                                 self);
 312 }
 313
 314 /* object type */
 315
 316 #define qstringType_doc \
 317 "QuotedString(str, enc) -> new quoted object with 'enc' encoding"
 318
 319 PyTypeObject qstringType = {
 320     PyObject_HEAD_INIT(NULL)
 321     0,
 322     "psycopg2._psycopg.QuotedString",
 323     sizeof(qstringObject),
 324     0,
 325     qstring_dealloc, /*tp_dealloc*/
 326     0,          /*tp_print*/
 327     0,          /*tp_getattr*/
 328     0,          /*tp_setattr*/
 329
 330     0,          /*tp_compare*/
 331     (reprfunc)qstring_repr, /*tp_repr*/
 332     0,          /*tp_as_number*/
 333     0,          /*tp_as_sequence*/
 334     0,          /*tp_as_mapping*/
 335     0,          /*tp_hash */
 336
 337     0,          /*tp_call*/
 338     (reprfunc)qstring_str, /*tp_str*/
 339     0,          /*tp_getattro*/
 340     0,          /*tp_setattro*/
 341     0,          /*tp_as_buffer*/
 342
 343     Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
 344
 345     qstringType_doc, /*tp_doc*/
 346
 347     0,          /*tp_traverse*/
 348     0,          /*tp_clear*/
 349
 350     0,          /*tp_richcompare*/
 351     0,          /*tp_weaklistoffset*/
 352
 353     0,          /*tp_iter*/
 354     0,          /*tp_iternext*/
 355
 356     /* Attribute descriptor and subclassing stuff */
 357
 358     qstringObject_methods, /*tp_methods*/
 359     qstringObject_members, /*tp_members*/
 360     0,          /*tp_getset*/
 361     0,          /*tp_base*/
 362     0,          /*tp_dict*/
 363
 364     0,          /*tp_descr_get*/
 365     0,          /*tp_descr_set*/
 366     0,          /*tp_dictoffset*/
 367
 368     qstring_init, /*tp_init*/
 369     0, /*tp_alloc  will be set to PyType_GenericAlloc in module init*/
 370     qstring_new, /*tp_new*/
 371     (freefunc)qstring_del, /*tp_free  Low-level free-memory routine */
 372     0,          /*tp_is_gc For PyObject_IS_GC */
 373     0,          /*tp_bases*/
 374     0,          /*tp_mro method resolution order */
 375     0,          /*tp_cache*/
 376     0,          /*tp_subclasses*/
 377     0           /*tp_weaklist*/
 378 };
 379
 380
 381 /** module-level functions **/
 382
 383 PyObject *
 384 psyco_QuotedString(PyObject *module, PyObject *args)
 385 {
 386     PyObject *str;
 387     char *enc = "latin-1"; /* default encoding as in Python */
 388
 389     if (!PyArg_ParseTuple(args, "O|s", &str, &enc))
 390         return NULL;
 391
 392     return PyObject_CallFunction((PyObject *)&qstringType, "Os", str, enc);
 393 }