psycopg2/psycopg/adapter_binary.c

   1 /* adapter_binary.c - Binary objects
   2  *
   3  * Copyright (C) 2003 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_binary.h"
  35 #include "psycopg/microprotocols_proto.h"
  36
  37 /** the quoting code */
  38
  39 #ifndef PSYCOPG_OWN_QUOTING
  40 static unsigned char *
  41 binary_escape(unsigned char *from, unsigned int from_length,
  42                unsigned int *to_length, 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     if (conn)
  48         return PQescapeByteaConn(conn, from, from_length, to_length);
  49     else
  50 #endif
  51         return PQescapeBytea(from, from_length, to_length);
  52 }
  53 #else
  54 static unsigned char *
  55 binary_escape(unsigned char *from, unsigned int from_length,
  56                unsigned int *to_length, PGconn *conn)
  57 {
  58     unsigneed char *quoted, *chptr, *newptr;
  59     int i, space, new_space;
  60
  61     space = from_length + 2;
  62
  63     Py_BEGIN_ALLOW_THREADS;
  64
  65     quoted = (unsigned char*)calloc(space, sizeof(char));
  66     if (quoted == NULL) return NULL;
  67
  68     chptr = quoted;
  69
  70     for (i=0; i < len; i++) {
  71         if (chptr - quoted > space - 6) {
  72             new_space  =  space * ((space) / (i + 1)) + 2 + 6;
  73             if (new_space - space < 1024) space += 1024;
  74             else space = new_space;
  75             newptr = (unsigned char *)realloc(quoted, space);
  76             if (newptr == NULL) {
  77                 free(quoted);
  78                 return NULL;
  79             }
  80             /* chptr has to be moved to the new location*/
  81             chptr = newptr + (chptr - quoted);
  82             quoted = newptr;
  83             Dprintf("binary_escape: reallocated %i bytes at %p", space,quoted);
  84         }
  85         if (from[i]) {
  86             if (from[i] >= ' ' && from[i] <= '~') {
  87                 if (from[i] == '\'') {
  88                     *chptr = '\'';
  89                     chptr++;
  90                     *chptr = '\'';
  91                     chptr++;
  92                 }
  93                 else if (from[i] == '\\') {
  94                     memcpy(chptr, "\\\\\\\\", 4);
  95                     chptr += 4;
  96                 }
  97                 else {
  98                     /* leave it as it is if ascii printable */
  99                     *chptr = from[i];
 100                     chptr++;
 101                 }
 102             }
 103             else {
 104                 unsigned char c;
 105
 106                 /* escape to octal notation \nnn */
 107                 *chptr++ = '\\';
 108                 *chptr++ = '\\';
 109                 c = from[i];
 110                 *chptr = ((c >> 6) & 0x07) + 0x30; chptr++;
 111                 *chptr = ((c >> 3) & 0x07) + 0x30; chptr++;
 112                 *chptr = ( c       & 0x07) + 0x30; chptr++;
 113             }
 114         }
 115         else {
 116             /* escape null as \\000 */
 117             memcpy(chptr, "\\\\000", 5);
 118             chptr += 5;
 119         }
 120     }
 121     *chptr = '\0';
 122
 123     Py_END_ALLOW_THREADS;
 124
 125     *to_size = chptr - quoted + 1;
 126     return quoted;
 127 }
 128 #endif
 129
 130 /* binary_quote - do the quote process on plain and unicode strings */
 131
 132 static PyObject *
 133 binary_quote(binaryObject *self)
 134 {
 135     char *to;
 136     const char *buffer;
 137     int buffer_len;
 138     size_t len = 0;
 139
 140     /* if we got a plain string or a buffer we escape it and save the buffer */
 141     if (PyString_Check(self->wrapped) || PyBuffer_Check(self->wrapped)) {
 142         /* escape and build quoted buffer */
 143         PyObject_AsCharBuffer(self->wrapped, &buffer, &buffer_len);
 144
 145         to = (char *)binary_escape((unsigned char*)buffer, buffer_len, &len,
 146             self->conn ? ((connectionObject*)self->conn)->pgconn : NULL);
 147         if (to == NULL) {
 148             PyErr_NoMemory();
 149             return NULL;
 150         }
 151
 152         if (len > 0)
 153             self->buffer = PyString_FromFormat("'%s'", to);
 154         else
 155             self->buffer = PyString_FromString("''");
 156         PQfreemem(to);
 157     }
 158
 159     /* if the wrapped object is not a string or a buffer, this is an error */
 160     else {
 161         PyErr_SetString(PyExc_TypeError, "can't escape non-string object");
 162         return NULL;
 163     }
 164
 165     return self->buffer;
 166 }
 167
 168 /* binary_str, binary_getquoted - return result of quoting */
 169
 170 static PyObject *
 171 binary_str(binaryObject *self)
 172 {
 173     if (self->buffer == NULL) {
 174         binary_quote(self);
 175     }
 176     Py_XINCREF(self->buffer);
 177     return self->buffer;
 178 }
 179
 180 PyObject *
 181 binary_getquoted(binaryObject *self, PyObject *args)
 182 {
 183     if (!PyArg_ParseTuple(args, "")) return NULL;
 184     return binary_str(self);
 185 }
 186
 187 PyObject *
 188 binary_prepare(binaryObject *self, PyObject *args)
 189 {
 190     connectionObject *conn;
 191
 192     if (!PyArg_ParseTuple(args, "O", &conn))
 193         return NULL;
 194
 195     Py_XDECREF(self->conn);
 196     if (conn) {
 197         self->conn = (PyObject*)conn;
 198         Py_INCREF(self->conn);
 199     }
 200
 201     Py_INCREF(Py_None);
 202     return Py_None;
 203 }
 204
 205 PyObject *
 206 binary_conform(binaryObject *self, PyObject *args)
 207 {
 208     PyObject *res, *proto;
 209
 210     if (!PyArg_ParseTuple(args, "O", &proto)) return NULL;
 211
 212     if (proto == (PyObject*)&isqlquoteType)
 213         res = (PyObject*)self;
 214     else
 215         res = Py_None;
 216
 217     Py_INCREF(res);
 218     return res;
 219 }
 220
 221 /** the Binary object **/
 222
 223 /* object member list */
 224
 225 static struct PyMemberDef binaryObject_members[] = {
 226     {"adapted", T_OBJECT, offsetof(binaryObject, wrapped), RO},
 227     {"buffer", T_OBJECT, offsetof(binaryObject, buffer), RO},
 228     {NULL}
 229 };
 230
 231 /* object method table */
 232
 233 static PyMethodDef binaryObject_methods[] = {
 234     {"getquoted", (PyCFunction)binary_getquoted, METH_VARARGS,
 235      "getquoted() -> wrapped object value as SQL-quoted binary string"},
 236     {"prepare", (PyCFunction)binary_prepare, METH_VARARGS,
 237      "prepare(conn) -> prepare for binary encoding using conn"},
 238     {"__conform__", (PyCFunction)binary_conform, METH_VARARGS, NULL},
 239     {NULL}  /* Sentinel */
 240 };
 241
 242 /* initialization and finalization methods */
 243
 244 static int
 245 binary_setup(binaryObject *self, PyObject *str)
 246 {
 247     Dprintf("binary_setup: init binary object at %p, refcnt = %d",
 248             self, ((PyObject *)self)->ob_refcnt);
 249
 250     self->buffer = NULL;
 251     self->conn = NULL;
 252     self->wrapped = str;
 253     Py_INCREF(self->wrapped);
 254
 255     Dprintf("binary_setup: good binary object at %p, refcnt = %d",
 256             self, ((PyObject *)self)->ob_refcnt);
 257     return 0;
 258 }
 259
 260 static void
 261 binary_dealloc(PyObject* obj)
 262 {
 263     binaryObject *self = (binaryObject *)obj;
 264
 265     Py_XDECREF(self->wrapped);
 266     Py_XDECREF(self->buffer);
 267     Py_XDECREF(self->conn);
 268
 269     Dprintf("binary_dealloc: deleted binary object at %p, refcnt = %d",
 270             obj, obj->ob_refcnt);
 271
 272     obj->ob_type->tp_free(obj);
 273 }
 274
 275 static int
 276 binary_init(PyObject *obj, PyObject *args, PyObject *kwds)
 277 {
 278     PyObject *str;
 279
 280     if (!PyArg_ParseTuple(args, "O", &str))
 281         return -1;
 282
 283     return binary_setup((binaryObject *)obj, str);
 284 }
 285
 286 static PyObject *
 287 binary_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 288 {
 289     return type->tp_alloc(type, 0);
 290 }
 291
 292 static void
 293 binary_del(PyObject* self)
 294 {
 295     PyObject_Del(self);
 296 }
 297
 298 static PyObject *
 299 binary_repr(binaryObject *self)
 300 {
 301     return PyString_FromFormat("<psycopg2._psycopg.Binary object at %p>", self);
 302 }
 303
 304 /* object type */
 305
 306 #define binaryType_doc \
 307 "Binary(buffer) -> new binary object"
 308
 309 PyTypeObject binaryType = {
 310     PyObject_HEAD_INIT(NULL)
 311     0,
 312     "psycopg2._psycopg.Binary",
 313     sizeof(binaryObject),
 314     0,
 315     binary_dealloc, /*tp_dealloc*/
 316     0,          /*tp_print*/
 317     0,          /*tp_getattr*/
 318     0,          /*tp_setattr*/
 319
 320     0,          /*tp_compare*/
 321     (reprfunc)binary_repr, /*tp_repr*/
 322     0,          /*tp_as_number*/
 323     0,          /*tp_as_sequence*/
 324     0,          /*tp_as_mapping*/
 325     0,          /*tp_hash */
 326
 327     0,          /*tp_call*/
 328     (reprfunc)binary_str, /*tp_str*/
 329     0,          /*tp_getattro*/
 330     0,          /*tp_setattro*/
 331     0,          /*tp_as_buffer*/
 332
 333     Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE, /*tp_flags*/
 334
 335     binaryType_doc, /*tp_doc*/
 336
 337     0,          /*tp_traverse*/
 338     0,          /*tp_clear*/
 339
 340     0,          /*tp_richcompare*/
 341     0,          /*tp_weaklistoffset*/
 342
 343     0,          /*tp_iter*/
 344     0,          /*tp_iternext*/
 345
 346     /* Attribute descriptor and subclassing stuff */
 347
 348     binaryObject_methods, /*tp_methods*/
 349     binaryObject_members, /*tp_members*/
 350     0,          /*tp_getset*/
 351     0,          /*tp_base*/
 352     0,          /*tp_dict*/
 353
 354     0,          /*tp_descr_get*/
 355     0,          /*tp_descr_set*/
 356     0,          /*tp_dictoffset*/
 357
 358     binary_init, /*tp_init*/
 359     0, /*tp_alloc  will be set to PyType_GenericAlloc in module init*/
 360     binary_new, /*tp_new*/
 361     (freefunc)binary_del, /*tp_free  Low-level free-memory routine */
 362     0,          /*tp_is_gc For PyObject_IS_GC */
 363     0,          /*tp_bases*/
 364     0,          /*tp_mro method resolution order */
 365     0,          /*tp_cache*/
 366     0,          /*tp_subclasses*/
 367     0           /*tp_weaklist*/
 368 };
 369
 370
 371 /** module-level functions **/
 372
 373 PyObject *
 374 psyco_Binary(PyObject *module, PyObject *args)
 375 {
 376     PyObject *str;
 377
 378     if (!PyArg_ParseTuple(args, "O", &str))
 379         return NULL;
 380
 381     return PyObject_CallFunction((PyObject *)&binaryType, "O", str);
 382 }