Set a memory exception in psycopg_escape_string

...otherwise all the callers should set it.

3 files changed, 25 insertions, 13 deletions

diff --git a/psycopg/adapter_qstring.c b/psycopg/adapter_qstring.c
index f2efac3..6f7bf50 100644
--- a/psycopg/adapter_qstring.c
+++ b/psycopg/adapter_qstring.c
@@ -75,7 +75,6 @@ qstring_quote(qstringObject *self)
     Bytes_AsStringAndSize(str, &s, &len);
     if (!(buffer = psycopg_escape_string(self->conn, s, len, NULL, &qlen))) {
         Py_DECREF(str);
-        PyErr_NoMemory();
         return NULL;
     }
 
diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c
index 983c4a9..9351a3c 100644
--- a/psycopg/cursor_type.c
+++ b/psycopg/cursor_type.c
@@ -1375,13 +1375,11 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
 
     if (!(quoted_delimiter = psycopg_escape_string(
             (PyObject*)self->conn, sep, 0, NULL, NULL))) {
-        PyErr_NoMemory();
         goto exit;
     }
 
     if (!(quoted_null = psycopg_escape_string(
             (PyObject*)self->conn, null, 0, NULL, NULL))) {
-        PyErr_NoMemory();
         goto exit;
     }
 
@@ -1471,13 +1469,11 @@ psyco_curs_copy_to(cursorObject *self, PyObject *args, PyObject *kwargs)
 
     if (!(quoted_delimiter = psycopg_escape_string(
             (PyObject*)self->conn, sep, 0, NULL, NULL))) {
-        PyErr_NoMemory();
         goto exit;
     }
 
     if (!(quoted_null = psycopg_escape_string(
             (PyObject*)self->conn, null, 0, NULL, NULL))) {
-        PyErr_NoMemory();
         goto exit;
     }
 
diff --git a/psycopg/utils.c b/psycopg/utils.c
index 698dcb4..ce10469 100644
--- a/psycopg/utils.c
+++ b/psycopg/utils.c
@@ -32,21 +32,34 @@
 #include <string.h>
 #include <stdlib.h>
 
+/* Escape a string for sql inclusion.
+ *
+ * The function must be called holding the GIL.
+ *
+ * Return a pointer to a new string on the Python heap on success, else NULL
+ * and set an exception. The returned string includes quotes and leading E if
+ * needed.
+ *
+ * If tolen is set, it will contain the length of the escaped string,
+ * including quotes.
+ */
 char *
 psycopg_escape_string(PyObject *obj, const char *from, Py_ssize_t len,
                        char *to, Py_ssize_t *tolen)
 {
     Py_ssize_t ql;
     connectionObject *conn = (connectionObject*)obj;
-    int eq = (conn && (conn->equote)) ? 1 : 0;   
+    int eq = (conn && (conn->equote)) ? 1 : 0;
 
     if (len == 0)
         len = strlen(from);
-    
+
     if (to == NULL) {
         to = (char *)PyMem_Malloc((len * 2 + 4) * sizeof(char));
-        if (to == NULL)
+        if (to == NULL) {
+            PyErr_NoMemory();
             return NULL;
+        }
     }
 
     {
@@ -59,15 +72,19 @@ psycopg_escape_string(PyObject *obj, const char *from, Py_ssize_t len,
                 ql = PQescapeString(to+eq+1, from, len);
     }
 
-    if (eq)
+    if (eq) {
         to[0] = 'E';
-    to[eq] = '\''; 
-    to[ql+eq+1] = '\'';
-    to[ql+eq+2] = '\0';
+        to[1] = to[ql+2] = '\'';
+        to[ql+3] = '\0';
+    }
+    else {
+        to[0] = to[ql+1] = '\'';
+        to[ql+2] = '\0';
+    }
 
     if (tolen)
         *tolen = ql+eq+2;
-        
+
     return to;
 }