pglib 5.7.1__tar.gz → 5.8.1__tar.gz

{pglib-5.7.1/pglib.egg-info → pglib-5.8.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pglib
-Version: 5.7.1
+Version: 5.8.1
 Summary: A PostgreSQL interface
 Home-page: https://gitlab.com/mkleehammer/pglib
 Maintainer: Michael Kleehammer

{pglib-5.7.1 → pglib-5.8.1/pglib.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pglib
-Version: 5.7.1
+Version: 5.8.1
 Summary: A PostgreSQL interface
 Home-page: https://gitlab.com/mkleehammer/pglib
 Maintainer: Michael Kleehammer

{pglib-5.7.1 → pglib-5.8.1}/setup.py

@@ -130,7 +130,7 @@ def _get_settings():
 
 setup(
     name='pglib',
-    version='5.7.1',
+    version='5.8.1',
     description='A PostgreSQL interface',
     long_description=long_description,
     maintainer='Michael Kleehammer',

{pglib-5.7.1 → pglib-5.8.1}/src/connection.cpp

@@ -176,6 +176,7 @@ static PGresult* internal_execute(PyObject* self, PyObject* args)
     return result;
 }
 
+
 static const char doc_script[] = "Connection.script(sql) --> None\n\n"
     "Executes a script which can contain multiple statements separated by semicolons.";
 
@@ -204,6 +205,7 @@ static PyObject* Connection_script(PyObject* self, PyObject* args)
     }
 }
 
+
 const char* doc_copy_from =
     "Connection.copy_from(command, source) --> int\n"
     "\n"
@@ -234,7 +236,7 @@ static PyObject* Connection_copy_from(PyObject* self, PyObject* args)
     // an object with a read method (e.g. file).
     const char* buffer = 0;
     Py_ssize_t buffer_size = 0;
-    PyObject* read_method = 0;
+    Object read_method;
 
     if (PyUnicode_Check(source))
     {
@@ -246,7 +248,7 @@ static PyObject* Connection_copy_from(PyObject* self, PyObject* args)
     {
         if (!PyObject_HasAttrString(source, "read"))
             return PyErr_Format(Error, "CSV source must be a string or file-like object.");
-        read_method = PyObject_GetAttrString(source, "read");
+        read_method.Attach(PyObject_GetAttrString(source, "read"));
     }
 
     Connection* cnxn = CastConnection(self, REQUIRE_OPEN);
@@ -354,6 +356,140 @@ static PyObject* Connection_copy_from(PyObject* self, PyObject* args)
 }
 
 
+const char* doc_copy_to_csv =
+    "Connection.copy_to_csv(table, dest, header=0, delimiter=',', quote='\"')\n"
+    "\n"
+    "Execute a COPY TO command and return the number of records copied.\n"
+    "\n"
+    "table\n"
+    "  The table to copy from.\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"
+    "header\n"
+    "  If non-zero, a CSV header will be written.\n";
+
+
+static PyObject* Connection_copy_to_csv(PyObject* self, PyObject* args, PyObject* kwargs)
+{
+    // This is not nearly as efficient as I'd like since newer Python versions no longer give
+    // us access to underlying file objects.  We have to write strings through a write method
+    // since there are io layers involved.
+    //
+    // For maximum performance, we should probably offer an option where we open the file given
+    // a filename.  We can either check the parameter type here or we could make a separate
+    // method with "file" in the name like copy_to_file.
+
+    static const char* kwlist[] = {"table", "dest", "header", "delimiter", "quote", 0};
+
+    PyObject* table;
+    PyObject* dest;
+    int header = 0;
+    char* szDelimiter = 0;
+    char* szQuote = 0;
+
+    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "UO|pzz", (char**)kwlist, &table, &dest, &header,
+                                     &szDelimiter, &szQuote)) {
+        return 0;
+    }
+
+    Connection* cnxn = CastConnection(self, REQUIRE_OPEN);
+    if (!cnxn)
+        return 0;
+
+    if (!PyObject_HasAttrString(dest, "write"))
+        return PyErr_Format(Error, "CSV destination must be a file-like object.");
+    Object write_method(PyObject_GetAttrString(dest, "write"));
+
+    char header_token[] = "header";
+    if (header == 0) {
+        header_token[0] = 0;
+    }
+
+    const char* pszDelimiter = szDelimiter ? szDelimiter : ",";
+    const char* pszQuote = szQuote ? szQuote : "\"";
+
+    Object sql(PyUnicode_FromFormat("copy %U to stdout with csv %s delimiter '%s' quote '%s'",
+                                    table, header_token, pszDelimiter, pszQuote));
+
+    const char* szSQL = PyUnicode_AsUTF8(sql);
+    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_IN: %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, 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);
+
+        //  while (cb > 0) {
+        //      PyObject* res = PyObject_CallObject(write_method)
+        //  }
+
+        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_from_csv =
     "Connection.copy_from_csv(table, source, header=0) --> int\n"
     "\n"
@@ -390,15 +526,16 @@ static PyObject* Connection_copy_from_csv(PyObject* self, PyObject* args, PyObje
 
     const char* pszDelimiter = szDelimiter ? szDelimiter : ",";
     const char* pszQuote = szQuote ? szQuote : "\"";
-    PyObject* sql = PyUnicode_FromFormat("copy %U from stdin with csv %s delimiter '%s' quote '%s'",
-                                         table, header_token, pszDelimiter, pszQuote);
+    Object sql(PyUnicode_FromFormat("copy %U from stdin with csv %s delimiter '%s' quote '%s'",
+                                    table, header_token, pszDelimiter, pszQuote));
 
     // If source is a string (Unicode), store the UTF-encoded value in buffer. If a byte
     // object, store directly in buffer.  Otherwise, buffer will be zero and `source` must be
     // an object with a read method (e.g. file).
     const char* buffer = 0;
     Py_ssize_t buffer_size = 0;
-    PyObject* read_method = 0;
+    Object read_method;
+    //  PyObject* read_method = 0;
 
     if (PyUnicode_Check(source))
     {
@@ -410,7 +547,7 @@ static PyObject* Connection_copy_from_csv(PyObject* self, PyObject* args, PyObje
     {
         if (!PyObject_HasAttrString(source, "read"))
             return PyErr_Format(Error, "CSV source must be a string or file-like object.");
-        read_method = PyObject_GetAttrString(source, "read");
+        read_method.Attach(PyObject_GetAttrString(source, "read"));
     }
 
     Connection* cnxn = CastConnection(self, REQUIRE_OPEN);
@@ -1462,6 +1599,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_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 },
     { "rollback", Connection_rollback,   METH_NOARGS, doc_rollback },

{pglib-5.7.1 → pglib-5.8.1}/src/pglib.h

@@ -135,7 +135,7 @@ struct Object
 {
     PyObject* p;
 
-    // Borrows the reference (takes ownership without adding a new referencing).
+    // Borrows the reference (takes ownership without adding a new reference).
     Object(PyObject* _p = 0) { p = _p; }
 
     // If it still has the pointer, it dereferences it.

{pglib-5.7.1 → pglib-5.8.1}/test/test_sync.py

@@ -3,11 +3,11 @@
 
 # pylint: disable=missing-function-docstring,redefined-outer-name,unidiomatic-typecheck
 
-import sys, threading, gzip, uuid, locale
+import sys, threading, gzip, uuid, locale, tempfile, csv
 from time import sleep
 from os.path import join, dirname, exists
 from decimal import Decimal
-from datetime import date, time, datetime, timedelta, timezone
+from datetime import date, time, datetime, timedelta
 from pathlib import Path
 import pytest
 
@@ -32,7 +32,7 @@ def cnxn():
     cnxn = pglib.connect(CONNINFO)
     for i in range(3):
         try:
-            cnxn.execute("drop table t%d" % i)
+            cnxn.execute(f"drop table t{i}")
         except:                 # noqa
             pass
     yield cnxn
@@ -44,9 +44,9 @@ def _test_strtype(cnxn, sqltype, value, resulttype=None, colsize=None):
     assert colsize is None or (value is None or colsize >= len(value))
 
     if colsize:
-        sql = "create table t1(s %s(%s))" % (sqltype, colsize)
+        sql = f"create table t1(s {sqltype}({colsize}))"
     else:
-        sql = "create table t1(s %s)" % sqltype
+        sql = f"create table t1(s {sqltype})"
 
     if resulttype is None:
         resulttype = type(value)
@@ -148,7 +148,37 @@ def test_script(cnxn):
 
 
 #
-# copy
+# copy to
+#
+
+def test_copytocsv(cnxn):
+    # Write a table to a CSV file.
+    #
+    # It isn't as efficient as I'd like since we can no longer get ahold of the file
+    # descriptor.  I think we'll want to make this accept a string filename for a more
+    # efficient write.  For now, make sure the output file is opened in text mode.
+
+    cnxn.execute("create table t1(a text, b text)")
+    cnxn.execute("insert into t1 values ('1', 'one'), ('2', 'two'), ('3', 'three')")
+
+    with tempfile.NamedTemporaryFile(mode='w', encoding='utf8') as tf:
+        print(tf.name)
+        count = cnxn.copy_to_csv('t1', tf, header=1)
+        assert count == 3
+
+        tf.seek(0)
+        with open(tf.name, mode='r', encoding='utf8') as fd:
+            reader = csv.reader(fd)
+            rows = []
+            for row in reader:
+                print('ROW=', row)
+                rows.append(row)
+
+            assert rows == [['a', 'b'], ['1', 'one'], ['2', 'two'], ['3', 'three']]
+
+
+#
+# copy from
 #
 
 def _datapath(filename):
@@ -159,7 +189,8 @@ def _datapath(filename):
 
 def test_copyfromcsv_csv(cnxn):
     cnxn.execute("create table t1(a int, b varchar(20))")
-    count = cnxn.copy_from_csv("t1", open(_datapath('test-noheader.csv')))
+    with open(_datapath('test-noheader.csv')) as fd:
+        count = cnxn.copy_from_csv("t1", fd)
     assert count == 2
     assert cnxn.fetchval("select count(*) from t1") == 2
     row = cnxn.fetchrow("select a,b from t1 where a=2")
@@ -169,7 +200,8 @@ def test_copyfromcsv_csv(cnxn):
 
 def test_copyfromcsv_csv_header(cnxn):
     cnxn.execute("create table t1(a int, b varchar(20))")
-    count = cnxn.copy_from_csv("t1", open(_datapath('test-header.csv')), header=True)
+    with open(_datapath('test-header.csv')) as fd:
+        count = cnxn.copy_from_csv("t1", fd, header=True)
     assert count == 2
     assert cnxn.fetchval("select count(*) from t1") == 2
     row = cnxn.fetchrow("select a,b from t1 where a=2")
@@ -184,13 +216,15 @@ def test_copyfromcsv_csv_error(cnxn):
     # We'll make the second column too small.
     cnxn.execute("create table t1(a int, b varchar(1) not null)")
     with pytest.raises(pglib.Error):
-        cnxn.copy_from_csv("t1", open(_datapath('test-noheader.csv')))
+        with open(_datapath('test-noheader.csv')) as fd:
+            cnxn.copy_from_csv("t1", fd)
 
 
 def test_copyfromcsv_csv_gzip(cnxn):
     # I don't remember why this test is here.  We're feeding it unzipped data.
     cnxn.execute("create table t1(a int, b varchar(20))")
-    cnxn.copy_from_csv("t1", gzip.open(_datapath('test-header.csv.gz')), header=True)
+    with gzip.open(_datapath('test-header.csv.gz')) as fd:
+        cnxn.copy_from_csv("t1", fd, header=True)
     assert cnxn.fetchval("select count(*) from t1") == 2
     row = cnxn.fetchrow("select a,b from t1 where a=2")
     assert row.a == 2
@@ -1063,7 +1097,7 @@ def test_closed_error(cnxn):
         cnxn.rollback()
 
 
-def test_count():
+def test_connection_count():
     before = pglib.connection_count()
     cnxn = pglib.connect(CONNINFO)
     assert pglib.connection_count() == (before + 1)
@@ -1073,6 +1107,18 @@ def test_count():
     assert pglib.connection_count() == before
 
 
+def test_count(cnxn):
+    "Ensure delete statements return affected row count."
+    cnxn.execute(
+        """
+        select generate_series(1, 3) as id
+          into t1
+        """)
+
+    count = cnxn.execute("delete from t1 where id in (1, 2, 3)")
+    assert count == 3
+
+
 def test_hstore(cnxn):
     cnxn.execute("create extension if not exists hstore")
     row = cnxn.fetchrow("select oid, typname from pg_type where typname='hstore'")