[MongoEngine#2685] implement a thread-safe switch db context manager

Author: Jesse McLaughlin

mongoengine/connection.py

@@ -1,8 +1,10 @@
 import warnings
+from threading import local
 
 from pymongo import MongoClient, ReadPreference, uri_parser
 from pymongo.database import _check_name
 
+from mongoengine.errors import DatabaseAliasError
 from mongoengine.pymongo_support import PYMONGO_VERSION
 
 __all__ = [
@@ -15,6 +17,8 @@
     "get_connection",
     "get_db",
     "register_connection",
+    "set_local_db_alias",
+    "del_local_db_alias"
 ]
 
 
@@ -26,6 +30,8 @@
 _connection_settings = {}
 _connections = {}
 _dbs = {}
+_local = local()
+_local.db_alias = {}
 
 READ_PREFERENCE = ReadPreference.PRIMARY
 
@@ -372,7 +378,29 @@ def _clean_settings(settings_dict):
             return _connections[db_alias]
 
 
+def set_local_db_alias(local_alias, alias=DEFAULT_CONNECTION_NAME):
+    if not alias or not local_alias:
+        raise DatabaseAliasError(f"db alias and local_alias cannot be empty")
+
+    if alias not in _local.db_alias:
+        _local.db_alias[alias] = local_alias
+    else:
+        raise DatabaseAliasError(f"local db alias already set: {alias}")
+
+
+def del_local_db_alias(alias):
+    if not alias:
+        raise DatabaseAliasError(f"db alias cannot be empty")
+    if alias in _local.db_alias:
+        del _local.db_alias[alias]
+    else:
+        raise DatabaseAliasError(f"local db alias not set: {alias}")
+
+
 def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
+    if alias in _local.db_alias:
+        alias = _local.db_alias[alias]
+
     if reconnect:
         disconnect(alias)
 

mongoengine/context_managers.py

@@ -4,10 +4,11 @@
 from pymongo.write_concern import WriteConcern
 
 from mongoengine.common import _import_class
-from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db
+from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db, set_local_db_alias, del_local_db_alias
 from mongoengine.pymongo_support import count_documents
 
 __all__ = (
+    "switch_db_local",
     "switch_db",
     "switch_collection",
     "no_dereference",
@@ -18,6 +19,36 @@
 )
 
 
+class switch_db_local:
+    """switch_db_local alias context manager.
+
+    Switches a db alias in a thread-safe way.
+
+    Example ::
+        register_connection('testdb-1', 'mongoenginetest1')
+        register_connection('testdb-2', 'mongoenginetest2')
+
+        class Group(Document):
+            name = StringField()
+
+        # The following two calls to save() could be run concurrently
+        with switch_db_local('testdb-1'):
+            Group(name='test').save()
+        with switch_db_local('testdb-2'):
+            Group(name='test').save()
+    """
+
+    def __init__(self, local_alias, alias=DEFAULT_CONNECTION_NAME):
+        self.local_alias = local_alias
+        self.alias = alias
+
+    def __enter__(self):
+        set_local_db_alias(self.local_alias, self.alias)
+
+    def __exit__(self, t, value, traceback):
+        del_local_db_alias(self.alias)
+
+
 class switch_db:
     """switch_db alias context manager.
 

mongoengine/errors.py

@@ -1,6 +1,7 @@
 from collections import defaultdict
 
 __all__ = (
+    "DatabaseAliasError",
     "NotRegistered",
     "InvalidDocumentError",
     "LookUpError",
@@ -21,6 +22,10 @@ class MongoEngineException(Exception):
     pass
 
 
+class DatabaseAliasError(MongoEngineException):
+    pass
+
+
 class NotRegistered(MongoEngineException):
     pass