This commit is contained in:
Harry Marr 2010-09-19 19:53:12 +01:00
commit ee72845701

View File

@ -1,5 +1,5 @@
from pymongo import Connection from pymongo import Connection
import multiprocessing
__all__ = ['ConnectionError', 'connect'] __all__ = ['ConnectionError', 'connect']
@ -8,12 +8,12 @@ _connection_settings = {
'host': 'localhost', 'host': 'localhost',
'port': 27017, 'port': 27017,
} }
_connection = None _connection = {}
_db_name = None _db_name = None
_db_username = None _db_username = None
_db_password = None _db_password = None
_db = None _db = {}
class ConnectionError(Exception): class ConnectionError(Exception):
@ -22,31 +22,38 @@ class ConnectionError(Exception):
def _get_connection(): def _get_connection():
global _connection global _connection
identity = get_identity()
# Connect to the database if not already connected # Connect to the database if not already connected
if _connection is None: if _connection.get(identity) is None:
try: try:
_connection = Connection(**_connection_settings) _connection[identity] = Connection(**_connection_settings)
except: except:
raise ConnectionError('Cannot connect to the database') raise ConnectionError('Cannot connect to the database')
return _connection return _connection[identity]
def _get_db(): def _get_db():
global _db, _connection global _db, _connection
identity = get_identity()
# Connect if not already connected # Connect if not already connected
if _connection is None: if _connection.get(identity) is None:
_connection = _get_connection() _connection[identity] = _get_connection()
if _db is None: if _db.get(identity) is None:
# _db_name will be None if the user hasn't called connect() # _db_name will be None if the user hasn't called connect()
if _db_name is None: if _db_name is None:
raise ConnectionError('Not connected to the database') raise ConnectionError('Not connected to the database')
# Get DB from current connection and authenticate if necessary # Get DB from current connection and authenticate if necessary
_db = _connection[_db_name] _db[identity] = _connection[identity][_db_name]
if _db_username and _db_password: if _db_username and _db_password:
_db.authenticate(_db_username, _db_password) _db[identity].authenticate(_db_username, _db_password)
return _db return _db[identity]
def get_identity():
identity = multiprocessing.current_process()._identity
identity = 0 if not identity else identity[0]
return identity
def connect(db, username=None, password=None, **kwargs): def connect(db, username=None, password=None, **kwargs):
"""Connect to the database specified by the 'db' argument. Connection """Connect to the database specified by the 'db' argument. Connection