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