Merge lp:~ccx/txpostgres/notify into lp:~wulczer/txpostgres/trunk

[Jan Pobrislo]

Support for NOTIFY/LISTEN events as supported by psycopg2 by connection.notifies.

Unmerged revisions

17. By Jan Pobrislo on 2010-10-13

add test case for notify events

16. By Jan Pobrislo on 2010-10-13

handle adding/removing reader and writer more elegantly

15. By Jan Pobrislo on 2010-10-13

Add support for NOTIFY/LISTEN events.

=== modified file 'test/test_txpostgres.py'
--- test/test_txpostgres.py	2010-07-02 00:55:29 +0000
+++ test/test_txpostgres.py	2010-10-14 00:11:49 +0000
@@ -305,6 +305,21 @@
         d = c.execute("drop table simple")
         return d.addCallback(lambda _: self.conn.close())
 
+class TxPostgresNotifyObserverTestCase(_SimpleDBSetupMixin, Psycopg2TestCase):
+
+    def test_notifyObserver(self):
+        notify_d = defer.Deferred()
+        self.conn.addNotifyObserver(
+                lambda pid, name: notify_d.callback((pid,name))
+                )
+        c = self.conn.cursor()
+        d = c.execute("LISTEN FOO"
+                ).addCallback(
+                lambda c: c.execute("NOTIFY FOO")
+                ).addCallback(
+                lambda c: notify_d
+                )
+        return d
 
 class TxPostgresManualQueryTestCase(_SimpleDBSetupMixin, Psycopg2TestCase):
 
 
=== modified file 'txpostgres/txpostgres.py'
--- txpostgres/txpostgres.py	2010-07-02 00:55:29 +0000
+++ txpostgres/txpostgres.py	2010-10-14 00:11:49 +0000
@@ -54,6 +54,10 @@
     prefix = "pollable"
     _pollingD = None
 
+    def __init__(self):
+        self._is_reading = False
+        self._is_writing = False
+
     def pollable(self):
         """
         Return the pollable object. Subclasses should override this.
@@ -62,6 +66,26 @@
         """
         raise NotImplementedError()
 
+    def addReader(self):
+        if not self._is_reading:
+            self.reactor.addReader(self)
+            self._is_reading = True
+
+    def addWriter(self):
+        if not self._is_writing:
+            self.reactor.addWriter(self)
+            self._is_writing = True
+
+    def removeReader(self):
+        if self._is_reading:
+            self.reactor.removeReader(self)
+            self._is_reading = False
+
+    def removeWriter(self):
+        if self._is_writing:
+            self.reactor.removeWriter(self)
+            self._is_writing = False
+
     def poll(self):
         """
         Start polling the wrapped pollable.
@@ -85,9 +109,9 @@
             d, self._pollingD = self._pollingD, None
             d.callback(self)
         elif self._pollingState == psycopg2.extensions.POLL_WRITE:
-            self.reactor.addWriter(self)
+            self.addWriter()
         elif self._pollingState == psycopg2.extensions.POLL_READ:
-            self.reactor.addReader(self)
+            self.addReader()
         else:
             d, self._pollingD = self._pollingD, None
             d.errback(UnexpectedPollResult())
@@ -95,11 +119,11 @@
         return ret
 
     def doRead(self):
-        self.reactor.removeReader(self)
+        self.removeReader()
         self.poll()
 
     def doWrite(self):
-        self.reactor.removeWriter(self)
+        self.removeWriter()
         self.poll()
 
     def logPrefix(self):
@@ -132,6 +156,8 @@
     """
 
     def __init__(self, cursor, connection):
+        super(Cursor, self).__init__()
+
         self.reactor = connection.reactor
         self.prefix = "cursor"
 
@@ -237,6 +263,8 @@
     cursorFactory = Cursor
 
     def __init__(self, reactor=None):
+        super(Connection, self).__init__()
+
         if not reactor:
             from twisted.internet import reactor
         self.reactor = reactor
@@ -245,10 +273,39 @@
         # this lock will be used to prevent concurrent query execution
         self.lock = defer.DeferredLock()
         self._connection = None
+        self._notify_observers = set()
 
     def pollable(self):
         return self._connection
 
+    def addNotifyObserver(self, callback):
+        """
+        Add callable to be called back when NOTIFY event are delivered.
+
+        @type callback: Callable that takes pid, name as parameters.
+        """
+        self._notify_observers.add(callback)
+        if self.pollable():
+            if self not in self.reactor.getReaders():
+                self.reactor.addReader(self)
+
+    def removeNotifyObserver(self, callback):
+        """
+        Remove callable from NOTIFY event are delivery.
+
+        @type callback: Callable previously registered by addNotifyObserver.
+        """
+        self._notify_observers.remove(callback)
+
+    def doRead(self):
+        super(Connection, self).doRead()
+        if self._notify_observers and self.pollable():
+            self.addReader()
+            while self.notifies:
+                pid, name = self.notifies.pop(0)
+                for observer in self._notify_observers:
+                    observer(pid, name)
+
     def connect(self, *args, **kwargs):
         """
         Connect to the database.
@@ -276,6 +333,8 @@
         """
         _connection, self._connection = self._connection, None
         _connection.close()
+        self.removeReader()
+        self.removeWriter()
 
     def cursor(self):
         """
@@ -575,3 +634,5 @@
         c = self.connections.pop()
         d = c.runInteraction(interaction, *args, **kwargs)
         return d.addBoth(self._putBackAndPassthrough, c)
+
+# vim: ft=python et sts=4 ts=4