fix(c): fix copy buffer corruption on dumping error

Restore the buffer so that if we trap the exception and continue in
Python the buffer will not be corrupted.

Improve on the #1147 fix.

diff --git a/psycopg/psycopg/_copy_base.py b/psycopg/psycopg/_copy_base.py
index 05180950e6e66b4fb5f164b17ccf0fde3b0875d6..19d016d9a4e2776daa3d8a1ae8779035b5882b98 100644 (file)
--- a/psycopg/psycopg/_copy_base.py
+++ b/psycopg/psycopg/_copy_base.py
@@ -348,8 +348,8 @@ def _format_row_binary(
     if out is None:
         out = bytearray()
 
-    out += _pack_int2(len(row))
     adapted = tx.dump_sequence(row, [PY_BINARY] * len(row))
+    out += _pack_int2(len(row))
     for b in adapted:
         if b is not None:
             out += _pack_int4(len(b))

diff --git a/psycopg_c/psycopg_c/_psycopg/copy.pyx b/psycopg_c/psycopg_c/_psycopg/copy.pyx
index d28e2e8a48545af697906be0a2c9ad183c12fe3b..bd3ac10ac8e49bed14a1357f6aa37771ac32e5b1 100644 (file)
--- a/psycopg_c/psycopg_c/_psycopg/copy.pyx
+++ b/psycopg_c/psycopg_c/_psycopg/copy.pyx
@@ -21,8 +21,8 @@ from psycopg import errors as e
 cdef int32_t _binary_null = -1
 
 
-def format_row_binary(
-    row: Sequence[Any], tx: Transformer, out: bytearray = None
+def _format_row_binary(
+    row: Sequence[Any], tx: Transformer, out: bytearray
 ) -> bytearray:
     """Convert a row of adapted data to the data to send for binary copy"""
     cdef Py_ssize_t rowlen
@@ -34,12 +34,8 @@ def format_row_binary(
         rowlen = len(row)
     cdef uint16_t berowlen = endian.htobe16(<int16_t>rowlen)
 
-    cdef Py_ssize_t pos  # offset in 'out' where to write
-    if out is None:
-        out = PyByteArray_FromStringAndSize("", 0)
-        pos = 0
-    else:
-        pos = PyByteArray_GET_SIZE(out)
+    # offset in 'out' where to write
+    cdef Py_ssize_t pos = PyByteArray_GET_SIZE(out)
 
     cdef char *target = CDumper.ensure_size(out, pos, sizeof(berowlen))
 
@@ -102,6 +98,25 @@ def format_row_binary(
     return out
 
 
+def format_row_binary(
+    row: Sequence[Any], tx: Transformer, out: bytearray = None
+) -> bytearray:
+    cdef Py_ssize_t size
+    if out is None:
+        out = PyByteArray_FromStringAndSize("", 0)
+        size = 0
+    else:
+        size = PyByteArray_GET_SIZE(out)
+
+    try:
+        _format_row_binary(row, tx, out)
+    except Exception as e:
+        PyByteArray_Resize(out, size)
+        raise e
+
+    return out
+
+
 cdef int _append_binary_none(bytearray out, Py_ssize_t *pos) except -1:
     cdef char *target
     target = CDumper.ensure_size(out, pos[0], sizeof(_binary_null))
@@ -110,28 +125,11 @@ cdef int _append_binary_none(bytearray out, Py_ssize_t *pos) except -1:
     return 0
 
 
-def format_row_text(
-    row: Sequence[Any], tx: Transformer, out: bytearray = None
+def _format_row_text(
+    row: Sequence[Any], tx: Transformer, out: bytearray
 ) -> bytearray:
-    cdef Py_ssize_t pos  # offset in 'out' where to write
-    if out is None:
-        out = PyByteArray_FromStringAndSize("", 0)
-        pos = 0
-    else:
-        pos = PyByteArray_GET_SIZE(out)
-
-    cdef Py_ssize_t rowlen
-    if type(row) is list:
-        rowlen = PyList_GET_SIZE(row)
-    elif type(row) is tuple:
-        rowlen = PyTuple_GET_SIZE(row)
-    else:
-        rowlen = len(row)
-
-    if rowlen == 0:
-        PyByteArray_Resize(out, pos + 1)
-        out[pos] = b"\n"
-        return out
+    # offset in 'out' where to write
+    cdef Py_ssize_t pos = PyByteArray_GET_SIZE(out)
 
     cdef Py_ssize_t size, tmpsize
     cdef char *buf
@@ -141,6 +139,7 @@ def format_row_text(
     cdef int with_tab
     cdef PyObject *fmt = <PyObject *>PG_TEXT
     cdef PyObject *row_dumper
+    cdef Py_ssize_t rowlen = len(row)
 
     # try to get preloaded dumpers from set_types
     if not tx._row_dumpers:
@@ -219,6 +218,32 @@ def format_row_text(
     return out
 
 
+def format_row_text(
+    row: Sequence[Any], tx: Transformer, out: bytearray = None
+) -> bytearray:
+    cdef Py_ssize_t size
+    if out is None:
+        out = PyByteArray_FromStringAndSize("", 0)
+        size = 0
+    else:
+        size = PyByteArray_GET_SIZE(out)
+
+    # exit early, if the row is empty
+    cdef Py_ssize_t rowlen = len(row)
+    if rowlen == 0:
+        PyByteArray_Resize(out, size + 1)
+        out[size] = b"\n"
+        return out
+
+    try:
+        _format_row_text(row, tx, out)
+    except Exception as e:
+        PyByteArray_Resize(out, size)
+        raise e
+
+    return out
+
+
 cdef int _append_text_none(bytearray out, Py_ssize_t *pos, int with_tab) except -1:
     cdef char *target