psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2025-01-31 09:24:07 +03:00
Code Issues Packages Projects Releases Wiki Activity

Avoid PyOS_snprintf to calculate the copy command buffer size

Browse Source 
On windows it returns -1 instead of sometihing portable. So just ditch
the static buffer and just use a dynamic one to compose the command.

Also squashed a couple of buglets in copy_to: copyfile was decremented
before being set to null, size_t was used instead of Py_ssize_t.
This commit is contained in:
Daniele Varrazzo 2011-10-19 21:01:53 +01:00
parent ff8158d7c0
commit 60b49f5c45
1 changed files with 42 additions and 49 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

77
psycopg/cursor_type.c
View File

@ -1267,20 +1267,23 @@ _psyco_curs_has_read_check(PyObject* o, void* var)
static PyObject * static PyObject *
psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs) psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
{ {
char *query = NULL; static char *kwlist[] = {
char query_buffer[DEFAULT_COPYBUFF]; "file", "table", "sep", "null", "size", "columns", NULL};
Py_ssize_t query_size;
const char *table_name;
const char *sep = "\t"; const char *sep = "\t";
const char *null = "\\N"; const char *null = "\\N";
Py_ssize_t bufsize = DEFAULT_COPYBUFF; const char *command =
PyObject *file, *columns = NULL, *res = NULL; "COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s";
Py_ssize_t query_size;
char *query = NULL;
char *columnlist = NULL; char *columnlist = NULL;
char *quoted_delimiter = NULL; char *quoted_delimiter = NULL;
char *quoted_null = NULL; char *quoted_null = NULL;
static char *kwlist[] = { const char *table_name;
"file", "table", "sep", "null", "size", "columns", NULL}; Py_ssize_t bufsize = DEFAULT_COPYBUFF;
PyObject *file, *columns = NULL, *res = NULL;
if (!PyArg_ParseTupleAndKeywords(args, kwargs, if (!PyArg_ParseTupleAndKeywords(args, kwargs,
"O&s|ss" CONV_CODE_PY_SSIZE_T "O", kwlist, "O&s|ss" CONV_CODE_PY_SSIZE_T "O", kwlist,
@ -1310,20 +1313,15 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
goto exit; goto exit;
} }
query = query_buffer; query_size = strlen(command) + strlen(table_name) + strlen(columnlist)
query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF, + strlen(quoted_delimiter) + strlen(quoted_null) + 1;
"COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s",
table_name, columnlist, quoted_delimiter, quoted_null);
if (query_size >= DEFAULT_COPYBUFF) {
/* Got truncated, allocate dynamically */
if (!(query = PyMem_New(char, query_size + 1))) { if (!(query = PyMem_New(char, query_size + 1))) {
PyErr_NoMemory(); PyErr_NoMemory();
goto exit; goto exit;
} }
PyOS_snprintf(query, query_size + 1,
"COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s", PyOS_snprintf(query, query_size, command,
table_name, columnlist, quoted_delimiter, quoted_null); table_name, columnlist, quoted_delimiter, quoted_null);
}
Dprintf("psyco_curs_copy_from: query = %s", query); Dprintf("psyco_curs_copy_from: query = %s", query);
@ -1336,14 +1334,13 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
Py_INCREF(Py_None); Py_INCREF(Py_None);
} }
self->copyfile = NULL; Py_CLEAR(self->copyfile);
Py_DECREF(file);
exit: exit:
PyMem_Free(columnlist); PyMem_Free(columnlist);
PyMem_Free(quoted_delimiter); PyMem_Free(quoted_delimiter);
PyMem_Free(quoted_null); PyMem_Free(quoted_null);
if (query != query_buffer) { PyMem_Free(query); } PyMem_Free(query);
return res; return res;
} }
@ -1370,18 +1367,21 @@ _psyco_curs_has_write_check(PyObject* o, void* var)
static PyObject * static PyObject *
psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs) psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
{ {
char *query = NULL; static char *kwlist[] = {"file", "table", "sep", "null", "columns", NULL};
char query_buffer[DEFAULT_COPYBUFF];
size_t query_size;
char *columnlist = NULL;
const char *table_name;
const char *sep = "\t"; const char *sep = "\t";
const char *null = "\\N"; const char *null = "\\N";
PyObject *file, *columns = NULL, *res = NULL; const char *command =
"COPY %s%s TO stdout WITH DELIMITER AS %s NULL AS %s";
Py_ssize_t query_size;
char *query = NULL;
char *columnlist = NULL;
char *quoted_delimiter = NULL; char *quoted_delimiter = NULL;
char *quoted_null = NULL; char *quoted_null = NULL;
static char *kwlist[] = {"file", "table", "sep", "null", "columns", NULL}; const char *table_name;
PyObject *file, *columns = NULL, *res = NULL;
if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O&s|ssO", kwlist, if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O&s|ssO", kwlist,
_psyco_curs_has_write_check, &file, _psyco_curs_has_write_check, &file,
@ -1409,20 +1409,15 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
goto exit; goto exit;
} }
query = query_buffer; query_size = strlen(command) + strlen(table_name) + strlen(columnlist)
query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF, + strlen(quoted_delimiter) + strlen(quoted_null) + 1;
"COPY %s%s TO stdout WITH DELIMITER AS %s NULL AS %s", if (!(query = PyMem_New(char, query_size))) {
table_name, columnlist, quoted_delimiter, quoted_null);
if (query_size >= DEFAULT_COPYBUFF) {
/* Got truncated, allocate dynamically */
if (!(query = PyMem_New(char, query_size + 1))) {
PyErr_NoMemory(); PyErr_NoMemory();
goto exit; goto exit;
} }
PyOS_snprintf(query, query_size + 1,
"COPY %s%s TO stdout WITH DELIMITER AS %s NULL AS %s", PyOS_snprintf(query, query_size, command,
table_name, columnlist, quoted_delimiter, quoted_null); table_name, columnlist, quoted_delimiter, quoted_null);
}
Dprintf("psyco_curs_copy_to: query = %s", query); Dprintf("psyco_curs_copy_to: query = %s", query);
@ -1435,14 +1430,13 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
Py_INCREF(Py_None); Py_INCREF(Py_None);
} }
Py_DECREF(file); Py_CLEAR(self->copyfile);
self->copyfile = NULL;
exit: exit:
PyMem_Free(columnlist); PyMem_Free(columnlist);
PyMem_Free(quoted_delimiter); PyMem_Free(quoted_delimiter);
PyMem_Free(quoted_null); PyMem_Free(quoted_null);
if (query != query_buffer) { PyMem_Free(query); } PyMem_Free(query);
return res; return res;
} }
@ -1510,8 +1504,7 @@ psyco_curs_copy_expert(cursorObject *self, PyObject *args, PyObject *kwargs)
Py_INCREF(res); Py_INCREF(res);
} }
self->copyfile = NULL; Py_CLEAR(self->copyfile);
Py_DECREF(file);
exit: exit:
Py_XDECREF(sql); Py_XDECREF(sql);