Updated patch to quote delimiter and null characters in copy_* functions

Alejandro Dubrovsky <[email protected]>
Newsgroups gmane.comp.python.db.psycopg.devel
Message-ID <[email protected]> 
Refer to my previous email for rational.

This patch enlarges the default size of the copy buffer and also checks
to see if the copy query got truncated.  ie copy_to(somefile, '(some
long query)') is very handy and I went crazy with it until I broke the
first patch.

_______________________________________________
Psycopg mailing list
Psycopg-IAPFreCvJWPBWskQ1e/[email protected]
http://lists.initd.org/mailman/listinfo/psycopg
psycocopypatch2.diff (text/x-patch, 8.2 KB) 
=== modified file 'psycopg/adapter_qstring.c'
--- psycopg/adapter_qstring.c	2008-07-26 13:10:02 +0000
+++ psycopg/adapter_qstring.c	2008-09-19 04:00:50 +0000
@@ -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 @@
 }
 #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;
 

=== modified file 'psycopg/cursor_type.c'
--- psycopg/cursor_type.c	2008-07-21 05:41:54 +0000
+++ psycopg/cursor_type.c	2008-09-23 04:03:32 +0000
@@ -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 @@
 
 #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 @@
 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 @@
 
     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 @@
         Py_INCREF(Py_None);
     }
 
+    if (query && (query != query_buffer)) {
+        free(query);
+    }
     self->copyfile =NULL;
 
     return res;
@@ -1260,11 +1296,14 @@
 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 @@
         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 @@
         res = Py_None;
         Py_INCREF(Py_None);
     }
-
+    if (query && (query != query_buffer)) {
+        free(query);
+    }
     self->copyfile = NULL;
 
     return res;

=== modified file 'psycopg/psycopg.h'
--- psycopg/psycopg.h	2008-08-02 08:30:36 +0000
+++ psycopg/psycopg.h	2008-09-19 04:00:50 +0000
@@ -143,7 +143,7 @@
 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 \

=== modified file 'setup.py'
--- setup.py	2008-09-19 19:25:16 +0000
+++ setup.py	2008-09-23 02:22:42 +0000
@@ -330,7 +330,7 @@
     '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')
lmpx.com only provides a reader for public news (NNTP) servers. It is not affiliated with the servers or forums shown here and is not responsible for the content of articles, which is written by their respective authors.