src/pyams_utils/zodb.py
changeset 124 53dc81f933ed
parent 100 119b9c2f3022
child 130 9dd00059a845
--- a/src/pyams_utils/zodb.py	Thu Jan 11 16:56:18 2018 +0100
+++ b/src/pyams_utils/zodb.py	Thu Jan 11 16:57:23 2018 +0100
@@ -29,7 +29,7 @@
 from persistent import Persistent
 from pyams_utils.adapter import adapter_config
 from pyams_utils.property import DocFieldProperty
-from pyams_utils.registry import get_utilities_for
+from pyams_utils.registry import get_utilities_for, get_global_registry
 from pyams_utils.vocabulary import vocabulary_config
 from pyramid.events import subscriber
 from pyramid_zodbconn import get_uris, db_from_uri
@@ -42,7 +42,7 @@
 
 
 @adapter_config(context=IPersistent, provides=IConnection)
-def get_connection(obj):
+def persistent_connection(obj):
     """An adapter which gets a ZODB connection from a persistent object
 
     We are assuming the object has a parent if it has been created in
@@ -61,7 +61,7 @@
 # IPersistent adapters copied from zc.twist package
 # also register this for adapting from IConnection
 @adapter_config(context=IPersistent, provides=ITransactionManager)
-def get_transaction_manager(obj):
+def persistent_transaction_manager(obj):
     conn = IConnection(obj)  # typically this will be
                              # zope.keyreference.persistent.connectionOfPersistent
     try:
@@ -73,15 +73,22 @@
 
 
 @adapter_config(context=object, provides=ICacheKeyValue)
-def persistent_key_adapter(obj):
+def object_key_adapter(obj):
     return '{0!r}'.format(obj)
 
 
+#
+# ZEO connection management
+#
+
 @implementer(IZEOConnection)
 class ZEOConnection(object):
     """ZEO connection object
 
     This object can be used to store all settings to be able to open a ZEO connection.
+    Note that this class is required only for tasks specifically targeting a ZEO database connection (like a ZEO
+    packer scheduler task); for generic ZODB operations, just use a :class:`ZODBConnection` class defined through
+    Pyramid's configuration file.
 
     Note that a ZEO connection object is a context manager, so you can use it like this:
 
@@ -201,8 +208,86 @@
         super(ZEOConnectionVocabulary, self).__init__(terms)
 
 
-def get_connection_from_settings(settings):
+def get_connection_from_settings(settings=None):
     """Load connection matching registry settings"""
+    if settings is None:
+        settings = get_global_registry().settings
     for name, uri in get_uris(settings):
         db = db_from_uri(uri, name, {})
         return db.open()
+
+
+class ZODBConnection(object):
+    """ZODB connection wrapper
+
+    Connections are extracted from Pyramid's settings file in *zodbconn.uri* entries.
+
+    Note that a ZODB connection object is a context manager, so you can use it like this:
+
+    .. code-block:: python
+
+        from pyams_utils.zodb import ZODBConnection
+
+        def my_method(zodb_name):
+            zodb_connection = ZODBConnection(zodb_name)
+            with zodb_connection as root:
+                # *root* is then the ZODB root object
+                # do whatever you want with ZODB connection,
+                # which is closed automatically
+    """
+
+    def __init__(self, name='', settings=None):
+        self.name = name or ''
+        if not settings:
+            settings = get_global_registry().settings
+        self.settings = settings
+
+    _connection = None
+    _db = None
+    _storage = None
+
+    @property
+    def connection(self):
+        return self._connection
+
+    @property
+    def db(self):
+        return self._db
+
+    @property
+    def storage(self):
+        return self._storage
+
+    def get_connection(self):
+        """Load named connection matching registry settings"""
+        for name, uri in get_uris(self.settings):
+            if name == self.name:
+                db = db_from_uri(uri, name, {})
+                connection = self._connection = db.open()
+                self._db = connection.db()
+                self._storage = self.db.storage
+                return connection
+
+    def close(self):
+        self._connection.close()
+        self._db.close()
+        self._storage.close()
+
+    # Context manager methods
+
+    def __enter__(self):
+        connection = self.get_connection()
+        return connection.root()
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+
+
+@vocabulary_config(name='PyAMS ZODB connections')
+class ZODBConnectionVocabulary(SimpleVocabulary):
+    """ZODB connections vocabulary"""
+
+    def __init__(self, context=None):
+        settings = get_global_registry().settings
+        terms = [SimpleTerm(name, title=name) for name, uri in get_uris(settings)]
+        super(ZODBConnectionVocabulary, self).__init__(terms)