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

Added support for CSV and TXT formats

Browse Source
This commit is contained in:
Sean Harrington 2016-07-29 12:21:36 -04:00
parent a1c7f4bf5d
commit e61e3d0e7b
1 changed files with 38 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

52
psycopg/cursor_type.c
View File

@ -1339,14 +1339,13 @@ static PyObject *
psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs) psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
{ {
static char *kwlist[] = { static char *kwlist[] = {
"file", "table", "sep", "null", "size", "columns", NULL}; "file", "table", "sep", "null", "size", "columns", "quote", "format", NULL};
const char *sep = "\t"; const char *sep = "\t";
const char *null = "\\N"; const char *null = "\\N";
const char *quote = "\""; const char *quote = "\"";
const char *format = "TXT";
const char *command =
"COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s QUOTE AS %s CSV";
Py_ssize_t query_size; Py_ssize_t query_size;
char *query = NULL; char *query = NULL;
@ -1356,16 +1355,19 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
char *quoted_quote = NULL; char *quoted_quote = NULL;
const char *table_name; const char *table_name;
const char *command;
Py_ssize_t bufsize = DEFAULT_COPYBUFF; Py_ssize_t bufsize = DEFAULT_COPYBUFF;
PyObject *file, *columns = NULL, *res = NULL; PyObject *file, *columns = NULL, *res = NULL;
if (!PyArg_ParseTupleAndKeywords(args, kwargs, if (!PyArg_ParseTupleAndKeywords(args, kwargs,
"O&s|ssnO", kwlist, "O&s|ssnO", kwlist,
_psyco_curs_has_read_check, &file, &table_name, &sep, &null, &bufsize, _psyco_curs_has_read_check, &file, &table_name, &sep, &null, &bufsize,
&columns)) &columns, &quote, &format))
{ {
return NULL; return NULL;
} }
EXC_IF_CURS_CLOSED(self); EXC_IF_CURS_CLOSED(self);
EXC_IF_CURS_ASYNC(self, copy_from); EXC_IF_CURS_ASYNC(self, copy_from);
@ -1384,23 +1386,45 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
self->conn, null, 0, NULL, NULL))) { self->conn, null, 0, NULL, NULL))) {
goto exit; goto exit;
} }
if (!(quoted_quote = psycopg_escape_string( if (!(quoted_quote = psycopg_escape_string(
self->conn, quote, 0, NULL, NULL))) { self->conn, quote, 0, NULL, NULL))) {
goto exit; goto exit;
} }
query_size = strlen(command) + strlen(table_name) + strlen(columnlist)
+ strlen(quoted_delimiter) + strlen(quoted_null) + strlen(quoted_quote) + 1;
if (!(query = PyMem_New(char, query_size))) { if(strcmp("TXT", format) == 0){
PyErr_NoMemory(); // Load by default TXT file type
goto exit; command =
"COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s";
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);
}else{
// Load from .CSV
command =
"COPY %s%s FROM stdin WITH DELIMITER AS %s NULL AS %s QUOTE AS %s %s";
query_size = strlen(command) + strlen(table_name) + strlen(columnlist)
+ strlen(quoted_delimiter) + strlen(quoted_null) + strlen(quoted_quote)
+ strlen(format) + 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, quoted_quote, format);
} }
PyOS_snprintf(query, query_size, command,
table_name, columnlist, quoted_delimiter, quoted_null, quoted_quote);
Dprintf("psyco_curs_copy_from: query = %s", query); Dprintf("psyco_curs_copy_from: query = %s", query);
self->copysize = bufsize; self->copysize = bufsize;
Py_INCREF(file); Py_INCREF(file);
self->copyfile = file; self->copyfile = file;