Merge branch 'fix-copy-on-win' into devel

This commit is contained in:
Daniele Varrazzo 2011-10-20 11:12:43 +01:00
commit 5728649944
2 changed files with 59 additions and 94 deletions

View File

@ -457,7 +457,7 @@ The ``cursor`` class
The :sql:`COPY` command is a PostgreSQL extension to the SQL standard. The :sql:`COPY` command is a PostgreSQL extension to the SQL standard.
As such, its support is a Psycopg extension to the |DBAPI|. As such, its support is a Psycopg extension to the |DBAPI|.
.. method:: copy_from(file, table, sep='\\t', null='\\N', size=8192, columns=None) .. method:: copy_from(file, table, sep='\\t', null='\\\\N', size=8192, columns=None)
Read data *from* the file-like object *file* appending them to Read data *from* the file-like object *file* appending them to
the table named *table*. See :ref:`copy` for an overview. the table named *table*. See :ref:`copy` for an overview.
@ -467,6 +467,7 @@ The ``cursor`` class
:param table: name of the table to copy data into. :param table: name of the table to copy data into.
:param sep: columns separator expected in the file. Defaults to a tab. :param sep: columns separator expected in the file. Defaults to a tab.
:param null: textual representation of :sql:`NULL` in the file. :param null: textual representation of :sql:`NULL` in the file.
The default is the two character string ``\N``.
:param size: size of the buffer used to read from the file. :param size: size of the buffer used to read from the file.
:param columns: iterable with name of the columns to import. :param columns: iterable with name of the columns to import.
The length and types should match the content of the file to read. The length and types should match the content of the file to read.
@ -489,7 +490,7 @@ The ``cursor`` class
are encoded in the connection `~connection.encoding` when sent to are encoded in the connection `~connection.encoding` when sent to
the backend. the backend.
.. method:: copy_to(file, table, sep='\\t', null='\\N', columns=None) .. method:: copy_to(file, table, sep='\\t', null='\\\\N', columns=None)
Write the content of the table named *table* *to* the file-like Write the content of the table named *table* *to* the file-like
object *file*. See :ref:`copy` for an overview. object *file*. See :ref:`copy` for an overview.
@ -499,6 +500,7 @@ The ``cursor`` class
:param table: name of the table to copy data from. :param table: name of the table to copy data from.
:param sep: columns separator expected in the file. Defaults to a tab. :param sep: columns separator expected in the file. Defaults to a tab.
:param null: textual representation of :sql:`NULL` in the file. :param null: textual representation of :sql:`NULL` in the file.
The default is the two character string ``\N``.
:param columns: iterable with name of the columns to export. :param columns: iterable with name of the columns to export.
If not specified, export all the columns. If not specified, export all the columns.

View File

@ -1241,7 +1241,7 @@ exit:
/* extension: copy_from - implements COPY FROM */ /* extension: copy_from - implements COPY FROM */
#define psyco_curs_copy_from_doc \ #define psyco_curs_copy_from_doc \
"copy_from(file, table, sep='\\t', null='\\N', size=8192, columns=None) -- Copy table from file." "copy_from(file, table, sep='\\t', null='\\\\N', size=8192, columns=None) -- Copy table from file."
static int static int
_psyco_curs_has_read_check(PyObject* o, void* var) _psyco_curs_has_read_check(PyObject* o, void* var)
@ -1267,19 +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};
const char *sep = "\t";
const char *null = "\\N";
const char *command =
"COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s";
Py_ssize_t query_size; Py_ssize_t query_size;
const char *table_name; char *query = NULL;
const char *sep = "\t", *null = NULL;
Py_ssize_t bufsize = DEFAULT_COPYBUFF;
PyObject *file, *columns = NULL, *res = 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,
@ -1303,43 +1307,22 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
goto exit; goto exit;
} }
query = query_buffer; if (!(quoted_null = psycopg_escape_string(
if (null) { (PyObject*)self->conn, null, 0, NULL, NULL))) {
if (!(quoted_null = psycopg_escape_string( PyErr_NoMemory();
(PyObject*)self->conn, null, 0, NULL, NULL))) { goto exit;
PyErr_NoMemory();
goto exit;
}
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);
}
} }
else {
query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF, query_size = strlen(command) + strlen(table_name) + strlen(columnlist)
"COPY %s%s FROM stdin WITH DELIMITER AS %s", + strlen(quoted_delimiter) + strlen(quoted_null) + 1;
table_name, columnlist, quoted_delimiter); if (!(query = PyMem_New(char, query_size))) {
if (query_size >= DEFAULT_COPYBUFF) { PyErr_NoMemory();
/* Got truncated, allocate dynamically */ goto exit;
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",
table_name, columnlist, quoted_delimiter);
}
} }
PyOS_snprintf(query, query_size, command,
table_name, columnlist, quoted_delimiter, quoted_null);
Dprintf("psyco_curs_copy_from: query = %s", query); Dprintf("psyco_curs_copy_from: query = %s", query);
self->copysize = bufsize; self->copysize = bufsize;
@ -1351,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;
} }
@ -1366,7 +1348,7 @@ exit:
/* extension: copy_to - implements COPY TO */ /* extension: copy_to - implements COPY TO */
#define psyco_curs_copy_to_doc \ #define psyco_curs_copy_to_doc \
"copy_to(file, table, sep='\\t', null='\\N', columns=None) -- Copy table to file." "copy_to(file, table, sep='\\t', null='\\\\N', columns=None) -- Copy table to file."
static int static int
_psyco_curs_has_write_check(PyObject* o, void* var) _psyco_curs_has_write_check(PyObject* o, void* var)
@ -1385,17 +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)
{ {
static char *kwlist[] = {"file", "table", "sep", "null", "columns", NULL};
const char *sep = "\t";
const char *null = "\\N";
const char *command =
"COPY %s%s TO stdout WITH DELIMITER AS %s NULL AS %s";
Py_ssize_t query_size;
char *query = NULL; char *query = NULL;
char query_buffer[DEFAULT_COPYBUFF];
size_t query_size;
char *columnlist = NULL; char *columnlist = NULL;
const char *table_name;
const char *sep = "\t", *null = NULL;
PyObject *file, *columns = NULL, *res = 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,
@ -1417,43 +1403,22 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
goto exit; goto exit;
} }
query = query_buffer; if (!(quoted_null = psycopg_escape_string(
if (null) { (PyObject*)self->conn, null, 0, NULL, NULL))) {
if (!(quoted_null = psycopg_escape_string( PyErr_NoMemory();
(PyObject*)self->conn, null, 0, NULL, NULL))) { goto exit;
PyErr_NoMemory();
goto exit;
}
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);
}
} }
else {
query_size = PyOS_snprintf(query, DEFAULT_COPYBUFF, query_size = strlen(command) + strlen(table_name) + strlen(columnlist)
"COPY %s%s TO stdout WITH DELIMITER AS %s", + strlen(quoted_delimiter) + strlen(quoted_null) + 1;
table_name, columnlist, quoted_delimiter); if (!(query = PyMem_New(char, query_size))) {
if (query_size >= DEFAULT_COPYBUFF) { PyErr_NoMemory();
/* Got truncated, allocate dynamically */ goto exit;
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",
table_name, columnlist, quoted_delimiter);
}
} }
PyOS_snprintf(query, query_size, command,
table_name, columnlist, quoted_delimiter, quoted_null);
Dprintf("psyco_curs_copy_to: query = %s", query); Dprintf("psyco_curs_copy_to: query = %s", query);
self->copysize = 0; self->copysize = 0;
@ -1465,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;
} }
@ -1540,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);