This commit is contained in:
Milosz 2015-11-30 06:02:46 +00:00
commit 92c5356b3f
4 changed files with 202 additions and 28 deletions

View File

@ -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
input sequence. Input parameters are left untouched, output and
input/output parameters replaced with possibly new values.
expects. Overloaded procedures are supported. Named parameters can be
used with a PostgreSQL 9.0+ client by supplying the sequence of
parameters as a Dict.
The procedure may also provide a result set as output. This must then
be made available through the standard |fetch*|_ methods.
This function is, at present, not DBAPI-compliant. The return value is
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])

View File

@ -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",
&procname, &procname_len, &parameters
))
{ goto exit; }
int using_dict;
PyObject *pname = NULL;
PyObject *pnames = NULL;
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 */
sl = procname_len + 17 + nparameters*3 - (nparameters ? 1 : 0);
using_dict = nparameters > 0 && PyDict_Check(parameters);
/* 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,
self->conn->async, 0)) {
Py_INCREF(parameters);
res = parameters;
if (0 <= _psyco_curs_execute(
self, operation, pvals, self->conn->async, 0)) {
/* 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;
}

View File

@ -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__)

View File

@ -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