Large objects landing..

ChangeLog

@@ -1,5 +1,7 @@
 2006-09-01  Federico Di Gregorio  <fog@initd.org>
 
+	* Implemented large objects support.
+
 	* psycopg/connection_int.c: removed increment of self->mark,
 	now it is done directly in pqpath.c to make sure even the
 	large object support gets it.

examples/lobject.py

@@ -28,13 +28,52 @@ print "Opening connection using dns:", DSN
 conn = psycopg2.connect(DSN)
 print "Encoding for this connection is", conn.encoding
 
-# this will create a large object with a new random oid
+# this will create a large object with a new random oid, we'll
+# use it to make some basic tests about read/write and seek.
 lobj = conn.lobject()
-print "lobject oid =", lobj.oid
+loid = lobj.oid
+print "Created a new large object with oid", loid
 
-# this will create a large object with the given oid
-lobj = conn.lobject(0, 0, 666)
-print "lobject oid =", lobj.oid
+print "Manually importing some binary data into the object:"
+data = open("somehackers.jpg").read()
+len = lobj.write(data)
+print "  imported", len, "bytes of data"
 
-lobj = conn.lobject(0, 0, 666)
+conn.commit()
 
+print "Trying to (re)open large object with oid", loid
+lobj = conn.lobject(loid)
+print "Manually exporting the data from the lobject:"
+data1 = lobj.read()
+len = lobj.tell()
+lobj.seek(0, 0)
+data2 = lobj.read()
+if data1 != data2:
+    print "ERROR: read after seek returned different data"
+open("somehackers_lobject1.jpg", 'wb').write(data1)
+print "  written", len, "bytes of data to somehackers_lobject1.jpg"
+
+lobj.unlink()
+print "Large object with oid", loid, "removed"
+
+conn.commit()
+
+# now we try to use the import and export functions to do the same
+lobj = conn.lobject(0, 'n', 0, "somehackers.jpg")
+loid = lobj.oid
+print "Imported a new large object with oid", loid
+
+conn.commit()
+
+print "Trying to (re)open large object with oid", loid
+lobj = conn.lobject(loid, 'n')
+print "Using export() to export the data from the large object:"
+lobj.export("somehackers_lobject2.jpg")
+print "  exported large object to somehackers_lobject2.jpg"
+
+lobj.unlink()
+print "Large object with oid", loid, "removed"
+
+conn.commit()
+
+print "\nNow try to load the new images, to check it worked!"

lib/extensions.py

@@ -61,14 +61,6 @@ STATUS_ASYNC    = 4
 # This is a usefull mnemonic to check if the connection is in a transaction
 STATUS_IN_TRANSACTION = STATUS_BEGIN
 
-"""Large object flags"""
-LOBJECT_INV_READ  = 0
-LOBJECT_INV_WRITE = 1
-
-LOBJECT_SEEK_SET = 0
-LOBJECT_SEEK_CUR = 1
-LOBJECT_SEEK_END = 2
-
 
 def register_adapter(typ, callable):
     """Register 'callable' as an ISQLQuote adapter for type 'typ'."""

psycopg/connection_type.c

@@ -209,15 +209,16 @@ psyco_conn_set_client_encoding(connectionObject *self, PyObject *args)
 static PyObject *
 psyco_conn_lobject(connectionObject *self, PyObject *args, PyObject *keywds)
 {
-    int oid=0, new_oid=0, mode=0;
-    char *new_file = NULL;
+    Oid oid=InvalidOid, new_oid=InvalidOid;
+    char *smode = NULL, *new_file = NULL;
+    int mode=0;
     PyObject *obj, *factory = NULL;
 
     static char *kwlist[] = {"oid", "mode", "new_oid", "new_file",
                              "cursor_factory", NULL};
     
-    if (!PyArg_ParseTupleAndKeywords(args, keywds, "|iiisO", kwlist,
-                                     &oid, &mode, &new_oid, &new_file, 
+    if (!PyArg_ParseTupleAndKeywords(args, keywds, "|izisO", kwlist,
+                                     &oid, &smode, &new_oid, &new_file, 
 				     &factory)) {
         return NULL;
     }
@@ -225,11 +226,30 @@ psyco_conn_lobject(connectionObject *self, PyObject *args, PyObject *keywds)
     EXC_IF_CONN_CLOSED(self);
 
     Dprintf("psyco_conn_lobject: new lobject for connection at %p", self);
-    Dprintf("psyco_conn_lobject:     parameters: oid = %d, mode = %d",
-            oid, mode);
+    Dprintf("psyco_conn_lobject:     parameters: oid = %d, mode = %s",
+            oid, smode);
     Dprintf("psyco_conn_lobject:     parameters: new_oid = %d, new_file = %s",
             new_oid, new_file);
     
+    /* build a mode number out of the mode string: right now we only accept
+       'r', 'w' and 'rw' (but note that 'w' implies 'rw' because PostgreSQL
+       backend does that. */
+    if (smode) {
+        if (strncmp("rw", smode, 2) == 0)
+            mode = INV_READ+INV_WRITE;
+        else if (smode[0] == 'r')
+           mode = INV_READ;
+        else if (smode[0] == 'w')
+           mode = INV_WRITE;
+	else if (smode[0] == 'n')
+	   mode = -1;
+        else {
+            PyErr_SetString(PyExc_TypeError,
+                "mode should be one of 'r', 'w' or 'rw'");
+	    return NULL;
+	}
+    }
+
     if (factory == NULL) factory = (PyObject *)&lobjectType;
     if (new_file)
         obj = PyObject_CallFunction(factory, "Oiiis", 

psycopg/cursor_type.c

@@ -1088,13 +1088,13 @@ psyco_curs_copy_from(cursorObject *self, PyObject *args, PyObject *kwargs)
     
     if (columns != NULL && columns != Py_None) {
         PyObject* collistiter = PyObject_GetIter(columns);
-        if (collistiter == NULL) {
-            return NULL;
-        }
         PyObject* col;
         int collistlen = 2;
         int colitemlen;
         char* colname;
+        if (collistiter == NULL) {
+            return NULL;
+        }
         strcpy(columnlist, " (");
         while ((col = PyIter_Next(collistiter)) != NULL) {
             if (!PyString_Check(col)) {

psycopg/lobject.h

@@ -51,8 +51,17 @@ typedef struct {
 
 extern int lobject_open(lobjectObject *self, connectionObject *conn,
                          Oid oid, int mode, Oid new_oid, char *new_file);
+extern int lobject_unlink(lobjectObject *self);
+extern int lobject_export(lobjectObject *self, char *filename);
+
+extern size_t lobject_read(lobjectObject *self, char *buf, size_t len);
+extern size_t lobject_write(lobjectObject *self, char *buf, size_t len);
+extern int lobject_seek(lobjectObject *self, int pos, int whence);
+extern int lobject_tell(lobjectObject *self);
+extern void lobject_close(lobjectObject *self);
 
 /* exception-raising macros */
+
 #define EXC_IF_LOBJ_CLOSED(self) \
 if ((self)->closed || ((self)->conn && (self)->conn->closed)) { \
     PyErr_SetString(InterfaceError, "lobject already closed");  \

psycopg/lobject_int.c

@@ -63,17 +63,25 @@ lobject_open(lobjectObject *self, connectionObject *conn,
         if (mode == 0) mode = INV_READ;
     }
 
-    /* if the oid is a real one we try to open with the given mode */
+    /* if the oid is a real one we try to open with the given mode,
+       unless the mode is -1, meaning "don't open!" */
+    if (mode != -1) {
         self->fd = lo_open(self->conn->pgconn, self->oid, mode);
         Dprintf("lobject_open: large object opened with fd = %d",
             self->fd);
+    }
+    else {
+        /* this is necessary to make sure no function that needs and
+	   fd is called on unopened lobjects */
+        self->closed = 1;
+    }
 
  end:
     pthread_mutex_unlock(&(self->conn->lock));
     Py_END_ALLOW_THREADS;
 
     /* here we check for errors before returning 0 */
-    if (self->fd == -1 || self->oid == InvalidOid) {
+    if ((self->fd == -1 && mode != -1) || self->oid == InvalidOid) {
         pq_raise(conn, NULL, NULL, NULL);
         return -1;
     }

psycopg/lobject_type.c

@@ -49,14 +49,156 @@ psyco_lobj_close(lobjectObject *self, PyObject *args)
 {
     if (!PyArg_ParseTuple(args, "")) return NULL;
    
-    EXC_IF_LOBJ_CLOSED(self);
-    EXC_IF_LOBJ_LEVEL0(self);
-    EXC_IF_LOBJ_UNMARKED(self);
-
+    /* file-like objects can be closed multiple times and remember that
+       closing the current transaction is equivalent to close all the 
+       opened large objects */
+    if (!self->closed
+        && self->conn->isolation_level > 0
+	&& self->conn->mark == self->mark)
+    {
         self->closed = 1;
         lobject_close(self);
 
         Dprintf("psyco_lobj_close: lobject at %p closed", self);
+    }
+
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+/* write method - write data to the lobject */
+
+#define psyco_lobj_write_doc \
+"write(str) -- Write a string to the large object."
+
+static PyObject *
+psyco_lobj_write(lobjectObject *self, PyObject *args)
+{
+    int len, res=0;
+    char *buffer;
+
+    if (!PyArg_ParseTuple(args, "s#", &buffer, &len)) return NULL;
+
+    EXC_IF_LOBJ_CLOSED(self);
+    EXC_IF_LOBJ_LEVEL0(self);
+    EXC_IF_LOBJ_UNMARKED(self);
+
+    if ((res = lobject_write(self, buffer, len)) < 0) return NULL;
+
+    return PyInt_FromLong((long)res);
+}
+
+/* read method - read data from the lobject */
+
+#define psyco_lobj_read_doc \
+"read(size=-1) -- Read at most size bytes or to the end of the large object."
+
+static PyObject *
+psyco_lobj_read(lobjectObject *self, PyObject *args)
+{
+    int where, end, size = -1;
+    char *buffer;
+
+    if (!PyArg_ParseTuple(args, "|i", &size)) return NULL;
+
+    EXC_IF_LOBJ_CLOSED(self);
+    EXC_IF_LOBJ_LEVEL0(self);
+    EXC_IF_LOBJ_UNMARKED(self);
+
+    if (size < 0) {
+        if ((where = lobject_tell(self)) < 0) return NULL;
+        if ((end = lobject_seek(self, 0, SEEK_END)) < 0) return NULL;
+        if (lobject_seek(self, where, SEEK_SET) < 0) return NULL;
+        size = end - where;
+    }
+
+    if ((buffer = PyMem_Malloc(size)) == NULL) return NULL;
+    if ((size = lobject_read(self, buffer, size)) < 0) {
+        PyMem_Free(buffer);
+        return NULL;
+    }
+
+    return PyString_FromStringAndSize(buffer, size);
+}
+
+/* seek method - seek in the lobject */
+
+#define psyco_lobj_seek_doc \
+"seek(offset, whence=0) -- Set the lobject's current position."
+
+static PyObject *
+psyco_lobj_seek(lobjectObject *self, PyObject *args)
+{
+    int offset, whence=0;
+    int pos=0;
+
+    if (!PyArg_ParseTuple(args, "i|i", &offset, &whence))
+    	return NULL;
+    
+    EXC_IF_LOBJ_CLOSED(self);
+    EXC_IF_LOBJ_LEVEL0(self);
+    EXC_IF_LOBJ_UNMARKED(self);
+ 
+    if ((pos = lobject_seek(self, pos, whence)) < 0)
+    	return NULL;
+    
+    return PyInt_FromLong((long)pos);
+}
+
+/* tell method - tell current position in the lobject */
+
+#define psyco_lobj_tell_doc \
+"tell() -- Return the lobject's current position."
+
+static PyObject *
+psyco_lobj_tell(lobjectObject *self, PyObject *args)
+{
+    int pos;
+
+    if (!PyArg_ParseTuple(args, "")) return NULL;
+    
+    EXC_IF_LOBJ_CLOSED(self);
+    EXC_IF_LOBJ_LEVEL0(self);
+    EXC_IF_LOBJ_UNMARKED(self);
+    
+    if ((pos = lobject_tell(self)) < 0)
+    	return NULL;
+    
+    return PyInt_FromLong((long)pos);
+}
+
+/* unlink method - unlink (destroy) the lobject */
+
+#define psyco_lobj_unlink_doc \
+"unlink() -- Close and then remove the lobject."
+
+static PyObject *
+psyco_lobj_unlink(lobjectObject *self, PyObject *args)
+{
+    if (!PyArg_ParseTuple(args, "")) return NULL;
+    
+    if (lobject_unlink(self) < 0)
+    	return NULL;
+    
+    Py_INCREF(Py_None);
+    return Py_None;
+}
+
+/* export method - export lobject's content to given file */
+
+#define psyco_lobj_export_doc \
+"export(filename) -- Export large object to given file."
+
+static PyObject *
+psyco_lobj_export(lobjectObject *self, PyObject *args)
+{
+    char *filename;
+
+    if (!PyArg_ParseTuple(args, "s", &filename))
+    	return NULL;
+    
+    if (lobject_export(self, filename) < 0)
+    	return NULL;
     
     Py_INCREF(Py_None);
     return Py_None;  
@@ -68,8 +210,20 @@ psyco_lobj_close(lobjectObject *self, PyObject *args)
 /* object method list */
 
 static struct PyMethodDef lobjectObject_methods[] = {
+    {"read", (PyCFunction)psyco_lobj_read,
+     METH_VARARGS, psyco_lobj_read_doc},
+    {"write", (PyCFunction)psyco_lobj_write,
+     METH_VARARGS, psyco_lobj_write_doc},
+    {"seek", (PyCFunction)psyco_lobj_seek,
+     METH_VARARGS, psyco_lobj_seek_doc},
+    {"tell", (PyCFunction)psyco_lobj_tell,
+     METH_VARARGS, psyco_lobj_tell_doc},
     {"close", (PyCFunction)psyco_lobj_close,
      METH_VARARGS, psyco_lobj_close_doc},
+    {"unlink",(PyCFunction)psyco_lobj_unlink,
+     METH_VARARGS, psyco_lobj_unlink_doc},
+    {"export",(PyCFunction)psyco_lobj_export,
+     METH_VARARGS, psyco_lobj_export_doc},
     {NULL}
 };