Add optional columns list to inserttable (#24)

docs/contents/changelog.rst

@@ -14,16 +14,18 @@ Version 5.2 (to be released)
       Additional connection parameter ``nowait``, and connection methods
       `send_query()`, `poll()`, `set_non_blocking()`, `is_non_blocking()`.
       Generously contributed by Patrick TJ McPhee (#19).
+    - The `types` parameter of `format_query` can now be passed as a string
+      that will be split on whitespace when values are passed as a sequence,
+      and the types can now also be specified using actual Python types
+      instead of type names. Suggested by Justin Pryzby (#38).
+    - The `inserttable()` method now accepts an optional column list that will
+      be passed on to the COPY command. Contributed by Justin Pryzby (#24).
 
 - Changes to the DB-API 2 module (pgdb):
     - When using Python 2, errors are now derived from StandardError
       instead of Exception, as required by the DB-API 2 compliance test.
     - Connection arguments containing single quotes caused problems
       (reported and fixed by Tyler Ramer and Jamie McAtamney).
-    - The `types` parameter of `format_query` can now be passed as a string
-      that will be split on whitespace when values are passed as a sequence,
-      and the types can now also be specified using actual Python types
-      instead of type names. Suggested by Justin Pryzby (#38).
 
 Version 5.1.2 (2020-04-19)
 --------------------------

docs/contents/pg/connection.rst

@@ -487,12 +487,13 @@ first, otherwise :meth:`Connection.getnotify` will always return ``None``.
 inserttable -- insert a list into a table
 -----------------------------------------
 
-.. method:: Connection.inserttable(table, values)
+.. method:: Connection.inserttable(table, values, [columns])
 
     Insert a Python list into a database table
 
     :param str table: the table name
     :param list values: list of rows values
+    :param list columns: list of column names
     :rtype: None
     :raises TypeError: invalid connection, bad argument type, or too many arguments
     :raises MemoryError: insert buffer could not be allocated
@@ -503,6 +504,8 @@ It inserts the whole values list into the given table. Internally, it
 uses the COPY command of the PostgreSQL database. The list is a list
 of tuples/lists that define the values for each inserted row. The rows
 values may contain string, integer, long or double (real) values.
+``columns`` is a optional sequence of column names to be passed on
+to the COPY command.
 
 .. warning::
 

pgconn.c

@@ -686,8 +686,9 @@ conn_is_non_blocking(connObject *self, PyObject *args)
 
 /* Insert table */
 static char conn_inserttable__doc__[] =
-"inserttable(table, data) -- insert list into table\n\n"
-"The fields in the list must be in the same order as in the table.\n";
+"inserttable(table, data, [columns]) -- insert list into table\n\n"
+"The fields in the list must be in the same order as in the table\n"
+"or in the list of columns if one is specified.\n";
 
 static PyObject *
 conn_inserttable(connObject *self, PyObject *args)
@@ -696,18 +697,19 @@ conn_inserttable(connObject *self, PyObject *args)
     char *table, *buffer, *bufpt;
     int encoding;
     size_t bufsiz;
-    PyObject *list, *sublist, *item;
+    PyObject *list, *sublist, *item, *columns = NULL;
     PyObject *(*getitem) (PyObject *, Py_ssize_t);
     PyObject *(*getsubitem) (PyObject *, Py_ssize_t);
-    Py_ssize_t i, j, m, n;
+    PyObject *(*getcolumn) (PyObject *, Py_ssize_t);
+    Py_ssize_t i, j, m, n = 0;
 
     if (!self->cnx) {
         PyErr_SetString(PyExc_TypeError, "Connection is not valid");
         return NULL;
     }
 
     /* gets arguments */
-    if (!PyArg_ParseTuple(args, "sO:filter", &table, &list)) {
+    if (!PyArg_ParseTuple(args, "sO|O", &table, &list, &columns)) {
         PyErr_SetString(
             PyExc_TypeError,
             "Method inserttable() expects a string and a list as arguments");
@@ -731,12 +733,68 @@ conn_inserttable(connObject *self, PyObject *args)
         return NULL;
     }
 
+    /* checks columns type */
+    if (columns) {
+        if (PyList_Check(columns)) {
+            n = PyList_Size(columns);
+            getcolumn = PyList_GetItem;
+        }
+        else if (PyTuple_Check(columns)) {
+            n = PyTuple_Size(columns);
+            getcolumn = PyTuple_GetItem;
+        }
+        else {
+            PyErr_SetString(
+                PyExc_TypeError,
+                "Method inserttable() expects a list or a tuple"
+                " as third argument");
+            return NULL;
+        }
+        if (!n) {
+            /* no columns specified, nothing to do */
+            Py_INCREF(Py_None);
+            return Py_None;
+        }
+    }
+
     /* allocate buffer */
     if (!(buffer = PyMem_Malloc(MAX_BUFFER_SIZE)))
         return PyErr_NoMemory();
 
+    encoding = PQclientEncoding(self->cnx);
+
     /* starts query */
-    sprintf(buffer, "copy %s from stdin", table);
+    bufpt = buffer;
+    table = PQescapeIdentifier(self->cnx, table, strlen(table));
+    bufpt += sprintf(bufpt, "copy %s", table);
+    PQfreemem(table);
+    if (columns) {
+        /* adds a string like f" ({','.join(columns)})" */
+        bufpt += sprintf(bufpt, " (");
+        for (int i = 0; i < n; ++i) {
+            PyObject *obj = getcolumn(columns, i);
+            ssize_t slen;
+            char *col;
+
+            if (PyBytes_Check(obj)) {
+                PyBytes_AsStringAndSize(obj, &col, &slen);
+            }
+            else if (PyUnicode_Check(obj)) {
+                obj = get_encoded_string(obj, encoding);
+                if (!obj) return NULL; /* pass the UnicodeEncodeError */
+                PyBytes_AsStringAndSize(obj, &col, &slen);
+                Py_DECREF(obj);
+            } else {
+                PyErr_SetString(
+                    PyExc_TypeError,
+                    "The third argument must contain only strings");
+            }
+            col = PQescapeIdentifier(self->cnx, col, (size_t) slen);
+            bufpt += sprintf(bufpt, "%s%s", col, i == n - 1 ? ")" : ",");
+            PQfreemem(col);
+        }
+    }
+    sprintf(bufpt, " from stdin");
 
     Py_BEGIN_ALLOW_THREADS
     result = PQexec(self->cnx, buffer);
@@ -748,12 +806,8 @@ conn_inserttable(connObject *self, PyObject *args)
         return NULL;
     }
 
-    encoding = PQclientEncoding(self->cnx);
-
     PQclear(result);
 
-    n = 0; /* not strictly necessary but avoids warning */
-
     /* feed table */
     for (i = 0; i < m; ++i) {
         sublist = getitem(list, i);

tests/test_classic_connection.py

@@ -1889,6 +1889,42 @@ def testInserttableNullValues(self):
         self.c.inserttable('test', data)
         self.assertEqual(self.get_back(), data)
 
+    def testInserttableNoColumn(self):
+        data = [()] * 10
+        self.c.inserttable('test', data, [])
+        self.assertEqual(self.get_back(), [])
+
+    def testInserttableOnlyOneColumn(self):
+        data = [(42,)] * 50
+        self.c.inserttable('test', data, ['i4'])
+        data = [tuple([42 if i == 1 else None for i in range(14)])] * 50
+        self.assertEqual(self.get_back(), data)
+
+    def testInserttableOnlyTwoColumns(self):
+        data = [(bool(i % 2), i * .5) for i in range(20)]
+        self.c.inserttable('test', data, ('b', 'f4'))
+        # noinspection PyTypeChecker
+        data = [(None,) * 3 + (bool(i % 2),) + (None,) * 3 + (i * .5,)
+                + (None,) * 6 for i in range(20)]
+        self.assertEqual(self.get_back(), data)
+
+    def testInserttableWithInvalidTableName(self):
+        data = [(42,)]
+        # check that the table name is not inserted unescaped
+        # (this would pass otherwise since there is a column named i4)
+        self.assertRaises(Exception, self.c.inserttable, 'test (i4)', data)
+        # make sure that it works if parameters are passed properly
+        self.c.inserttable('test', data, ['i4'])
+
+    def testInserttableWithInvalidColumnName(self):
+        data = [(2, 4)]
+        # check that the column names are not inserted unescaped
+        # (this would pass otherwise since there are columns i2 and i4)
+        self.assertRaises(
+            Exception, self.c.inserttable, 'test', data, ['i2,i4'])
+        # make sure that it works if parameters are passed properly
+        self.c.inserttable('test', data, ['i2', 'i4'])
+
     def testInserttableMaxValues(self):
         data = [(2 ** 15 - 1, int(2 ** 31 - 1), long(2 ** 31 - 1),
                  True, '2999-12-31', '11:59:59', 1e99,