Allocate dynamically memory for the list of columns in COPY

Some bloke finds the limit of 8K too restrictive... ticket #68.

NEWS

@@ -10,6 +10,7 @@ What's new in psycopg 2.4.3
   - Fixed --static-libpq setup option (ticket #64).
   - Fixed interaction between RealDictCursor and named cursors
     (ticket #67).
+  - Dropped limit on the columns length in COPY operations (ticket #68).
   - 'errorcodes' map updated to PostgreSQL 9.1.
 
 

psycopg/cursor_type.c

@@ -1161,32 +1161,53 @@ psyco_curs_scroll(cursorObject *self, PyObject *args, PyObject *kwargs)
 
 #ifdef PSYCOPG_EXTENSIONS
 
-static int _psyco_curs_copy_columns(PyObject *columns, char *columnlist)
+/* Return a newly allocated buffer containing the list of columns to be
+ * copied. On error return NULL and set an exception.
+ */
+static char *_psyco_curs_copy_columns(PyObject *columns)
 {
     PyObject *col, *coliter;
     Py_ssize_t collen;
     char *colname;
+    char *columnlist = NULL;
+    Py_ssize_t bufsize = 512;
     Py_ssize_t offset = 1;
 
+    if (columns == NULL || columns == Py_None) {
+        if (NULL == (columnlist = PyMem_Malloc(2))) {
+            PyErr_NoMemory();
+            goto error;
+        }
         columnlist[0] = '\0';
-    if (columns == NULL || columns == Py_None) return 0;
+        goto exit;
+    }
 
-    coliter = PyObject_GetIter(columns);
+    if (NULL == (coliter = PyObject_GetIter(columns))) {
-    if (coliter == NULL) return 0;
+        goto error;
+    }
 
+    if (NULL == (columnlist = PyMem_Malloc(bufsize))) {
+        PyErr_NoMemory();
+        goto error;
+    }
     columnlist[0] = '(';
 
     while ((col = PyIter_Next(coliter)) != NULL) {
         if (!(col = psycopg_ensure_bytes(col))) {
             Py_DECREF(coliter);
-            return -1;
+            goto error;
         }
         Bytes_AsStringAndSize(col, &colname, &collen);
-        if (offset + collen > DEFAULT_COPYBUFF - 2) {
+        while (offset + collen > bufsize - 2) {
+            char *tmp;
+            bufsize *= 2;
+            if (NULL == (tmp = PyMem_Realloc(columnlist, bufsize))) {
                 Py_DECREF(col);
                 Py_DECREF(coliter);
-            PyErr_SetString(PyExc_ValueError, "column list too long");
+                PyErr_NoMemory();
-            return -1;
+                goto error;
+            }
+            columnlist = tmp;
         }
         strncpy(&columnlist[offset], colname, collen);
         offset += collen;
@@ -1197,17 +1218,24 @@ static int _psyco_curs_copy_columns(PyObject *columns, char *columnlist)
 
     /* Error raised by the coliter generator */
     if (PyErr_Occurred()) {
-        return -1;
+        goto error;
     }
 
     if (offset == 2) {
-        return 0;
+        goto exit;
     }
     else {
         columnlist[offset - 1] = ')';
         columnlist[offset] = '\0';
-        return 1;
+        goto exit;
     }
+
+error:
+    PyMem_Free(columnlist);
+    columnlist = NULL;
+
+exit:
+    return columnlist;
 }
 
 /* extension: copy_from - implements COPY FROM */
@@ -1246,7 +1274,7 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
     const char *sep = "\t", *null = NULL;
     Py_ssize_t bufsize = DEFAULT_COPYBUFF;
     PyObject *file, *columns = NULL, *res = NULL;
-    char columnlist[DEFAULT_COPYBUFF];
+    char *columnlist = NULL;
     char *quoted_delimiter = NULL;
     char *quoted_null = NULL;
 
@@ -1261,14 +1289,14 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
         return NULL;
     }
 
-    if (_psyco_curs_copy_columns(columns, columnlist) == -1)
-        return NULL;
-
     EXC_IF_CURS_CLOSED(self);
     EXC_IF_CURS_ASYNC(self, copy_from);
     EXC_IF_GREEN(copy_from);
     EXC_IF_TPC_PREPARED(self->conn, copy_from);
 
+    if (NULL == (columnlist = _psyco_curs_copy_columns(columns)))
+        goto exit;
+
     if (!(quoted_delimiter = psycopg_escape_string(
             (PyObject*)self->conn, sep, 0, NULL, NULL))) {
         PyErr_NoMemory();
@@ -1327,6 +1355,7 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
     Py_DECREF(file);
 
 exit:
+    PyMem_Free(columnlist);
     PyMem_Free(quoted_delimiter);
     PyMem_Free(quoted_null);
     if (query != query_buffer) { PyMem_Free(query); }
@@ -1359,7 +1388,7 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
     char *query = NULL;
     char query_buffer[DEFAULT_COPYBUFF];
     size_t query_size;
-    char columnlist[DEFAULT_COPYBUFF];
+    char *columnlist = NULL;
     const char *table_name;
     const char *sep = "\t", *null = NULL;
     PyObject *file, *columns = NULL, *res = NULL;
@@ -1374,14 +1403,14 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
         return NULL;
     }
 
-    if (_psyco_curs_copy_columns(columns, columnlist) == -1)
-        return NULL;
-
     EXC_IF_CURS_CLOSED(self);
     EXC_IF_CURS_ASYNC(self, copy_to);
     EXC_IF_GREEN(copy_to);
     EXC_IF_TPC_PREPARED(self->conn, copy_to);
 
+    if (NULL == (columnlist = _psyco_curs_copy_columns(columns)))
+        goto exit;
+
     if (!(quoted_delimiter = psycopg_escape_string(
             (PyObject*)self->conn, sep, 0, NULL, NULL))) {
         PyErr_NoMemory();
@@ -1440,6 +1469,7 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
     self->copyfile = NULL;
 
 exit:
+    PyMem_Free(columnlist);
     PyMem_Free(quoted_delimiter);
     PyMem_Free(quoted_null);
     if (query != query_buffer) { PyMem_Free(query); }

tests/test_copy.py

@@ -253,6 +253,25 @@ class CopyTests(unittest.TestCase):
         self.assertRaises(TypeError,
             curs.copy_expert, 'COPY tcopy (data) FROM STDIN', f)
 
+    def test_copy_no_column_limit(self):
+        cols = [ "c%050d" % i for i in range(200) ]
+
+        curs = self.conn.cursor()
+        curs.execute('CREATE TEMPORARY TABLE manycols (%s)' % ',\n'.join(
+            [ "%s int" % c for c in cols]))
+        curs.execute("INSERT INTO manycols DEFAULT VALUES")
+
+        f = StringIO()
+        curs.copy_to(f, "manycols", columns = cols)
+        f.seek(0)
+        self.assertEqual(f.read().split(), ['\\N'] * len(cols))
+
+        f.seek(0)
+        curs.copy_from(f, "manycols", columns = cols)
+        curs.execute("select count(*) from manycols;")
+        self.assertEqual(curs.fetchone()[0], 2)
+
+
 decorate_all_tests(CopyTests, skip_if_green)