Skip to content

inserttable: remove arbitrary buffer limit (solves #79) #91

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Cito merged 1 commit into from
Jan 2, 2026
Merged

inserttable: remove arbitrary buffer limit (solves #79) #91

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
179 changes: 75 additions & 104 deletions ext/pgconn.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -716,9 +716,9 @@ static PyObject *
conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
{
PGresult *result;
char *table, *buffer, *bufpt, *bufmax, *s, *t;
char *table, *s, *t;
struct CharBuffer buffer = {0};
int freeze = 0, encoding, ret;
size_t bufsiz;
PyObject *rows, *iter_row, *item, *columns = NULL;
Py_ssize_t i, j, m, n;

Expand Down Expand Up @@ -775,37 +775,35 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
n = -1; /* number of columns not yet known */
}

/* allocate buffer */
if (!(buffer = PyMem_Malloc(MAX_BUFFER_SIZE))) {
Py_DECREF(iter_row);
return PyErr_NoMemory();
}

encoding = PQclientEncoding(self->cnx);

/* starts query */
bufpt = buffer;
bufmax = bufpt + MAX_BUFFER_SIZE;
bufpt += snprintf(bufpt, (size_t)(bufmax - bufpt), "copy ");
ext_char_buffer_s(&buffer, "copy ");

/* return early if there is no buffer */
if (buffer.error) {
if (buffer.data)
PyMem_Free(buffer.data);
Py_DECREF(iter_row);
return PyErr_NoMemory();
}

s = table;
do {
t = strchr(s, '.');
if (!t)
t = s + strlen(s);
table = PQescapeIdentifier(self->cnx, s, (size_t)(t - s));
if (bufpt < bufmax)
bufpt += snprintf(bufpt, (size_t)(bufmax - bufpt), "%s", table);
ext_char_buffer_s(&buffer, table);
PQfreemem(table);
s = t;
if (*s && bufpt < bufmax)
*bufpt++ = *s++;
if (*s)
ext_char_buffer_c(&buffer, *s++);
} while (*s);

if (columns) {
/* adds a string like f" ({','.join(columns)})" */
if (bufpt < bufmax)
bufpt += snprintf(bufpt, (size_t)(bufmax - bufpt), " (");
ext_char_buffer_s(&buffer, " (");
for (j = 0; j < n; ++j) {
PyObject *obj = PySequence_Fast_GET_ITEM(columns, j);
Py_ssize_t slen;
Expand All @@ -817,7 +815,7 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
else if (PyUnicode_Check(obj)) {
obj = get_encoded_string(obj, encoding);
if (!obj) {
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(iter_row);
return NULL; /* pass the UnicodeEncodeError */
}
Expand All @@ -826,37 +824,34 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
PyErr_SetString(
PyExc_TypeError,
"The third argument must contain only strings");
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(iter_row);
return NULL;
}
PyBytes_AsStringAndSize(obj, &col, &slen);
col = PQescapeIdentifier(self->cnx, col, (size_t)slen);
Py_DECREF(obj);
if (bufpt < bufmax)
bufpt += snprintf(bufpt, (size_t)(bufmax - bufpt), "%s%s", col,
j == n - 1 ? ")" : ",");
ext_char_buffer_s(&buffer, col);
ext_char_buffer_c(&buffer, j == n - 1 ? ')' : ',');
PQfreemem(col);
}
}
if (bufpt < bufmax)
bufpt += snprintf(bufpt, (size_t)(bufmax - bufpt), " from stdin");
if (freeze && bufpt < bufmax) {
bufpt += snprintf(bufpt, (size_t)(bufmax - bufpt), " freeze");
}
ext_char_buffer_s(&buffer, " from stdin");
if (freeze)
ext_char_buffer_s(&buffer, " freeze");

if (bufpt >= bufmax) {
PyMem_Free(buffer);
if (buffer.error) {
PyMem_Free(buffer.data);
Py_DECREF(iter_row);
return PyErr_NoMemory();
}

Py_BEGIN_ALLOW_THREADS
result = PQexec(self->cnx, buffer);
result = PQexec(self->cnx, buffer.data);
Py_END_ALLOW_THREADS

if (!result || PQresultStatus(result) != PGRES_COPY_IN) {
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(iter_row);
PyErr_SetString(PyExc_ValueError, PQerrorMessage(self->cnx));
return NULL;
Expand All @@ -871,7 +866,7 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)

if (!(PyTuple_Check(columns) || PyList_Check(columns))) {
PQputCopyEnd(self->cnx, "Invalid arguments");
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(columns);
Py_DECREF(columns);
Py_DECREF(iter_row);
Expand All @@ -887,7 +882,7 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
}
else if (j != n) {
PQputCopyEnd(self->cnx, "Invalid arguments");
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(columns);
Py_DECREF(iter_row);
PyErr_SetString(
Expand All @@ -896,65 +891,54 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
return NULL;
}

/* reset buffer to empty */
buffer.len = 0;

/* builds insert line */
bufpt = buffer;
bufsiz = MAX_BUFFER_SIZE - 1;

for (j = 0; j < n; ++j) {
if (j) {
*bufpt++ = '\t';
--bufsiz;
}
if (j)
ext_char_buffer_c(&buffer, '\t');

item = PySequence_Fast_GET_ITEM(columns, j);

/* convert item to string and append to buffer */
if (item == Py_None) {
if (bufsiz > 2) {
*bufpt++ = '\\';
*bufpt++ = 'N';
bufsiz -= 2;
}
else
bufsiz = 0;
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, 'N');
}
else if (PyBytes_Check(item)) {
const char *t = PyBytes_AsString(item);

while (*t && bufsiz) {
while (*t) {
switch (*t) {
case '\\':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = '\\';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\\');
break;
case '\t':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 't';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\t');
break;
case '\r':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 'r';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\r');
break;
case '\n':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 'n';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\n');
break;
default:
*bufpt++ = *t;
ext_char_buffer_c(&buffer, *t);
}
++t;
--bufsiz;
}
}
else if (PyUnicode_Check(item)) {
PyObject *s = get_encoded_string(item, encoding);
if (!s) {
PQputCopyEnd(self->cnx, "Encoding error");
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(item);
Py_DECREF(columns);
Py_DECREF(iter_row);
Expand All @@ -963,33 +947,28 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
else {
const char *t = PyBytes_AsString(s);

while (*t && bufsiz) {
while (*t) {
switch (*t) {
case '\\':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = '\\';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\\');
break;
case '\t':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 't';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\t');
break;
case '\r':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 'r';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\r');
break;
case '\n':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 'n';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\n');
break;
default:
*bufpt++ = *t;
ext_char_buffer_c(&buffer, *t);
}
++t;
--bufsiz;
}
Py_DECREF(s);
}
Expand All @@ -998,50 +977,42 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)
PyObject *s = PyObject_Str(item);
const char *t = PyUnicode_AsUTF8(s);

while (*t && bufsiz) {
*bufpt++ = *t++;
--bufsiz;
}
ext_char_buffer_s(&buffer, t);
Py_DECREF(s);
}
else {
PyObject *s = PyObject_Repr(item);
const char *t = PyUnicode_AsUTF8(s);

while (*t && bufsiz) {
while (*t) {
switch (*t) {
case '\\':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = '\\';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\\');
break;
case '\t':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 't';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\t');
break;
case '\r':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 'r';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\r');
break;
case '\n':
*bufpt++ = '\\';
if (--bufsiz)
*bufpt++ = 'n';
ext_char_buffer_c(&buffer, '\\');
ext_char_buffer_c(&buffer, '\n');
break;
default:
*bufpt++ = *t;
ext_char_buffer_c(&buffer, *t);
}
++t;
--bufsiz;
}
Py_DECREF(s);
}

if (bufsiz <= 0) {
if (buffer.error) {
PQputCopyEnd(self->cnx, "Memory error");
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(columns);
Py_DECREF(iter_row);
return PyErr_NoMemory();
Expand All @@ -1050,36 +1021,36 @@ conn_inserttable(connObject *self, PyObject *args, PyObject *kwds)

Py_DECREF(columns);

*bufpt++ = '\n';
ext_char_buffer_c(&buffer, '\n');

/* sends data */
ret = PQputCopyData(self->cnx, buffer, (int)(bufpt - buffer));
ret = PQputCopyData(self->cnx, buffer.data, (int)buffer.len);
if (ret != 1) {
char *errormsg = ret == -1 ? PQerrorMessage(self->cnx)
: "Data cannot be queued";
PyErr_SetString(PyExc_IOError, errormsg);
PQputCopyEnd(self->cnx, errormsg);
PyMem_Free(buffer);
PyMem_Free(buffer.data);
Py_DECREF(iter_row);
return NULL;
}
}

Py_DECREF(iter_row);
if (PyErr_Occurred()) {
PyMem_Free(buffer);
PyMem_Free(buffer.data);
return NULL; /* pass the iteration error */
}

ret = PQputCopyEnd(self->cnx, NULL);
if (ret != 1) {
PyErr_SetString(PyExc_IOError, ret == -1 ? PQerrorMessage(self->cnx)
: "Data cannot be queued");
PyMem_Free(buffer);
PyMem_Free(buffer.data);
return NULL;
}

PyMem_Free(buffer);
PyMem_Free(buffer.data);

Py_BEGIN_ALLOW_THREADS
result = PQgetResult(self->cnx);
Expand Down
Loading