Added multidb support

No change required to upgrade to multiple databases. Aliases are used
to describe the database and these can be manually registered or fall
through to a default alias using connect.

Made get_connection and get_db first class members of the connection class.
Old style _get_connection and _get_db still supported.

Refs: #84 #87 #93 #215

mongoengine/connection.py

@@ -1,82 +1,106 @@
 from pymongo import Connection
-import multiprocessing
-import threading
-
-__all__ = ['ConnectionError', 'connect']
 
 
-_connection_defaults = {
-    'host': 'localhost',
-    'port': 27017,
-}
-_connection = {}
-_connection_settings = _connection_defaults.copy()
+__all__ = ['ConnectionError', 'connect', 'register_connection']
 
-_db_name = None
-_db_username = None
-_db_password = None
-_db = {}
+
+DEFAULT_CONNECTION_NAME = 'default'
 
 
 class ConnectionError(Exception):
     pass
 
 
-def _get_connection(reconnect=False):
-    """Handles the connection to the database
+_connection_settings = {}
+_connections = {}
+_dbs = {}
+
+
+def register_connection(alias, name, host='localhost', port=27017,
+                        is_slave=False, slaves=None, username=None,
+                        password=None):
+    """Add a connection.
+
+    :param alias: the name that will be used to refer to this connection
+        throughout MongoEngine
+    :param name: the name of the specific database to use
+    :param host: the host name of the :program:`mongod` instance to connect to
+    :param port: the port that the :program:`mongod` instance is running on
+    :param is_slave: whether the connection can act as a slave
+    :param slaves: a list of aliases of slave connections; each of these must
+        be a registered connection that has :attr:`is_slave` set to ``True``
+    :param username: username to authenticate with
+    :param password: password to authenticate with
     """
-    global _connection
-    identity = get_identity()
+    global _connection_settings
+    _connection_settings[alias] = {
+        'name': name,
+        'host': host,
+        'port': port,
+        'is_slave': is_slave,
+        'slaves': slaves or [],
+        'username': username,
+        'password': password,
+    }
+
+
+def get_connection(alias=DEFAULT_CONNECTION_NAME):
+    global _connections
     # Connect to the database if not already connected
-    if _connection.get(identity) is None or reconnect:
+    if alias not in _connections:
+        if alias not in _connection_settings:
+            msg = 'Connection with alias "%s" has not been defined'
+            if alias == DEFAULT_CONNECTION_NAME:
+                msg = 'You have not defined a default connection'
+            raise ConnectionError(msg)
+        conn_settings = _connection_settings[alias].copy()
+
+        # Get all the slave connections
+        slaves = []
+        for slave_alias in conn_settings['slaves']:
+            slaves.append(get_connection(slave_alias))
+        conn_settings['slaves'] = slaves
+
         try:
-            _connection[identity] = Connection(**_connection_settings)
+            _connections[alias] = Connection(**conn_settings)
         except Exception, e:
-            raise ConnectionError("Cannot connect to the database:\n%s" % e)
-    return _connection[identity]
+            raise e
+            raise ConnectionError('Cannot connect to database %s' % alias)
+    return _connections[alias]
 
-def _get_db(reconnect=False):
-    """Handles database connections and authentication based on the current
-    identity
+
+def get_db(alias=DEFAULT_CONNECTION_NAME):
+    global _dbs
+    if alias not in _dbs:
+        conn = get_connection(alias)
+        conn_settings = _connection_settings[alias]
+        _dbs[alias] = conn[conn_settings['name']]
+
+        # Authenticate if necessary
+        if conn_settings['username'] and conn_settings['password']:
+            _dbs[alias].authenticate(conn_settings['username'],
+                                     conn_settings['password'])
+    return _dbs[alias]
+
+
+def connect(db, alias=DEFAULT_CONNECTION_NAME, **kwargs):
+    """Connect to the database specified by the 'db' argument.
+
+    Connection settings may be provided here as well if the database is not
+    running on the default port on localhost. If authentication is needed,
+    provide username and password arguments as well.
+
+    Multiple databases are supported by using aliases.  Provide a separate
+    `alias` to connect to a different instance of :program:`mongod`.
+
+    .. versionchanged:: 0.6 - added multiple database support.
     """
-    global _db, _connection
-    identity = get_identity()
-    # Connect if not already connected
-    if _connection.get(identity) is None or reconnect:
-        _connection[identity] = _get_connection(reconnect=reconnect)
+    global _connections
+    if alias not in _connections:
+        register_connection(alias, db, **kwargs)
 
-    if _db.get(identity) is None or reconnect:
-        # _db_name will be None if the user hasn't called connect()
-        if _db_name is None:
-            raise ConnectionError('Not connected to the database')
-
-        # Get DB from current connection and authenticate if necessary
-        _db[identity] = _connection[identity][_db_name]
-        if _db_username and _db_password:
-            _db[identity].authenticate(_db_username, _db_password)
-
-    return _db[identity]
-
-def get_identity():
-    """Creates an identity key based on the current process and thread
-    identity.
-    """
-    identity = multiprocessing.current_process()._identity
-    identity = 0 if not identity else identity[0]
-
-    identity = (identity, threading.current_thread().ident)
-    return identity
-
-def connect(db, username=None, password=None, **kwargs):
-    """Connect to the database specified by the 'db' argument. Connection
-    settings may be provided here as well if the database is not running on
-    the default port on localhost. If authentication is needed, provide
-    username and password arguments as well.
-    """
-    global _connection_settings, _db_name, _db_username, _db_password, _db
-    _connection_settings = dict(_connection_defaults, **kwargs)
-    _db_name = db
-    _db_username = username
-    _db_password = password
-    return _get_db(reconnect=True)
+    return get_connection(alias)
 
+# Support old naming convention
+_get_connection = get_connection
+_get_db = get_db