125 lines
3.8 KiB
Python
125 lines
3.8 KiB
Python
from pymongo import Connection
|
|
|
|
|
|
__all__ = ['ConnectionError', 'connect', 'register_connection',
|
|
'DEFAULT_CONNECTION_NAME']
|
|
|
|
|
|
DEFAULT_CONNECTION_NAME = 'default'
|
|
|
|
|
|
class ConnectionError(Exception):
|
|
pass
|
|
|
|
|
|
_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_settings
|
|
_connection_settings[alias] = {
|
|
'name': name,
|
|
'host': host,
|
|
'port': port,
|
|
'is_slave': is_slave,
|
|
'slaves': slaves or [],
|
|
'username': username,
|
|
'password': password,
|
|
}
|
|
|
|
|
|
def disconnect(alias=DEFAULT_CONNECTION_NAME):
|
|
global _connections
|
|
global _dbs
|
|
|
|
if alias in _connections:
|
|
get_connection(alias=alias).disconnect()
|
|
del _connections[alias]
|
|
if alias in _dbs:
|
|
del _dbs[alias]
|
|
|
|
|
|
def get_connection(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
|
|
global _connections
|
|
# Connect to the database if not already connected
|
|
if reconnect:
|
|
disconnect(alias)
|
|
|
|
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:
|
|
_connections[alias] = Connection(**conn_settings)
|
|
except Exception, e:
|
|
raise e
|
|
raise ConnectionError('Cannot connect to database %s' % alias)
|
|
return _connections[alias]
|
|
|
|
|
|
def get_db(alias=DEFAULT_CONNECTION_NAME, reconnect=False):
|
|
global _dbs
|
|
if reconnect:
|
|
disconnect(alias)
|
|
|
|
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 _connections
|
|
if alias not in _connections:
|
|
register_connection(alias, db, **kwargs)
|
|
|
|
return get_connection(alias)
|
|
|
|
# Support old naming convention
|
|
_get_connection = get_connection
|
|
_get_db = get_db
|