Merge bc0b1e0195 into fe4cb0d493

doc/src/cursor.rst

@@ -201,12 +201,17 @@ The ``cursor`` class
             
         Call a stored database procedure with the given name. The sequence of
         parameters must contain one entry for each argument that the procedure
-        expects. The result of the call is returned as modified copy of the
+        expects. Overloaded procedures are supported. Named parameters can be
-        input sequence. Input parameters are left untouched, output and
+        used with a PostgreSQL 9.0+ client by supplying the sequence of
-        input/output parameters replaced with possibly new values.
+        parameters as a Dict.
 
-        The procedure may also provide a result set as output. This must then
+        This function is, at present, not DBAPI-compliant. The return value is
-        be made available through the standard |fetch*|_ methods.
+        supposed to consist of the sequence of parameters with modified output
+        and input/output parameters. In future versions, the DBAPI-compliant
+        return value may be implemented, but for now the function returns None.
+
+        The procedure may provide a result set as output. This is then made
+        available through the standard |fetch*|_ methods.
 
 
     .. method:: mogrify(operation [, parameters])

psycopg/cursor_type.c

@@ -1015,6 +1015,35 @@ exit:
 #define psyco_curs_callproc_doc \
 "callproc(procname, parameters=None) -- Execute stored procedure."
 
+/* Call PQescapeIdentifier.
+ *
+ * In case of error set a Python exception.
+ *
+ * TODO: this function can become more generic and go into utils
+ */
+static char *
+_escape_identifier(PGconn *pgconn, const char *str, size_t length)
+{
+    char *rv = NULL;
+
+#if PG_VERSION_NUM >= 90000
+    rv = PQescapeIdentifier(pgconn, str, length);
+    if (!rv) {
+        char *msg;
+        msg = PQerrorMessage(pgconn);
+        if (!msg || !msg[0]) {
+            msg = "no message provided";
+        }
+        PyErr_Format(InterfaceError, "failed to escape identifier: %s", msg);
+    }
+#else
+    PyErr_Format(PyExc_NotImplementedError,
+        "named parameters require psycopg2 compiled against libpq 9.0");
+#endif
+
+    return rv;
+}
+
 static PyObject *
 psyco_curs_callproc(cursorObject *self, PyObject *args)
 {
@@ -1025,10 +1054,17 @@ psyco_curs_callproc(cursorObject *self, PyObject *args)
     PyObject *operation = NULL;
     PyObject *res = NULL;
 
-    if (!PyArg_ParseTuple(args, "s#|O",
+    int using_dict;
-          &procname, &procname_len, &parameters
+    PyObject *pname = NULL;
-       ))
+    PyObject *pnames = NULL;
-    { goto exit; }
+    PyObject *pvals = NULL;
+    char *cpname = NULL;
+    char **scpnames = NULL;
+
+    if (!PyArg_ParseTuple(args, "s#|O", &procname, &procname_len,
+                &parameters)) {
+        goto exit;
+    }
 
     EXC_IF_CURS_CLOSED(self);
     EXC_IF_ASYNC_IN_PROGRESS(self, callproc);
@@ -1044,8 +1080,64 @@ psyco_curs_callproc(cursorObject *self, PyObject *args)
         if (-1 == (nparameters = PyObject_Length(parameters))) { goto exit; }
     }
 
-    /* allocate some memory, build the SQL and create a PyString from it */
+    using_dict = nparameters > 0 && PyDict_Check(parameters);
-    sl = procname_len + 17 + nparameters*3 - (nparameters ? 1 : 0);
+
+    /* a Dict is complicated; the parameter names go into the query */
+    if (using_dict) {
+        if (!(pnames = PyDict_Keys(parameters))) { goto exit; }
+        if (!(pvals = PyDict_Values(parameters))) { goto exit; }
+
+        sl = procname_len + 17 + nparameters * 5 - (nparameters ? 1 : 0);
+
+        if (!(scpnames = PyMem_New(char *, nparameters))) {
+            PyErr_NoMemory();
+            goto exit;
+        }
+
+        memset(scpnames, 0, sizeof(char *) * nparameters);
+
+        /* each parameter has to be processed; it's a few steps. */
+        for (i = 0; i < nparameters; i++) {
+            /* all errors are RuntimeErrors as they should never occur */
+
+            if (!(pname = PyList_GetItem(pnames, i))) { goto exit; }
+            Py_INCREF(pname);   /* was borrowed */
+
+            /* this also makes a check for keys being strings */
+            if (!(pname = psycopg_ensure_bytes(pname))) { goto exit; }
+            if (!(cpname = Bytes_AsString(pname))) { goto exit; }
+
+            if (!(scpnames[i] = _escape_identifier(
+                    self->conn->pgconn, cpname, strlen(cpname)))) {
+                goto exit;
+            }
+
+            Py_CLEAR(pname);
+
+            sl += strlen(scpnames[i]);
+        }
+
+        if (!(sql = (char*)PyMem_Malloc(sl))) {
+            PyErr_NoMemory();
+            goto exit;
+        }
+
+        sprintf(sql, "SELECT * FROM %s(", procname);
+        for (i = 0; i < nparameters; i++) {
+            strcat(sql, scpnames[i]);
+            strcat(sql, ":=%s,");
+        }
+        sql[sl-2] = ')';
+        sql[sl-1] = '\0';
+    }
+
+    /* a list (or None, or empty data structure) is a little bit simpler */
+    else {
+        Py_INCREF(parameters);
+        pvals = parameters;
+
+        sl = procname_len + 17 + nparameters * 3 - (nparameters ? 1 : 0);
+
         sql = (char*)PyMem_Malloc(sl);
         if (sql == NULL) {
             PyErr_NoMemory();
@@ -1053,22 +1145,38 @@ psyco_curs_callproc(cursorObject *self, PyObject *args)
         }
 
         sprintf(sql, "SELECT * FROM %s(", procname);
-    for(i=0; i<nparameters; i++) {
+        for (i = 0; i < nparameters; i++) {
             strcat(sql, "%s,");
         }
         sql[sl-2] = ')';
         sql[sl-1] = '\0';
+    }
 
-    if (!(operation = Bytes_FromString(sql))) { goto exit; }
+    if (!(operation = Bytes_FromString(sql))) {
+        goto exit;
+    }
 
-    if (0 <= _psyco_curs_execute(self, operation, parameters,
+    if (0 <= _psyco_curs_execute(
-            self->conn->async, 0)) {
+            self, operation, pvals, self->conn->async, 0)) {
-        Py_INCREF(parameters);
+
-        res = parameters;
+        /* return None from this until it's DBAPI compliant... */
+        Py_INCREF(Py_None);
+        res = Py_None;
     }
 
 exit:
+    if (scpnames != NULL) {
+        for (i = 0; i < nparameters; i++) {
+            if (scpnames[i] != NULL) {
+                PQfreemem(scpnames[i]);
+            }
+        }
+    }
+    PyMem_Del(scpnames);
+    Py_XDECREF(pname);
+    Py_XDECREF(pnames);
     Py_XDECREF(operation);
+    Py_XDECREF(pvals);
     PyMem_Free((void*)sql);
     return res;
 }

tests/test_cursor.py

@@ -490,6 +490,44 @@ class CursorTests(ConnectingTestCase):
         cur = self.conn.cursor()
         self.assertRaises(TypeError, cur.callproc, 'lower', 42)
 
+    # It would be inappropriate to test callproc's named parameters in the
+    # DBAPI2.0 test section because they are a psycopg2 extension.
+    @skip_before_postgres(9, 0)
+    def test_callproc_dict(self):
+        # This parameter name tests for injection and quote escaping
+        paramname = '''
+            Robert'); drop table "students" --
+        '''.strip()
+        escaped_paramname = '"%s"' % paramname.replace('"', '""')
+        procname = 'pg_temp.randall'
+
+        cur = self.conn.cursor()
+
+        # Set up the temporary function
+        cur.execute('''
+            CREATE FUNCTION %s(%s INT)
+            RETURNS INT AS
+                'SELECT $1 * $1'
+            LANGUAGE SQL
+        ''' % (procname, escaped_paramname));
+
+        # Make sure callproc works right
+        cur.callproc(procname, { paramname: 2 })
+        self.assertEquals(cur.fetchone()[0], 4)
+
+        # Make sure callproc fails right
+        failing_cases = [
+            ({ paramname: 2, 'foo': 'bar' }, psycopg2.ProgrammingError),
+            ({ paramname: '2' },             psycopg2.ProgrammingError),
+            ({ paramname: 'two' },           psycopg2.ProgrammingError),
+            ({ u'bj\xc3rn': 2 },             psycopg2.ProgrammingError),
+            ({ 3: 2 },                       TypeError),
+            ({ self: 2 },                    TypeError),
+        ]
+        for parameter_sequence, exception in failing_cases:
+            self.assertRaises(exception, cur.callproc, procname, parameter_sequence)
+            self.conn.rollback()
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)

tests/test_psycopg2_dbapi20.py

@@ -36,6 +36,29 @@ class Psycopg2Tests(dbapi20.DatabaseAPI20Test):
     connect_kw_args = {'dsn': dsn}
 
     lower_func = 'lower' # For stored procedure test
+    def test_callproc(self):
+        # Until DBAPI 2.0 compliance, callproc should return None or it's just
+        # misleading. Therefore, we will skip the return value test for
+        # callproc and only perform the fetch test.
+        #
+        # For what it's worth, the DBAPI2.0 test_callproc doesn't actually
+        # test for DBAPI2.0 compliance! It doesn't check for modified OUT and
+        # IN/OUT parameters in the return values!
+        con = self._connect()
+        try:
+            cur = con.cursor()
+            if self.lower_func and hasattr(cur,'callproc'):
+                cur.callproc(self.lower_func,('FOO',))
+                r = cur.fetchall()
+                self.assertEqual(len(r),1,'callproc produced no result set')
+                self.assertEqual(len(r[0]),1,
+                    'callproc produced invalid result set'
+                    )
+                self.assertEqual(r[0][0],'foo',
+                    'callproc produced invalid results'
+                    )
+        finally:
+            con.close()
 
     def test_setoutputsize(self):
         # psycopg2's setoutputsize() is a no-op