diff options
| -rw-r--r-- | psycopg/adapter_qstring.c | 4 | ||||
| -rw-r--r-- | psycopg/cursor_type.c | 102 | ||||
| -rw-r--r-- | psycopg/psycopg.h | 2 | ||||
| -rw-r--r-- | setup.py | 2 |
4 files changed, 90 insertions, 20 deletions
diff --git a/psycopg/adapter_qstring.c b/psycopg/adapter_qstring.c index e1eb4e1..bad5d81 100644 --- a/psycopg/adapter_qstring.c +++ b/psycopg/adapter_qstring.c @@ -40,7 +40,7 @@ #ifndef PSYCOPG_OWN_QUOTING size_t -qstring_escape(char *to, char *from, size_t len, PGconn *conn) +qstring_escape(char *to, const char *from, size_t len, PGconn *conn) { #if PG_MAJOR_VERSION > 8 || \ (PG_MAJOR_VERSION == 8 && PG_MINOR_VERSION > 1) || \ @@ -54,7 +54,7 @@ qstring_escape(char *to, char *from, size_t len, PGconn *conn) } #else size_t -qstring_escape(char *to, char *from, size_t len, PGconn *conn) +qstring_escape(char *to, const char *from, size_t len, PGconn *conn) { int i, j; diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c index 8000816..2730818 100644 --- a/psycopg/cursor_type.c +++ b/psycopg/cursor_type.c @@ -34,7 +34,9 @@ #include "psycopg/typecast.h" #include "psycopg/microprotocols.h" #include "psycopg/microprotocols_proto.h" +#include "psycopg/utils.h" #include "pgversion.h" +#include <stdlib.h> extern PyObject *pyPsycopgTzFixedOffsetTimezone; @@ -1113,7 +1115,7 @@ psyco_curs_scroll(cursorObject *self, PyObject *args, PyObject *kwargs) #ifdef PSYCOPG_EXTENSIONS -#define COPY_BUFFER_SIZE 1024 +#define COPY_BUFFER_SIZE 8192 static int _psyco_curs_copy_columns(PyObject *columns, char *columnlist) { @@ -1187,12 +1189,15 @@ _psyco_curs_has_read_check(PyObject* o, void* var) static PyObject * psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs) { - char query[DEFAULT_COPYBUFF]; + char query_buffer[COPY_BUFFER_SIZE]; + size_t query_size; + char *query; const char *table_name; const char *sep = "\t", *null = NULL; Py_ssize_t bufsize = DEFAULT_COPYBUFF; PyObject *file, *columns = NULL, *res = NULL; char columnlist[DEFAULT_COPYBUFF]; + char *quoted_delimiter; static char *kwlist[] = {"file", "table", "sep", "null", "size", "columns", NULL}; @@ -1211,16 +1216,44 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs) EXC_IF_CURS_CLOSED(self); + quoted_delimiter = psycopg_internal_escape_string(self->conn, sep); + if (NULL == quoted_delimiter) { + PyErr_SetString(PyExc_ValueError, "Failed to quote delimiter"); + return NULL; + } + query = query_buffer; if (null) { - PyOS_snprintf(query, DEFAULT_COPYBUFF-1, - "COPY %s%s FROM stdin USING DELIMITERS '%s'" - " WITH NULL AS '%s'", table_name, columnlist, sep, null); + char *quoted_null = psycopg_internal_escape_string(self->conn, null); + if (NULL == quoted_null) { + PyErr_SetString(PyExc_ValueError, "Failed to quote null-marker"); + return NULL; + } + query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE, + "COPY %s%s FROM stdin WITH DELIMITER AS %s" + " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null); + if (query_size >= COPY_BUFFER_SIZE) { + /* Got truncated, allocate dynamically */ + query = (char *) malloc((query_size + 1) * sizeof(char)); + 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); + } + free(quoted_null); } else { - PyOS_snprintf(query, DEFAULT_COPYBUFF-1, - "COPY %s%s FROM stdin USING DELIMITERS '%s'", - table_name, columnlist, sep); + query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE, + "COPY %s%s FROM stdin WITH DELIMITER AS %s", + table_name, columnlist, quoted_delimiter); + if (query_size >= COPY_BUFFER_SIZE) { + /* Got truncated, allocate dynamically */ + query = (char *) malloc((query_size + 1) * sizeof(char)); + PyOS_snprintf(query, query_size + 1, + "COPY %s%s FROM stdin WITH DELIMITER AS %s", + table_name, columnlist, quoted_delimiter); + } + } + free(quoted_delimiter); Dprintf("psyco_curs_copy_from: query = %s", query); self->copysize = bufsize; @@ -1231,6 +1264,9 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs) Py_INCREF(Py_None); } + if (query && (query != query_buffer)) { + free(query); + } self->copyfile =NULL; return res; @@ -1260,11 +1296,14 @@ _psyco_curs_has_write_check(PyObject* o, void* var) static PyObject * psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs) { - char query[DEFAULT_COPYBUFF]; + char *query = NULL; + char query_buffer[COPY_BUFFER_SIZE]; + size_t query_size; char columnlist[DEFAULT_COPYBUFF]; const char *table_name; const char *sep = "\t", *null = NULL; PyObject *file, *columns = NULL, *res = NULL; + char *quoted_delimiter; static char *kwlist[] = {"file", "table", "sep", "null", "columns", NULL}; @@ -1278,17 +1317,46 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs) return NULL; EXC_IF_CURS_CLOSED(self); + quoted_delimiter = psycopg_internal_escape_string(self->conn, sep); + if (NULL == quoted_delimiter) { + PyErr_SetString(PyExc_ValueError, "Failed to quote delimiter"); + return NULL; + } + query = query_buffer; if (null) { - PyOS_snprintf(query, DEFAULT_COPYBUFF-1, - "COPY %s%s TO stdout USING DELIMITERS '%s'" - " WITH NULL AS '%s'", table_name, columnlist, sep, null); + char *quoted_null = psycopg_internal_escape_string(self->conn, null); + if (NULL == quoted_null) { + PyErr_SetString(PyExc_ValueError, "Failed to quote null-marker"); + return NULL; + } + + query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE, + "COPY %s%s TO stdout WITH DELIMITER AS %s" + " NULL AS %s", table_name, columnlist, quoted_delimiter, quoted_null); + + if (query_size >= COPY_BUFFER_SIZE) { + /* Got truncated, allocate dynamically */ + query = (char *) malloc((query_size + 1) * sizeof(char)); + 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 { - PyOS_snprintf(query, DEFAULT_COPYBUFF-1, - "COPY %s%s TO stdout USING DELIMITERS '%s'", - table_name, columnlist, sep); + query_size = PyOS_snprintf(query, COPY_BUFFER_SIZE, + "COPY %s%s TO stdout WITH DELIMITER AS %s", + table_name, columnlist, quoted_delimiter); + if (query_size >= COPY_BUFFER_SIZE) { + /* Got truncated, allocate dynamically */ + query = (char *) malloc((query_size + 1) * sizeof(char)); + PyOS_snprintf(query, query_size + 1, + "COPY %s%s TO stdout WITH DELIMITER AS %s", + table_name, columnlist, quoted_delimiter); + } } + free(quoted_delimiter); self->copysize = 0; self->copyfile = file; @@ -1297,7 +1365,9 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs) res = Py_None; Py_INCREF(Py_None); } - + if (query && (query != query_buffer)) { + free(query); + } self->copyfile = NULL; return res; diff --git a/psycopg/psycopg.h b/psycopg/psycopg.h index 5706c5a..20e33ff 100644 --- a/psycopg/psycopg.h +++ b/psycopg/psycopg.h @@ -143,7 +143,7 @@ HIDDEN PyObject *psyco_GetDecimalType(void); HIDDEN void psyco_set_error(PyObject *exc, PyObject *curs, const char *msg, const char *pgerror, const char *pgcode); -HIDDEN size_t qstring_escape(char *to, char *from, size_t len, PGconn *conn); +HIDDEN size_t qstring_escape(char *to, const char *from, size_t len, PGconn *conn); /* Exceptions docstrings */ #define Error_doc \ @@ -330,7 +330,7 @@ sources = [ 'connection_type.c', 'connection_int.c', 'cursor_type.c', 'cursor_int.c', 'lobject_type.c', 'lobject_int.c', 'adapter_qstring.c', 'adapter_pboolean.c', 'adapter_binary.c', - 'adapter_asis.c', 'adapter_list.c'] + 'adapter_asis.c', 'adapter_list.c', 'utils.c'] parser = ConfigParser.ConfigParser() parser.read('setup.cfg') |