Unified string quoting

ChangeLog

@@ -1,3 +1,14 @@
+2008-11-25  Federico Di Gregorio  <fog@initd.org>
+
+	* psycopg/cursor_type.c: integrated patch from Alejandro Dubrovsky.
+	Note that the statically allocated buffer should probably go away
+	in favor of always allocating the buffer dinamically.
+
+	* psycopg/utils.c: modified patch from Alejandro Dubrovsky to
+	support quoted separators in COPY queries: now all the string
+	quoting code is in utils.c and the same function is used by
+	qstrings and everything else (like the COPY code.)
+
 2008-09-24  Federico Di Gregorio <fog@initd.org>
 
 	* lib/extras.py: added inet support and related tests.

NEWS

@@ -1,3 +1,10 @@
+What's new in psycopg 2.0.9
+---------------------------
+
+* New features:
+  - COPY TO/COPY FROM queries now can be of any size and psycopg will
+    correctly quote separators.
+
 What's new in psycopg 2.0.8
 ---------------------------
 

psycopg/adapter_qstring.c

@@ -36,56 +36,6 @@
 #include "psycopg/microprotocols_proto.h"
 
 
-/** the quoting code */
-
-#ifndef PSYCOPG_OWN_QUOTING
-size_t
-qstring_escape(char *to, const char *from, size_t len, PGconn *conn)
-{
-#if PG_MAJOR_VERSION > 8 || \
- (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION > 1) || \
- (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION == 1 && PG_PATCH_VERSION >= 4)
-    int err;
-    if (conn)
-        return PQescapeStringConn(conn, to, from, len, &err);
-    else
-#endif
-        return PQescapeString(to, from, len);
-}
-#else
-size_t
-qstring_escape(char *to, const char *from, size_t len, PGconn *conn)
-{
-    int i, j;
-
-    for (i=0, j=0; i<len; i++) {
-        switch(from[i]) {
-
-        case '\'':
-            to[j++] = '\'';
-            to[j++] = '\'';
-            break;
-
-        case '\\':
-            to[j++] = '\\';
-            to[j++] = '\\';
-            break;
-
-        case '\0':
-            /* do nothing, embedded \0 are discarded */
-            break;
-
-        default:
-            to[j++] = from[i];
-        }
-    }
-    to[j] = '\0';
-
-    Dprintf("qstring_quote: to = %s", to);
-    return strlen(to);
-}
-#endif
-
 /* qstring_quote - do the quote process on plain and unicode strings */
 
 static PyObject *
@@ -93,8 +43,7 @@ qstring_quote(qstringObject *self)
 {
     PyObject *str;
     char *s, *buffer;
-    Py_ssize_t len;
+    Py_ssize_t len, qlen;
-    int equote;         /* buffer offset if E'' quotes are needed */
 
     /* if the wrapped object is an unicode object we can encode it to match
        self->encoding but if the encoding is not specified we don't know what
@@ -139,39 +88,29 @@ qstring_quote(qstringObject *self)
     /* encode the string into buffer */
     PyString_AsStringAndSize(str, &s, &len);
 
-    buffer = (char *)PyMem_Malloc((len*2+4) * sizeof(char));
+    /* Call qstring_escape with the GIL released, then reacquire the GIL
+       before verifying that the results can fit into a Python string; raise
+       an exception if not. */        
+
+    Py_BEGIN_ALLOW_THREADS
+    buffer = psycopg_escape_string(self->conn, s, len, NULL, &qlen);
+    Py_END_ALLOW_THREADS
+    
     if (buffer == NULL) {
         Py_DECREF(str);
         PyErr_NoMemory();
         return NULL;
     }
 
-    equote = (self->conn && ((connectionObject*)self->conn)->equote) ? 1 : 0;
+    if (qlen > (size_t) PY_SSIZE_T_MAX) {
-
+        PyErr_SetString(PyExc_IndexError,
-    { /* Call qstring_escape with the GIL released, then reacquire the GIL
+            "PG buffer too large to fit in Python buffer.");
-       * before verifying that the results can fit into a Python string; raise
+        PyMem_Free(buffer);
-       * an exception if not. */
+        Py_DECREF(str);
-        size_t qstring_res;
+        return NULL;
-
-        Py_BEGIN_ALLOW_THREADS
-        qstring_res = qstring_escape(buffer+equote+1, s, len,
-          self->conn ? ((connectionObject*)self->conn)->pgconn : NULL);
-        Py_END_ALLOW_THREADS
-
-        if (qstring_res > (size_t) PY_SSIZE_T_MAX) {
-            PyErr_SetString(PyExc_IndexError, "PG buffer too large to fit in"
-              " Python buffer.");
-            PyMem_Free(buffer);
-            Py_DECREF(str);
-            return NULL;
-        }
-        len = (Py_ssize_t) qstring_res;
-        if (equote)
-            buffer[0] = 'E';
-        buffer[equote] = '\'' ; buffer[len+equote+1] = '\'';
     }
-
+    
-    self->buffer = PyString_FromStringAndSize(buffer, len+equote+2);
+    self->buffer = PyString_FromStringAndSize(buffer, qlen);
     PyMem_Free(buffer);
     Py_DECREF(str);
 

psycopg/connection_type.c

@@ -302,7 +302,8 @@ psyco_conn_lobject(connectionObject *self, PyObject *args, PyObject *keywds)
         return NULL;
     }
     
-    Dprintf("psyco_conn_lobject: new lobject at %p: refcnt = %d",
+    Dprintf("psyco_conn_lobject: new lobject at %p: refcnt = "
+            FORMAT_CODE_PY_SSIZE_T,
             obj, obj->ob_refcnt);
     return obj;
 }

psycopg/cursor_type.c

@@ -34,7 +34,6 @@
 #include "psycopg/typecast.h"
 #include "psycopg/microprotocols.h"
 #include "psycopg/microprotocols_proto.h"
-#include "psycopg/utils.h"
 #include "pgversion.h"
 #include <stdlib.h>
 
@@ -1190,7 +1189,7 @@ static PyObject *
 psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
 {
     char query_buffer[COPY_BUFFER_SIZE];
-    size_t query_size;
+    Py_ssize_t query_size;
     char *query;
     const char *table_name;
     const char *sep = "\t", *null = NULL;
@@ -1199,14 +1198,13 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
     char columnlist[DEFAULT_COPYBUFF];
     char *quoted_delimiter;
 
-    static char *kwlist[] = {"file", "table", "sep", "null", "size",
+    static char *kwlist[] = {
-                             "columns", NULL};
+            "file", "table", "sep", "null", "size", "columns", NULL};
 
     if (!PyArg_ParseTupleAndKeywords(args, kwargs,
         "O&s|ss" CONV_CODE_PY_SSIZE_T "O", kwlist,
         _psyco_curs_has_read_check, &file, &table_name, &sep, &null, &bufsize,
-        &columns)
+        &columns))
-      )
     {
         return NULL;
     }
@@ -1216,44 +1214,46 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
 
     EXC_IF_CURS_CLOSED(self);
 
-    quoted_delimiter = psycopg_internal_escape_string(self->conn, sep);
+    quoted_delimiter = psycopg_escape_string((PyObject*)self->conn, sep, 0, NULL, NULL);
-    if (NULL == quoted_delimiter) {
+    if (quoted_delimiter == NULL) {
-        PyErr_SetString(PyExc_ValueError, "Failed to quote delimiter");
+        PyErr_NoMemory();
         return NULL;
     }
+    
     query = query_buffer;
     if (null) {
-        char *quoted_null = psycopg_internal_escape_string(self->conn, null);
+        char *quoted_null = psycopg_escape_string((PyObject*)self->conn, null, 0, NULL, NULL);
-        if (NULL == quoted_null) {
+        if (quoted_null == NULL) {
-            PyErr_SetString(PyExc_ValueError, "Failed to quote null-marker");
+            PyMem_Free(quoted_delimiter);
+            PyErr_NoMemory();
             return NULL;
         }
         query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE,
-                                   "COPY %s%s FROM stdin WITH DELIMITER AS %s"
+            "COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s",
-                                   " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
+            table_name, columnlist, quoted_delimiter, quoted_null);
         if (query_size >= COPY_BUFFER_SIZE) {
             /* Got truncated, allocate dynamically */
-            query = (char *) malloc((query_size + 1) * sizeof(char));
+            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
             PyOS_snprintf(query, query_size + 1,
-                          "COPY %s%s FROM stdin WITH DELIMITER AS %s"
+                "COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s",
-                          " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
+                table_name, columnlist, quoted_delimiter, quoted_null);
         }
-        free(quoted_null);
+        PyMem_Free(quoted_null);
     }
     else {
         query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE,
-                                   "COPY %s%s FROM stdin WITH DELIMITER AS %s",
+           "COPY %s%s FROM stdin WITH DELIMITER AS %s",
-                                   table_name, columnlist, quoted_delimiter);
+           table_name, columnlist, quoted_delimiter);
         if (query_size >= COPY_BUFFER_SIZE) {
             /* Got truncated, allocate dynamically */
-            query = (char *) malloc((query_size + 1) * sizeof(char));
+            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
             PyOS_snprintf(query, query_size + 1,
-                          "COPY %s%s FROM stdin WITH DELIMITER AS %s",
+                "COPY %s%s FROM stdin WITH DELIMITER AS %s",
-                          table_name, columnlist, quoted_delimiter);
+                table_name, columnlist, quoted_delimiter);
         }
-        
+    }    
-    }
+    PyMem_Free(quoted_delimiter);
-    free(quoted_delimiter);
+    
     Dprintf("psyco_curs_copy_from: query = %s", query);
 
     self->copysize = bufsize;
@@ -1265,9 +1265,9 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
     }
 
     if (query && (query != query_buffer)) {
-        free(query);
+        PyMem_Free(query);
     }
-    self->copyfile =NULL;
+    self->copyfile = NULL;
 
     return res;
 }
@@ -1317,46 +1317,47 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
         return NULL;
 
     EXC_IF_CURS_CLOSED(self);
-    quoted_delimiter = psycopg_internal_escape_string(self->conn, sep);
+    quoted_delimiter = psycopg_escape_string((PyObject*)self->conn, sep, 0, NULL, NULL);
-    if (NULL == quoted_delimiter) {
+    if (quoted_delimiter == NULL) {
-        PyErr_SetString(PyExc_ValueError, "Failed to quote delimiter");
+        PyErr_NoMemory();
         return NULL;
     }
-
+    
     query = query_buffer;
     if (null) {
-        char *quoted_null = psycopg_internal_escape_string(self->conn, null);
+        char *quoted_null = psycopg_escape_string((PyObject*)self->conn, null, 0, NULL, NULL);
         if (NULL == quoted_null) {
-            PyErr_SetString(PyExc_ValueError, "Failed to quote null-marker");
+            PyMem_Free(quoted_delimiter);
+            PyErr_NoMemory();
             return NULL;
         }
-        
         query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE,
-                                   "COPY %s%s TO stdout WITH DELIMITER AS %s"
+            "COPY %s%s TO stdout WITH DELIMITER AS %s"
-                                   " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
+            " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
-
         if (query_size >= COPY_BUFFER_SIZE) {
             /* Got truncated, allocate dynamically */
-            query = (char *) malloc((query_size + 1) * sizeof(char));
+            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
             PyOS_snprintf(query, query_size + 1,
-                          "COPY %s%s TO stdout WITH DELIMITER AS %s"
+                "COPY %s%s TO stdout WITH DELIMITER AS %s"
-                          " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
+                " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null);
         }
-        
+        PyMem_Free(quoted_null);
     }
     else {
         query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE,
-                      "COPY %s%s TO stdout WITH DELIMITER AS %s",
+            "COPY %s%s TO stdout WITH DELIMITER AS %s",
-                      table_name, columnlist, quoted_delimiter);
+            table_name, columnlist, quoted_delimiter);
         if (query_size >= COPY_BUFFER_SIZE) {
             /* Got truncated, allocate dynamically */
-            query = (char *) malloc((query_size + 1) * sizeof(char));
+            query = (char *)PyMem_Malloc((query_size + 1) * sizeof(char));
             PyOS_snprintf(query, query_size + 1,
-                      "COPY %s%s TO stdout WITH DELIMITER AS %s",
+                "COPY %s%s TO stdout WITH DELIMITER AS %s",
-                      table_name, columnlist, quoted_delimiter);
+                table_name, columnlist, quoted_delimiter);
         }
     }
-    free(quoted_delimiter);
+    PyMem_Free(quoted_delimiter);
+    
+    Dprintf("psyco_curs_copy_to: query = %s", query);
 
     self->copysize = 0;
     self->copyfile = file;
@@ -1366,7 +1367,7 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
         Py_INCREF(Py_None);
     }
     if (query && (query != query_buffer)) {
-        free(query);
+        PyMem_Free(query);
     }
     self->copyfile = NULL;
 

psycopg/psycopg.h

@@ -143,7 +143,8 @@ HIDDEN PyObject *psyco_GetDecimalType(void);
 HIDDEN void psyco_set_error(PyObject *exc, PyObject *curs,  const char *msg,
                             const char *pgerror, const char *pgcode);
 
-HIDDEN size_t qstring_escape(char *to, const char *from, size_t len, PGconn *conn);
+HIDDEN char *psycopg_escape_string(PyObject *conn,
+              const char *from, Py_ssize_t len, char *to, Py_ssize_t *tolen);
 
 /* Exceptions docstrings */
 #define Error_doc \

psycopg/utils.c

@@ -3,7 +3,6 @@
  */
 
 #include "psycopg/config.h"
-#include "psycopg/utils.h"
 #include "psycopg/psycopg.h"
 #include "psycopg/connection.h"
 #include "psycopg/pgtypes.h"
@@ -11,33 +10,75 @@
 #include <string.h>
 #include <stdlib.h>
 
-char *psycopg_internal_escape_string(connectionObject *conn, const char *string)
+char *
+psycopg_escape_string(PyObject *obj, const char *from, Py_ssize_t len,
+                       char *to, Py_ssize_t *tolen)
 {
-    char *buffer;
+    Py_ssize_t ql;
-    size_t string_length;
+    connectionObject *conn = (connectionObject*)obj;
-    int equote;         /* buffer offset if E'' quotes are needed */  
+    int eq = (conn && (conn->equote)) ? 1 : 0;   
 
-    string_length = strlen(string);
+    if (len == 0)
+        len = strlen(from);
     
-    buffer = (char *) malloc((string_length * 2 + 4) * sizeof(char));
+    if (to == NULL) {
-    if (buffer == NULL) {
+        to = (char *)PyMem_Malloc((len * 2 + 4) * sizeof(char));
-        return NULL;
+        if (to == NULL)
+            return NULL;
     }
 
-    equote = (conn && (conn->equote)) ? 1 : 0;
+    #ifndef PSYCOPG_OWN_QUOTING
-
+    {
-    { 
+        #if PG_MAJOR_VERSION > 8 || \
-        size_t qstring_length;
+         (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION > 1) || \
-
+         (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION == 1 && PG_PATCH_VERSION >= 4)
-        qstring_length = qstring_escape(buffer + equote + 1, string, string_length,
+            int err;
-                                        (conn ? conn->pgconn : NULL));
+            if (conn && conn->pgconn)
-
+                ql = PQescapeStringConn(conn->pgconn, to+eq+1, from, len, &err);
-        if (equote)
+            else
-            buffer[0] = 'E';
+        #endif
-        buffer[equote] = '\''; 
+                ql = PQescapeString(to+eq+1, from, len);
-        buffer[qstring_length + equote + 1] = '\'';
-        buffer[qstring_length + equote + 2] = 0;
     }
+    #else
+    {
+        int i, j;
+    
+        for (i=0, j=eq+1; i<len; i++) {
+            switch(from[i]) {
+    
+            case '\'':
+                to[j++] = '\'';
+                to[j++] = '\'';
+                break;
+    
+            case '\\':
+                to[j++] = '\\';
+                to[j++] = '\\';
+                break;
+    
+            case '\0':
+                /* do nothing, embedded \0 are discarded */
+                break;
+    
+            default:
+                to[j++] = from[i];
+            }
+        }
+        to[j] = '\0';
+    
+        Dprintf("qstring_quote: to = %s", to);
+        ql = strlen(to);
+    }
+    #endif
 
-    return buffer;
+    if (eq)
+        to[0] = 'E';
+    to[eq] = '\''; 
+    to[ql+eq+1] = '\'';
+    to[ql+eq+2] = '\0';
+
+    if (tolen)
+        *tolen = ql+eq+2;
+        
+    return to;
 }

psycopg/utils.h

@@ -1,16 +0,0 @@
-/* utils.h - miscellaneous utility functions
- *
- */
-
-#ifndef PSYCOPG_UTILS_H
-#define PSYCOPG_UTILS_H 1
-
-#include "psycopg/config.h"
-#include "psycopg/connection.h"
-
-
-HIDDEN char *psycopg_internal_escape_string(connectionObject *conn, const char *string);
-
-#endif /* !defined(PSYCOPG_UTILS_H) */
-
-

psycopg2.mdp

@@ -133,7 +133,6 @@
     <File name="tests/bugX000.py" subtype="Code" buildaction="Nothing" />
     <File name="tests/types_extras.py" subtype="Code" buildaction="Nothing" />
     <File name="psycopg/utils.c" subtype="Code" buildaction="Compile" />
-    <File name="psycopg/utils.h" subtype="Code" buildaction="Nothing" />
     <File name="tests/test_connection.py" subtype="Code" buildaction="Nothing" />
     <File name="tests/test_dates.py" subtype="Code" buildaction="Nothing" />
     <File name="tests/test_lobject.py" subtype="Code" buildaction="Nothing" />