Commit 84143f12 by Jonathan Kamens

Be more careful about which attributes and keys get wrapped

Only wrap things that are callable, so that, e.g., properties (like
"db.name") will continue to work.

Also, __getitem__ should fetch item from self.conn, and __getattr__
should fetch attr from self.conn, and they shouldn't be mixed up.
parent 93d81f8f
...@@ -2,9 +2,15 @@ ...@@ -2,9 +2,15 @@
import time import time
import pymongo import pymongo
EXECUTABLE_MONGO_METHODS = set([typ for typ in dir(pymongo.collection.Collection) if not typ.startswith('_')]) def get_methods(obj):
EXECUTABLE_MONGO_METHODS.update(set([typ for typ in dir(pymongo.Connection) if not typ.startswith('_')])) attrs = (attr for attr in dir(obj))
EXECUTABLE_MONGO_METHODS.update(set([typ for typ in dir(pymongo) if not typ.startswith('_')])) attrs = (attr for attr in attrs if not attr.startswith('_'))
return set([attr for attr in attrs if hasattr(getattr(obj, attr), '__call__')])
EXECUTABLE_MONGO_METHODS = get_methods(pymongo.collection.Collection)
EXECUTABLE_MONGO_METHODS.update(get_methods(pymongo.Connection))
EXECUTABLE_MONGO_METHODS.update(get_methods(pymongo))
def safe_mongocall(call): def safe_mongocall(call):
""" Decorator for automatic handling of AutoReconnect-exceptions. """ Decorator for automatic handling of AutoReconnect-exceptions.
...@@ -53,20 +59,22 @@ class MongoProxy: ...@@ -53,20 +59,22 @@ class MongoProxy:
named "key". named "key".
""" """
attr = getattr(self.conn, key) item = self.conn[key]
if hasattr(attr, '__call__'): if hasattr(item, '__call__'):
return MongoProxy(attr) return MongoProxy(item)
return attr return item
def __getattr__(self, key): def __getattr__(self, key):
""" If key is the name of an executable method in the MongoDB connection, """ If key is the name of an executable method in the MongoDB connection,
for instance find or insert, wrap this method in Executable-class that for instance find or insert, wrap this method in Executable-class that
handles AutoReconnect-Exception. handles AutoReconnect-Exception.
""" """
if key in EXECUTABLE_MONGO_METHODS:
return Executable(getattr(self.conn, key)) attr = getattr(self.conn, key)
return self[key] if hasattr(attr, '__call__') and key in EXECUTABLE_MONGO_METHODS:
return Executable(attr)
return attr
def __call__(self, *args, **kwargs): def __call__(self, *args, **kwargs):
return self.conn(*args, **kwargs) return self.conn(*args, **kwargs)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment