Re: 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]> 
On Tue, 2008-09-23 at 15:34 +0200, Federico Di Gregorio wrote:
> Il giorno mar, 23/09/2008 alle 14.38 +1000, Alejandro Dubrovsky ha scritto:
> > 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.
> 
> Your patch seems fine to me and it surely add needed functionality but
> I don't like to have to different quoting functions. Do you have the
> time to move all functionality of qstring_escape() to
> psycopg_internal_escape_string() and then have qstring_escape call that?
> Also note that the patch you sent does not include utils.h and utils.c.
> 
I'm not sure I understand what you are suggesting.  In the patch I sent,
psycopg_internal_escape_string calls qstring_escape, there aren't two
quoting functions.  psycopg_internal_escape_string is just a wrapper
around qstring_escape that makes it easier on the caller.  

Would you prefer the calling structure to be in reverse? (qstring_escape
calling psycopg2_internal_escape_string)  That seems wrong to me since
psycopg_internal_escape_string creates a new buffer, but qstring_escape
is given one, so qstring_escape would just copy the result from
psycopg_internal_escape_string and free the newly allocated buffer, even
though there was no need to create or destroy the buffer in the first
place.

Getting rid of qstring_escape altogether, on the other hand, that I
could understand.


Regarding missing utils.[ch], I noticed when I received it, and I
promptly replied, to myself :(

Attached same version as yesterday but with utils.*

_______________________________________________
Psycopg mailing list
Psycopg-IAPFreCvJWPBWskQ1e/[email protected]
http://lists.initd.org/mailman/listinfo/psycopg
psycocopypatch2.diff (text/x-patch, 9.8 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 \

=== added file 'psycopg/utils.c'
--- psycopg/utils.c	1970-01-01 00:00:00 +0000
+++ psycopg/utils.c	2008-09-19 04:00:50 +0000
@@ -0,0 +1,43 @@
+/* utils.c - miscellaneous utility functions
+ *
+ */
+
+#include "psycopg/config.h"
+#include "psycopg/utils.h"
+#include "psycopg/psycopg.h"
+#include "psycopg/connection.h"
+#include "psycopg/pgtypes.h"
+#include "psycopg/pgversion.h"
+#include <string.h>
+#include <stdlib.h>
+
+char *psycopg_internal_escape_string(connectionObject *conn, const char *string)
+{
+    char *buffer;
+    size_t string_length;
+    int equote;         /* buffer offset if E'' quotes are needed */  
+
+    string_length = strlen(string);
+    
+    buffer = (char *) malloc((string_length * 2 + 4) * sizeof(char));
+    if (buffer == NULL) {
+        return NULL;
+    }
+
+    equote = (conn && (conn->equote)) ? 1 : 0;
+
+    { 
+        size_t qstring_length;
+
+        qstring_length = qstring_escape(buffer + equote + 1, string, string_length,
+                                        (conn ? conn->pgconn : NULL));
+
+        if (equote)
+            buffer[0] = 'E';
+        buffer[equote] = '\''; 
+        buffer[qstring_length + equote + 1] = '\'';
+        buffer[qstring_length + equote + 2] = 0;
+    }
+
+    return buffer;
+}

=== added file 'psycopg/utils.h'
--- psycopg/utils.h	1970-01-01 00:00:00 +0000
+++ psycopg/utils.h	2008-09-19 04:00:50 +0000
@@ -0,0 +1,16 @@
+/* utils.h - miscellaneous utility functions
+ *
+ */
+
+#ifndef PSYCOPG_UTILS_H
+#define PSYCOPG_UTILS_H 1
+
+#include "psycopg/config.h"
+#include "psycopg/connection.h"
+
+
+HIDDEN char *psycopg_internal_escape_string(connectionObject *conn, const char *string);
+
+#endif /* !defined(PSYCOPG_UTILS_H) */
+
+

=== 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.