pglib 5.11.0__tar.gz → 5.12.0__tar.gz

{pglib-5.11.0/pglib.egg-info → pglib-5.12.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pglib
-Version: 5.11.0
+Version: 5.12.0
 Summary: A PostgreSQL interface
 Author-email: Michael Kleehammer <michael@kleehammer.com>
 License-Expression: MIT

{pglib-5.11.0 → pglib-5.12.0/pglib.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pglib
-Version: 5.11.0
+Version: 5.12.0
 Summary: A PostgreSQL interface
 Author-email: Michael Kleehammer <michael@kleehammer.com>
 License-Expression: MIT

{pglib-5.11.0 → pglib-5.12.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pglib"
-version = "5.11.0"
+version = "5.12.0"
 description = "A PostgreSQL interface"
 readme = "README.rst"
 license = "MIT"

{pglib-5.11.0 → pglib-5.12.0}/src/connection.cpp

@@ -365,6 +365,113 @@ static PyObject* Connection_copy_from(PyObject* self, PyObject* args)
 }
 
 
+const char* doc_copy_to =
+    "Connection.copy_to(command, dest) --> int\n"
+    "\n"
+    "Executes the given COPY TO command and returns the number of records copied.\n"
+    "\n"
+    "command\n"
+    "  The copy command which must be 'to stdout'.\n"
+    "\n"
+    "dest\n"
+    "  The file-like object to write to.  Strings will be written, not bytes, so\n"
+    "  open in text mode.\n"
+    "\n"
+    "Examples:\n"
+    "  cnxn.copy_to('copy t1 to stdout csv header', open('test.csv', 'w'))\n"
+    "  cnxn.copy_to('copy t1(a, b, c) to stdout csv header', open('test.csv', 'w'))\n"
+    "  cnxn.copy_to('copy (select * from t1 where id > 10) to stdout csv', f)\n";
+
+static PyObject* Connection_copy_to(PyObject* self, PyObject* args)
+{
+    PyObject* command;
+    PyObject* dest;
+    if (!PyArg_ParseTuple(args, "UO", &command, &dest))
+        return 0;
+
+    Connection* cnxn = CastConnection(self, REQUIRE_OPEN);
+    if (!cnxn)
+        return 0;
+
+    if (!PyObject_HasAttrString(dest, "write"))
+        return PyErr_Format(Error, "Destination must be a file-like object.");
+    Object write_method(PyObject_GetAttrString(dest, "write"));
+
+    const char* szSQL = PyUnicode_AsUTF8(command);
+    if (!szSQL)
+        return 0;
+
+    ResultHolder result;
+    Py_BEGIN_ALLOW_THREADS
+    result = PQexec(cnxn->pgconn, szSQL);
+    Py_END_ALLOW_THREADS
+
+    if (result == 0)
+        return 0;
+
+    switch (PQresultStatus(result)) {
+    case PGRES_COPY_OUT:
+        // This is what we are expecting.
+        break;
+
+    case PGRES_BAD_RESPONSE:
+    case PGRES_NONFATAL_ERROR:
+    case PGRES_FATAL_ERROR:
+        return SetResultError(result.Detach());
+
+    default:
+        return PyErr_Format(Error, "Result was not PGRES_COPY_OUT: %d", (int)PQresultStatus(result));
+    }
+
+
+    for (;;) {
+        int cb = 0;
+        char* buffer;
+        Py_BEGIN_ALLOW_THREADS
+        cb = PQgetCopyData(cnxn->pgconn, &buffer, 0);
+        Py_END_ALLOW_THREADS
+
+        if (cb == -2) {
+            return SetResultError(result.Detach());
+        }
+
+        if (cb == -1) {
+            // The copy is complete.
+            break;
+        }
+
+        // We have a buffer of byte data.  We have the length (`cb`), but the libpq docs say
+        // that the string is also zero terminated, so we're going to try not calling 'write'.
+
+        int err = PyFile_WriteString(buffer, dest);
+
+        PQfreemem(buffer);
+        if (err) {
+            return 0;
+        }
+    }
+
+    // After a copy, you have to get another result to know if it was successful.
+
+    ResultHolder final_result;
+    ExecStatusType status = PGRES_COMMAND_OK;
+    Py_BEGIN_ALLOW_THREADS
+    final_result = PQgetResult(cnxn->pgconn);
+    status = PQresultStatus(final_result);
+    Py_END_ALLOW_THREADS
+
+    if (status != PGRES_COMMAND_OK) {
+      // SetResultError will take ownership of `result`.
+      return SetResultError(final_result.Detach());
+    }
+
+    const char* sz = PQcmdTuples(final_result);
+    if (sz == 0 || *sz == 0)
+      Py_RETURN_NONE;
+    return PyLong_FromLong(atoi(sz));
+}
+
+
 const char* doc_copy_to_csv =
     "Connection.copy_to_csv(table, dest, header=0, delimiter=',', quote='\"')\n"
     "\n"
@@ -1631,6 +1738,7 @@ static struct PyMethodDef Connection_methods[] =
     { "script",  Connection_script,  METH_VARARGS, doc_script },
     { "copy_from", (PyCFunction) Connection_copy_from, METH_VARARGS | METH_KEYWORDS, doc_copy_from },
     { "copy_from_csv", (PyCFunction) Connection_copy_from_csv, METH_VARARGS | METH_KEYWORDS, doc_copy_from_csv },
+    { "copy_to", (PyCFunction) Connection_copy_to, METH_VARARGS, doc_copy_to },
     { "copy_to_csv", (PyCFunction) Connection_copy_to_csv, METH_VARARGS | METH_KEYWORDS, doc_copy_to_csv},
     { "begin",    Connection_begin,   METH_NOARGS, doc_begin },
     { "commit",   Connection_commit,   METH_NOARGS, doc_commit },

{pglib-5.11.0 → pglib-5.12.0}/test/test_sync.py

@@ -176,6 +176,39 @@ def test_copytocsv(cnxn):
             assert rows == [['a', 'b'], ['1', 'one'], ['2', 'two'], ['3', 'three']]
 
 
+def test_copy_to(cnxn):
+    # Test copy_to with a full command, including a subquery.
+    cnxn.execute("create table t1(a int, b text)")
+    cnxn.execute("insert into t1 values (1, 'one'), (2, 'two'), (3, 'three')")
+
+    with tempfile.NamedTemporaryFile(mode='w', encoding='utf8') as tf:
+        # Use a subquery to only copy rows where a > 1
+        count = cnxn.copy_to('copy (select * from t1 where a > 1 order by a) to stdout csv header', tf)
+        assert count == 2
+
+        tf.flush()
+        with open(tf.name, mode='r', encoding='utf8') as fd:
+            reader = csv.reader(fd)
+            rows = list(reader)
+            assert rows == [['a', 'b'], ['2', 'two'], ['3', 'three']]
+
+
+def test_copy_to_table(cnxn):
+    # Test copy_to with a simple table copy (no subquery).
+    cnxn.execute("create table t1(a int, b text)")
+    cnxn.execute("insert into t1 values (1, 'one'), (2, 'two')")
+
+    with tempfile.NamedTemporaryFile(mode='w', encoding='utf8') as tf:
+        count = cnxn.copy_to('copy t1 to stdout csv', tf)
+        assert count == 2
+
+        tf.flush()
+        with open(tf.name, mode='r', encoding='utf8') as fd:
+            reader = csv.reader(fd)
+            rows = list(reader)
+            assert rows == [['1', 'one'], ['2', 'two']]
+
+
 #
 # copy from
 #