large object truncating support

From fcb10d5db0ee051c95d793a503536dc0805f4e14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Urba=C5=84ski?= <[email protected]>
Date: Mon, 29 Mar 2010 00:42:56 +0200
Subject: [PATCH] Support large objects truncating.

The lobject.truncate(len=0) method will be available if psycopg2 has
been built against libpq from 8.3 or later (which is when the lobject
truncating support has been introduced).
---
 doc/src/extensions.rst |   13 +++++++++
 psycopg/lobject.h      |    1 +
 psycopg/lobject_int.c  |   30 +++++++++++++++++++++
 psycopg/lobject_type.c |   30 +++++++++++++++++++++
 tests/test_lobject.py  |   66 +++++++++++++++++++++++++++++++++++++++++++++++-
 5 files changed, 139 insertions(+), 1 deletions(-)

diff --git a/doc/src/extensions.rst b/doc/src/extensions.rst
index c4dfd66..436bb7a 100644
--- a/doc/src/extensions.rst
+++ b/doc/src/extensions.rst
@@ -78,6 +78,19 @@ functionalities defined by the |DBAPI|_.
 
         Return the lobject current position.
 
+    .. method:: truncate(len=0)
+
+        .. versionadded:: 2.0.15
+
+        Truncate the lobject to the given size.
+
+	The method will only be available if psycopg has been built against libpq
+	from PostgreSQL 8.3 or later and can only be used with PostgreSQL servers
+	running these versions. It uses the |lo_truncate|_ libpq function.
+
+        .. |lo_truncate| replace:: `!lo_truncate()`
+        .. _lo_truncate: http://www.postgresql.org/docs/8.4/static/lo-interfaces.html#AEN36420
+
     .. method:: close()
 
         Close the object.
diff --git a/psycopg/lobject.h b/psycopg/lobject.h
index d0be0a9..2d6715c 100644
--- a/psycopg/lobject.h
+++ b/psycopg/lobject.h
@@ -64,6 +64,7 @@ HIDDEN Py_ssize_t lobject_write(lobjectObject *self, const char *buf,
                                 size_t len);
 HIDDEN int lobject_seek(lobjectObject *self, int pos, int whence);
 HIDDEN int lobject_tell(lobjectObject *self);
+HIDDEN int lobject_truncate(lobjectObject *self, size_t len);
 HIDDEN int lobject_close(lobjectObject *self);
 
 #define lobject_is_closed(self) \
diff --git a/psycopg/lobject_int.c b/psycopg/lobject_int.c
index a66c72b..0b38325 100644
--- a/psycopg/lobject_int.c
+++ b/psycopg/lobject_int.c
@@ -333,6 +333,36 @@ lobject_export(lobjectObject *self, const char *filename)
     return retvalue;
 }
 
+#if PG_VERSION_HEX >= 0x080300
+
+int
+lobject_truncate(lobjectObject *self, size_t len)
+{
+    int retvalue;
+    PGresult *pgres = NULL;
+    char *error = NULL;
+
+    Dprintf("lobject_truncate: fd = %d, len = %d",
+            self->fd, len);
+
+    Py_BEGIN_ALLOW_THREADS;
+    pthread_mutex_lock(&(self->conn->lock));
+
+    retvalue = lo_truncate(self->conn->pgconn, self->fd, len);
+    Dprintf("lobject_truncate: result = %d", retvalue);
+    if (retvalue < 0)
+        collect_error(self->conn, &error);
+
+    pthread_mutex_unlock(&(self->conn->lock));
+    Py_END_ALLOW_THREADS;
+
+    if (retvalue < 0)
+        pq_complete_error(self->conn, &pgres, &error);
+    return retvalue;
+
+}
+
+#endif /* PG_VERSION_HEX >= 0x080300 */
 
 #endif
 
diff --git a/psycopg/lobject_type.c b/psycopg/lobject_type.c
index b227294..7fe16ab 100644
--- a/psycopg/lobject_type.c
+++ b/psycopg/lobject_type.c
@@ -227,6 +227,32 @@ psyco_lobj_get_closed(lobjectObject *self, void *closure)
     return closed;
 }
 
+#if PG_VERSION_HEX >= 0x080300
+
+#define psyco_lobj_truncate_doc \
+"truncate(len=0) -- Truncate large object to given size."
+
+static PyObject *
+psyco_lobj_truncate(lobjectObject *self, PyObject *args)
+{
+    int len = 0;
+
+    if (!PyArg_ParseTuple(args, "|i", &len))
+        return NULL;
+
+    EXC_IF_LOBJ_CLOSED(self);
+    EXC_IF_LOBJ_LEVEL0(self);
+    EXC_IF_LOBJ_UNMARKED(self);
+
+    if (lobject_truncate(self, len) < 0)
+        return NULL;
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+#endif /* PG_VERSION_HEX >= 0x080300 */
+
 
 /** the lobject object **/
 
@@ -247,6 +273,10 @@ static struct PyMethodDef lobjectObject_methods[] = {
      METH_VARARGS, psyco_lobj_unlink_doc},
     {"export",(PyCFunction)psyco_lobj_export,
      METH_VARARGS, psyco_lobj_export_doc},
+#if PG_VERSION_HEX >= 0x080300
+    {"truncate",(PyCFunction)psyco_lobj_truncate,
+     METH_VARARGS, psyco_lobj_truncate_doc},
+#endif /* PG_VERSION_HEX >= 0x080300 */
     {NULL}
 };
 
diff --git a/tests/test_lobject.py b/tests/test_lobject.py
index ecb1eec..f4a684f 100644
--- a/tests/test_lobject.py
+++ b/tests/test_lobject.py
@@ -3,6 +3,7 @@ import os
 import shutil
 import tempfile
 import unittest
+import warnings
 
 import psycopg2
 import psycopg2.extensions
@@ -261,8 +262,71 @@ class LargeObjectTests(unittest.TestCase):
         self.assertEqual(open(filename, "rb").read(), "some data")
 
 
+class LargeObjectTruncateTests(LargeObjectTests):
+
+    skip = None
+
+    def setUp(self):
+        LargeObjectTests.setUp(self)
+
+        if self.skip is None:
+            self.skip = False
+            if self.conn.server_version < 80300:
+                warnings.warn("Large object truncate tests skipped, "
+                              "the server does not support them")
+                self.skip = True
+
+            if not hasattr(psycopg2.extensions.lobject, 'truncate'):
+                warnings.warn("Large object truncate tests skipped, "
+                              "psycopg2 has been built against an old library")
+                self.skip = True
+
+    def test_truncate(self):
+        if self.skip:
+            return
+
+        lo = self.conn.lobject()
+        lo.write("some data")
+        lo.close()
+
+        lo = self.conn.lobject(lo.oid, "w")
+        lo.truncate(4)
+
+        # seek position unchanged
+        self.assertEqual(lo.tell(), 0)
+        # data truncated
+        self.assertEqual(lo.read(), "some")
+
+        lo.truncate(6)
+        lo.seek(0)
+        # large object extended with zeroes
+        self.assertEqual(lo.read(), "some\x00\x00")
+
+        lo.truncate()
+        lo.seek(0)
+        # large object empty
+        self.assertEqual(lo.read(), "")
+
+    def test_truncate_after_close(self):
+        if self.skip:
+            return
+
+        lo = self.conn.lobject()
+        lo.close()
+        self.assertRaises(psycopg2.InterfaceError, lo.truncate)
+
+    def test_truncate_after_commit(self):
+        if self.skip:
+            return
+
+        lo = self.conn.lobject()
+        self.lo_oid = lo.oid
+        self.conn.commit()
+
+        self.assertRaises(psycopg2.ProgrammingError, lo.truncate)
+
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)
 
 if __name__ == "__main__":
-    unittest.main()
\ No newline at end of file
+    unittest.main()
-- 
1.7.0