Re: NOTIFY payload patch

"A.M." <agentm-/iWpWt6iY7eAP89PaY/[email protected]> Fri, 27 Aug 2010 12:25:16 -0400
Newsgroups gmane.comp.python.db.psycopg.devel
Message-ID <[email protected]>
On Aug 26, 2010, at 10:33 AM, A.M. wrote:

> 
> On Aug 26, 2010, at 5:31 AM, Federico Di Gregorio wrote:
> 
>> On 08/25/10 21:39, A.M. wrote:
>> [snip]
>>> Note that this will break code which unpacks the tuple in an
>>> assignment:
>>> 
>>> (pid,name)  = dbconn.notifies.pop()
>>> 
>>> must become
>>> 
>>> (pid,name,payload) = dbconn.notifies.pop()
>>> 
>>> but as long as pop() doesn't return a Notification object, I don't
>>> see a way around this.
>>> 
>>> This patch was tested successfully against 8.4 and 9.0b4 libpq
>>> against a 9.0b4 server with an without payloads passed.
>> 
>> Thank you very much for this patch. I don't like breakages so, probably,
>> the best thing would be to substitute the tuple with an instance of a
>> Notify class (as you suggest) that has both 3 attributes (for pid, name
>> and payload) and can be indexed like a tuple composed of two elements.
> 
> That's a good plan- I will work on this.

So here is my draft attempt at a Notification class which behaves like a backwards-compatible tuple. This is my first work with the Python C API, so a thorough code review would be much appreciated- perhaps I missed some locking or memory management. With this patch, one can do this:

>>> notification = dbconn.notifies.pop()
>>> print notification
(123,'test')
>>> print notification.name,notification.backendpid,notification.payload
test 123 payload

Please let me know how this can be improved. Thanks!

Cheers,
M

diff --git a/lib/notification.py b/lib/notification.py
new file mode 100644
index 0000000..c0814d9
--- /dev/null
+++ b/lib/notification.py
@@ -0,0 +1,12 @@
+"""
+psycopg2.notification - represents a PostgreSQL NOTIFY with support for 9.0 payloads
+- see psycopg/connection_int.c to see where this is generated
+"""
+class Notification(tuple):
+    def __new__(cls,backendpid,name,payload=''):
+        return super(Notification,cls).__new__(cls,[backendpid,name])
+
+    def __init__(self,backendpid,name,payload=''):
+        self.backendpid = backendpid
+        self.name = name
+        self.payload = payload
diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index 902fdbb..483712d 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -136,16 +136,49 @@ conn_notifies_process(connectionObject *self)
     PGnotify *pgn;
 
     while ((pgn = PQnotifies(self->pgconn)) != NULL) {
-        PyObject *notify;
-
-        Dprintf("conn_notifies_process: got NOTIFY from pid %d, msg = %s",
-                (int) pgn->be_pid, pgn->relname);
-
-        notify = PyTuple_New(2);
-        PyTuple_SET_ITEM(notify, 0, PyInt_FromLong((long)pgn->be_pid));
-        PyTuple_SET_ITEM(notify, 1, PyString_FromString(pgn->relname));
-        PyList_Append(self->notifies, notify);
-        Py_DECREF(notify);
+        PyObject *notificationModule;
+       PyObject *notificationType;
+       PyObject *newNotification;
+       PyObject *constructorArguments;
+
+        Dprintf("conn_notifies_process: got NOTIFY from pid %d, msg = %s, payload = %s",
+                (int) pgn->be_pid, pgn->relname, pgn->extra);
+
+       /* get python type by string psycopg2.Notification 
+          
+          from psycopg2.notification import Notification
+          newNotification = Notification(pid,name,payload)
+        */
+       notificationModule = PyImport_ImportModule("psycopg2.notification");
+       if(notificationModule == NULL) {
+         PyErr_SetString(PyExc_ImportError,"No module named psycopg2.notification- NOTIFY will not be reported");
+         return;
+       }
+       
+       notificationType = PyObject_GetAttrString(notificationModule,"Notification");
+       if(notificationType == NULL) {
+         PyErr_SetString(PyExc_AttributeError,"psycopg2.Notification class not found- NOTIFY will not be reported");
+         Py_DECREF(notificationModule);
+         return;
+       }
+        
+       constructorArguments = PyTuple_New(3);
+        PyTuple_SET_ITEM(constructorArguments, 0, PyInt_FromLong((long)pgn->be_pid));
+        PyTuple_SET_ITEM(constructorArguments, 1, PyString_FromString(pgn->relname));
+       PyTuple_SET_ITEM(constructorArguments, 2, PyString_FromString(pgn->extra));
+
+       newNotification = PyObject_CallObject(notificationType,constructorArguments);
+       if(newNotification == NULL) {
+         PyErr_SetString(PyErr_Occurred(),"Failed to create new psycopg2.Notification");
+       }
+       else { /* successfully created new Notification */
+         PyList_Append(self->notifies, newNotification);
+         Py_DECREF(newNotification);
+       }
+         
+        Py_DECREF(constructorArguments);
+       Py_DECREF(notificationType);
+       Py_DECREF(notificationModule);
         PQfreemem(pgn);
     }
 }