New class to implement "DB switcheable" documents in MongoEngine.

docs/db_switcheable.rst

@@ -0,0 +1,60 @@
+DBSwitchDocument class 
+======================
+
+DBSwitchDocument class allows documents to point to a new database in an easy
+way. By default, all document objects point to a specific database after
+connection. You can use *context manager* to switch a document from one
+database to another, but there is no way to do this for all documents at once
+in MongoEngine.
+
+This class provides that feadure. All the documents derived from
+DBSwitchDocument, will be switched to a new database every time you call the
+dbconnect() function.
+
+Example:
+
+    class MyDoc(DBSwitchDocument, Document):
+        field1 = StringField()
+        field2 = StringField()
+        field3 = StringField()
+
+    dbconnect("mydb1")
+
+    doc = MyDoc()
+    do_something(doc)
+
+    dbconnect("mydb2")
+
+    doc = MyDoc()
+    do_something(doc)
+
+In this example, all the documents derived from DBSwitchDocument will use the
+the *mydb1* database after dbconnect() first call. The second call redirects
+all the operatitions to *mydb2* database. You can switch anytime you want in the
+code to work with diferent DBs.
+
+Shared DB documents
+-------------------
+
+There is a special configuration in the DBSwitchDocument class that allows
+documents to live in a shared database. This documents won't switch from one DB
+to another on dbconnect() calls.
+
+Example:
+
+    class MySharedDoc(DBSwitchDocument, Document):
+        field1 = StringField()
+        field2 = StringField()
+        field3 = StringField()
+
+        shared_db = "my_shared_db"
+
+Then call dbconnect("my_shared_db") to establish the connection to the shared
+DB. Any other call to dbconnect() won't modify where this document points to.
+
+
+Authors
+-------
+
+- Diego Woitasen <[email protected]>
+

documents/db_switch.py

@@ -0,0 +1,94 @@
+
+import threading
+from collections import defaultdict
+
+from mongoengine import *
+from mongoengine.connection import get_db
+
+
+def dbconnect(dbname, prefix='', **kwargs):
+    connect(prefix + dbname, dbname, **kwargs)
+    DBSwitchDocument.switch_active_db(dbname)
+
+def current_db():
+    return DBSwitchDocument.active_db
+
+
+class DBSwitchDocument(object):
+    """DBSwitchDocument class, a special document class that allows DB switch 
+    for all its derived classes on the fly.
+
+    Example:
+
+    class MyDoc(DBSwitchDocument, Document):
+        field1 = StringField()
+        field2 = StringField()
+        field3 = StringField()
+
+    dbconnect("mydb1")
+
+    doc = MyDoc()
+    do_something(doc)
+
+    dbconnect("mydb2")
+
+    doc = MyDoc()
+    do_something(doc)
+
+    Another example, a Document that lives in a shared DB:
+
+    class MyDoc(DBSwitchDocument, Document):
+        field1 = StringField()
+        field2 = StringField()
+        field3 = StringField()
+
+        shared_db = "my_shared_db"
+
+    All the instances of this class will operate over the DB `my_shared_db`.
+
+    """
+
+    _collections = defaultdict(lambda: None)
+    active_db = None
+
+    #Set this to the name of the share DB in classes that live there
+    shared_db = None
+
+    @classmethod
+    def switch_active_db(cls, dbname):
+        cls.active_db = dbname
+
+    @classmethod
+    def _get_db(cls):
+        if cls.shared_db is None:
+            return get_db(cls.active_db)
+
+        return get_db(cls.shared_db)
+
+    @classmethod
+    def _build_key(cls, dbname, collection_name):
+        #XXX: check if we really need this per thread thing!!
+        thread_id = threading.currentThread().ident
+        key = '%s-%s-%s' % (thread_id, dbname, collection_name)
+        return key
+
+    @classmethod
+    def _get_collection(cls):
+        """Custom MongoEngine function to support global DB switch thread-safe.
+
+        """
+        collection_name = cls._get_collection_name()
+        if cls.shared_db is None:
+            key = cls._build_key(cls.active_db, collection_name)
+        else:
+            key = cls._build_key(cls.shared_db, collection_name)
+
+        if cls._collections[key] is None:
+            db = cls._get_db()
+            cls._collection = db[collection_name]
+            cls._collections[key] = db[collection_name]
+            if cls._meta.get('auto_create_index', True):
+                cls.ensure_indexes()
+
+        return cls._collections[key]
+

tests/documents/test_dbswitch.py

@@ -0,0 +1,54 @@
+
+import unittest
+import random
+import string
+
+from mongoengine import *
+from documents.db_switch import *
+
+class MyDoc(DBSwitchDocument, Document):
+    field1 = StringField()
+
+class MySharedDoc(DBSwitchDocument, Document):
+    field1 = StringField()
+
+    shared_db = "shared_test"
+
+def random_str(n=10):
+    return ''.join(random.choice(string.ascii_uppercase + string.digits) \
+        for x in range(n))
+
+class ConnectionTest(unittest.TestCase):
+    def test_document_switch(self):
+        """Test document class DB switching."""
+
+        dbconnect("dbswitch1")
+        doc1_fields = dict(field1=random_str())
+        doc1 = MyDoc(**doc1_fields)
+        doc1.save()
+
+        dbconnect("dbswitch2")
+        doc2_fields = dict(field1=random_str())
+        doc2 = MyDoc(**doc2_fields)
+        doc2.save()
+
+        dbconnect("dbswitch1")
+        doc1bis = MyDoc.objects(id=doc1.id).first()
+        self.assertEqual(doc1, doc1bis)
+        self.assertEqual(doc1.field1, doc1bis.field1)
+
+        dbconnect("dbswitch2")
+        doc2bis = MyDoc.objects(id=doc2.id).first()
+        self.assertEqual(doc2, doc2bis)
+        self.assertEqual(doc2.field1, doc2bis.field1)
+
+    def test_document_shared_db(self):
+        """Test a document in a shared db. This document musn't switch."""
+        dbconnect("shared_test")
+        shared_doc = MySharedDoc(field1=random_str())
+        shared_doc.save()
+
+        dbconnect("other_db_test")
+        shared_docbis = MySharedDoc.objects(id=shared_doc.id).first()
+        self.assertEqual(shared_doc, shared_docbis)
+