From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 27 Oct 2020 23:32:07 +0000 (+0100)
Subject: Added PQescapeString wrapper
X-Git-Tag: 3.0.dev0~445
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=a0362f19f8c96b4180289ed6e0b780a7cfa9ae05;p=thirdparty%2Fpsycopg.git

Added PQescapeString wrapper

`psycopg3.pq.Escaping.escape_string()` can now work without a connection
too, with the limitation explained in the libpq docs.

Also fixed memory leak. Oops.
---

diff --git a/psycopg3/psycopg3/pq/_pq_ctypes.py b/psycopg3/psycopg3/pq/_pq_ctypes.py
index fa45ef7fe..9aaa39fb8 100644
--- a/psycopg3/psycopg3/pq/_pq_ctypes.py
+++ b/psycopg3/psycopg3/pq/_pq_ctypes.py
@@ -402,7 +402,10 @@ PQescapeStringConn = pq.PQescapeStringConn
 # ]
 PQescapeStringConn.restype = c_size_t
 
-# won't wrap: PQescapeString
+PQescapeString = pq.PQescapeString
+# TODO: raises "wrong type" error
+# PQescapeString.argtypes = [c_char_p, c_char_p, c_size_t]
+PQescapeString.restype = c_size_t
 
 PQescapeByteaConn = pq.PQescapeByteaConn
 PQescapeByteaConn.argtypes = [
diff --git a/psycopg3/psycopg3/pq/_pq_ctypes.pyi b/psycopg3/psycopg3/pq/_pq_ctypes.pyi
index 90804f538..ea443727b 100644
--- a/psycopg3/psycopg3/pq/_pq_ctypes.pyi
+++ b/psycopg3/psycopg3/pq/_pq_ctypes.pyi
@@ -67,6 +67,7 @@ def PQescapeStringConn(
     arg4: int,
     arg5: pointer[c_int],
 ) -> int: ...
+def PQescapeString(arg1: c_char_p, arg2: bytes, arg3: int) -> int: ...
 def PQsendPrepare(
     arg1: Optional[PGconn_struct],
     arg2: bytes,
diff --git a/psycopg3/psycopg3/pq/pq_ctypes.py b/psycopg3/psycopg3/pq/pq_ctypes.py
index 588b9a83b..4147436d9 100644
--- a/psycopg3/psycopg3/pq/pq_ctypes.py
+++ b/psycopg3/psycopg3/pq/pq_ctypes.py
@@ -797,7 +797,13 @@ class Escaping:
             return out.value
 
         else:
-            raise PQerror("escape_identifier failed: no connection provided")
+            out = create_string_buffer(len(data) * 2 + 1)
+            impl.PQescapeString(
+                pointer(out),  # type: ignore
+                data,
+                len(data),
+            )
+            return out.value
 
     def escape_bytea(self, data: bytes) -> bytes:
         len_out = c_size_t()
diff --git a/psycopg3_c/psycopg3_c/libpq.pxd b/psycopg3_c/psycopg3_c/libpq.pxd
index 908f4e3f2..a9e241d59 100644
--- a/psycopg3_c/psycopg3_c/libpq.pxd
+++ b/psycopg3_c/psycopg3_c/libpq.pxd
@@ -183,12 +183,12 @@ cdef extern from "libpq-fe.h":
     Oid PQoidValue(const PGresult *res)
 
     # 33.3.4. Escaping Strings for Inclusion in SQL Commands
-    # TODO: PQescapeStringConn PQescapeString
     char *PQescapeIdentifier(PGconn *conn, const char *str, size_t length)
     char *PQescapeLiteral(PGconn *conn, const char *str, size_t length)
     size_t PQescapeStringConn(PGconn *conn,
                               char *to, const char *from_, size_t length,
                               int *error)
+    size_t PQescapeString(char *to, const char *from_, size_t length)
     unsigned char *PQescapeByteaConn(PGconn *conn,
                                      const unsigned char *src,
                                      size_t from_length,
diff --git a/psycopg3_c/psycopg3_c/pq_cython.pyx b/psycopg3_c/psycopg3_c/pq_cython.pyx
index ca0b0c576..567c3d8c8 100644
--- a/psycopg3_c/psycopg3_c/pq_cython.pyx
+++ b/psycopg3_c/psycopg3_c/pq_cython.pyx
@@ -838,6 +838,8 @@ cdef class Escaping:
         cdef size_t len_data = len(data)
         cdef char *out
         cdef size_t len_out
+        cdef bytes rv
+
         if self.conn is not None:
             if self.conn.pgconn_ptr is NULL:
                 raise PQerror("the connection is closed")
@@ -848,13 +850,22 @@ cdef class Escaping:
             )
 
             if error:
+                PyMem_Free(out)
                 raise PQerror(
                     f"escape_string failed: {error_message(self.conn)}"
                 )
-            return out[:len_out]
+
+            rv = out[:len_out]
+            PyMem_Free(out)
+            return rv
 
         else:
-            raise PQerror("escape_identifier failed: no connection provided")
+            out = <char *>PyMem_Malloc(len_data * 2 + 1)
+            len_out = impl.PQescapeString(out, data, len_data)
+            rv = out[:len_out]
+            PyMem_Free(out)
+            return rv
+
 
     def escape_bytea(self, data: bytes) -> bytes:
         cdef size_t len_out
diff --git a/tests/pq/test_escaping.py b/tests/pq/test_escaping.py
index 25fd1e072..b55c5132e 100644
--- a/tests/pq/test_escaping.py
+++ b/tests/pq/test_escaping.py
@@ -116,11 +116,22 @@ def test_escape_string_1char(pgconn, scs):
         assert rv == exp
 
 
-def test_escape_string_noconn(pgconn):
+@pytest.mark.parametrize(
+    "data, want",
+    [
+        (b"", b""),
+        (b"hello", b"hello"),
+        (b"foo'bar", b"foo''bar"),
+        (b"foo\\bar", b"foo\\\\bar"),
+    ],
+)
+def test_escape_string_noconn(data, want):
     esc = pq.Escaping()
-    with pytest.raises(psycopg3.OperationalError):
-        esc.escape_string(b"hi")
+    out = esc.escape_string(data)
+    assert out == want
+
 
+def test_escape_string_badconn(pgconn):
     esc = pq.Escaping(pgconn)
     pgconn.finish()
     with pytest.raises(psycopg3.OperationalError):