From 650f372b4ca4efa2a43e61b109a9569b3314a787 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 17 Dec 2020 13:55:56 +0100
Subject: [PATCH] Reduce re-allocations in objects quoting

---
 psycopg3/psycopg3/adapt.py          |  2 ++
 psycopg3/psycopg3/proto.py          |  3 ++
 psycopg3_c/psycopg3_c/adapt.pyx     | 53 ++++++++++++++++++-----------
 psycopg3_c/psycopg3_c/pq_cython.pxd |  5 +++
 psycopg3_c/psycopg3_c/pq_cython.pyx |  4 ++-
 5 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/psycopg3/psycopg3/adapt.py b/psycopg3/psycopg3/adapt.py
index 03eca5904..55dfbded4 100644
--- a/psycopg3/psycopg3/adapt.py
+++ b/psycopg3/psycopg3/adapt.py
@@ -36,6 +36,8 @@ class Dumper(ABC):
         """Convert the object *obj* to PostgreSQL representation."""
         ...
 
+    # TODO: the protocol signature should probably return a Buffer like object
+    # (the C implementation may return bytearray)
     def quote(self, obj: Any) -> bytes:
         """Convert the object *obj* to escaped representation."""
         value = self.dump(obj)
diff --git a/psycopg3/psycopg3/proto.py b/psycopg3/psycopg3/proto.py
index 328f716bc..3deb72dc2 100644
--- a/psycopg3/psycopg3/proto.py
+++ b/psycopg3/psycopg3/proto.py
@@ -53,6 +53,9 @@ LoadFunc = Callable[[bytes], Any]
 LoaderType = Type["Loader"]
 LoadersMap = Dict[Tuple[int, Format], LoaderType]
 
+# TODO: Loader, Dumper should probably become protocols
+# as there are both C and a Python implementation
+
 
 class Transformer(Protocol):
     def __init__(self, context: AdaptContext = None):
diff --git a/psycopg3_c/psycopg3_c/adapt.pyx b/psycopg3_c/psycopg3_c/adapt.pyx
index 303265478..62c928126 100644
--- a/psycopg3_c/psycopg3_c/adapt.pyx
+++ b/psycopg3_c/psycopg3_c/adapt.pyx
@@ -16,12 +16,16 @@ equivalent C implementations.
 from typing import Any
 
 from cpython.bytes cimport PyBytes_AsStringAndSize
+from cpython.bytearray cimport PyByteArray_FromStringAndSize, PyByteArray_Resize
+from cpython.bytearray cimport PyByteArray_AS_STRING
 
 from psycopg3_c cimport libpq as impl
 from psycopg3_c.adapt cimport cloader_func, get_context_func
-from psycopg3_c.pq_cython cimport Escaping
+from psycopg3_c.pq_cython cimport Escaping, _buffer_as_string_and_size
 
+from psycopg3 import errors as e
 from psycopg3.pq import Format
+from psycopg3.pq.misc import error_message
 
 import logging
 logger = logging.getLogger("psycopg3.adapt")
@@ -56,27 +60,38 @@ cdef class CDumper:
     def dump(self, obj: Any) -> bytes:
         raise NotImplementedError()
 
-    def quote(self, obj: Any) -> bytes:
-        # TODO: can be optimized
-        cdef object ovalue = self.dump(obj)
-
-        cdef bytes value
-        if isinstance(ovalue, bytes):
-            value = ovalue
+    def quote(self, obj: Any) -> bytearray:
+        cdef char *ptr
+        cdef char *ptr_out
+        cdef Py_ssize_t length, len_out
+        cdef int error
+        cdef bytearray rv
+
+        pyout = self.dump(obj)
+        _buffer_as_string_and_size(pyout, &ptr, &length)
+        rv = PyByteArray_FromStringAndSize("", 0)
+        PyByteArray_Resize(rv, length * 2 + 3)  # Must include the quotes
+        ptr_out = PyByteArray_AS_STRING(rv)
+
+        if self._pgconn is not None:
+            if self._pgconn.pgconn_ptr == NULL:
+                raise e.OperationalError("the connection is closed")
+
+            len_out = impl.PQescapeStringConn(
+                self._pgconn.pgconn_ptr, ptr_out + 1, ptr, length, &error
+            )
+            if error:
+                raise e.OperationalError(
+                    f"escape_string failed: {error_message(self.connection)}"
+                )
         else:
-            value = bytes(ovalue)
+            len_out = impl.PQescapeString(ptr_out + 1, ptr, length)
 
-        cdef bytes tmp
-        cdef Escaping esc
+        ptr_out[0] = b'\''
+        ptr_out[len_out + 1] = b'\''
+        PyByteArray_Resize(rv, len_out + 2)
 
-        if self.connection:
-            esc = Escaping(self._pgconn)
-            return bytes(esc.escape_literal(value))
-
-        else:
-            esc = Escaping()
-            tmp = bytes(esc.escape_string(value))
-            return b"'%s'" % tmp
+        return rv
 
     @property
     def oid(self) -> int:
diff --git a/psycopg3_c/psycopg3_c/pq_cython.pxd b/psycopg3_c/psycopg3_c/pq_cython.pxd
index faeb46094..b6537409e 100644
--- a/psycopg3_c/psycopg3_c/pq_cython.pxd
+++ b/psycopg3_c/psycopg3_c/pq_cython.pxd
@@ -44,3 +44,8 @@ cdef class PQBuffer:
 
     @staticmethod
     cdef PQBuffer _from_buffer(unsigned char *buf, Py_ssize_t len)
+
+
+cdef int _buffer_as_string_and_size(
+    data: "Buffer", char **ptr, Py_ssize_t *length
+) except -1
diff --git a/psycopg3_c/psycopg3_c/pq_cython.pyx b/psycopg3_c/psycopg3_c/pq_cython.pyx
index a86e540aa..fac39a5b2 100644
--- a/psycopg3_c/psycopg3_c/pq_cython.pyx
+++ b/psycopg3_c/psycopg3_c/pq_cython.pyx
@@ -962,7 +962,9 @@ cdef class PQBuffer:
         pass
 
 
-cdef int _buffer_as_string_and_size(data: "Buffer", char **ptr, Py_ssize_t *length) except -1:
+cdef int _buffer_as_string_and_size(
+    data: "Buffer", char **ptr, Py_ssize_t *length
+) except -1:
     cdef Py_buffer buf
 
     if isinstance(data, bytes):
-- 
2.47.2