src/pyams_utils/zodb.py
changeset 10 e87103c49c8a
parent 1 3f89629b9e54
child 17 c5a19fddac05
--- a/src/pyams_utils/zodb.py	Thu Mar 05 16:44:00 2015 +0100
+++ b/src/pyams_utils/zodb.py	Thu Mar 05 16:44:16 2015 +0100
@@ -17,11 +17,26 @@
 
 # import interfaces
 from persistent.interfaces import IPersistent
+from pyams_utils.interfaces.site import IOptionalUtility
+from pyams_utils.interfaces.zeo import IZEOConnection
 from transaction.interfaces import ITransactionManager
 from ZODB.interfaces import IConnection
+from zope.annotation.interfaces import IAttributeAnnotatable
+from zope.lifecycleevent.interfaces import IObjectAddedEvent, IObjectRemovedEvent
+from zope.schema.interfaces import IVocabularyFactory
 
 # import packages
+from persistent import Persistent
 from pyams_utils.adapter import adapter_config
+from pyramid.events import subscriber
+from ZEO import ClientStorage
+from ZODB import DB
+from zope.componentvocabulary.vocabulary import UtilityVocabulary
+from zope.container.contained import Contained
+from zope.interface import implementer, provider
+from zope.schema import getFieldNames
+from zope.schema.fieldproperty import FieldProperty
+from zope.schema.vocabulary import getVocabularyRegistry
 
 
 @adapter_config(context=IPersistent, provides=IConnection)
@@ -46,10 +61,98 @@
 @adapter_config(context=IPersistent, provides=ITransactionManager)
 def get_transaction_manager(obj):
     conn = IConnection(obj)  # typically this will be
-                             # zope.app.keyreference.persistent.connectionOfPersistent
+                             # zope.keyreference.persistent.connectionOfPersistent
     try:
         return conn.transaction_manager
     except AttributeError:
         return conn._txn_mgr
         # or else we give up; who knows.  transaction_manager is the more
         # recent spelling.
+
+
+@implementer(IZEOConnection)
+class ZEOConnection(object):
+    """ZEO connection object"""
+
+    _storage = None
+    _db = None
+    _connection = None
+
+    name = FieldProperty(IZEOConnection['name'])
+    server_name = FieldProperty(IZEOConnection['server_name'])
+    server_port = FieldProperty(IZEOConnection['server_port'])
+    storage = FieldProperty(IZEOConnection['storage'])
+    username = FieldProperty(IZEOConnection['username'])
+    password = FieldProperty(IZEOConnection['password'])
+    server_realm = FieldProperty(IZEOConnection['server_realm'])
+    blob_dir = FieldProperty(IZEOConnection['blob_dir'])
+    shared_blob_dir = FieldProperty(IZEOConnection['shared_blob_dir'])
+
+    def get_settings(self):
+        result = {}
+        for name in getFieldNames(IZEOConnection):
+            result[name] = getattr(self, name)
+        return result
+
+    def update(self, settings):
+        names = getFieldNames(IZEOConnection)
+        for key, value in settings.items():
+            if key in names:
+                setattr(self, key, value)
+
+    def get_connection(self, wait=False, get_storage=False):
+        storage = ClientStorage.ClientStorage((self.server_name, self.server_port),
+                                              storage=self.storage,
+                                              username=self.username or '',
+                                              password=self.password or '',
+                                              realm=self.server_realm,
+                                              blob_dir=self.blob_dir,
+                                              shared_blob_dir=self.shared_blob_dir,
+                                              wait=wait)
+        db = DB(storage)
+        return (storage, db) if get_storage else db
+
+    @property
+    def connection(self):
+        return self._connection
+
+    # Context manager methods
+    def __enter__(self):
+        self._storage, self._db = self.get_connection(get_storage=True)
+        self._connection = self._db.open()
+        return self._connection.root()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self._connection is not None:
+            self._connection.close()
+        if self._storage is not None:
+            self._storage.close()
+
+
+@implementer(IOptionalUtility, IAttributeAnnotatable)
+class ZEOConnectionUtility(ZEOConnection, Persistent, Contained):
+    """Persistent ZEO connection utility"""
+
+
+@subscriber(IObjectAddedEvent, context_selector=IZEOConnection)
+def handle_added_connection(event):
+    """Register new ZEO connection when added"""
+    manager = event.newParent
+    manager.registerUtility(event.object, IZEOConnection, name=event.object.name)
+
+
+@subscriber(IObjectRemovedEvent, context_selector=IZEOConnection)
+def handle_removed_connection(event):
+    """Un-register ZEO connection when deleted"""
+    manager = event.oldParent
+    manager.unregisterUtility(event.object, IZEOConnection, name=event.object.name)
+
+
+@provider(IVocabularyFactory)
+class ZEOConnectionVocabulary(UtilityVocabulary):
+    """ZEO connections vocabulary"""
+
+    interface = IZEOConnection
+    nameOnly = True
+
+getVocabularyRegistry().register('PyAMS ZEO connections', ZEOConnectionVocabulary)