diff --git a/django_statsd/patches/db.py b/django_statsd/patches/db.py index 4806d4c..82307ef 100644 --- a/django_statsd/patches/db.py +++ b/django_statsd/patches/db.py @@ -1,14 +1,18 @@ import django from django.db.backends import util - +from django_statsd.clients import statsd from django_statsd.patches.utils import wrap def key(db, attr): return 'db.%s.%s.%s' % (db.client.executable_name, db.alias, attr) - def __getattr__(self, attr): + """ + The CursorWrapper is a pretty small wrapper around the cursor. + If you are NOT in debug mode, this is the wrapper that's used. + Sadly if it's in debug mode, we get a different wrapper. + """ if django.VERSION < (1, 6) and self.db.is_managed(): # In Django 1.6 you can't put a connection in managed mode self.db.set_dirty() @@ -20,19 +24,27 @@ def __getattr__(self, attr): return getattr(self.cursor, attr) -def wrap_class(base): +def wrap_class(base, super_class=None): + """Returns a sub-class of the argument 'base' + The sub class has its execute and executemany methods will record query timings. + + If the execute and executemany methods you're overriding makes super() + calls, the argument super_class can be provided to not break things. + """ + class Wrapper(base): - def execute(self, *args, **kw): - return wrap(super(Wrapper, self).execute, - key(self.db, 'execute'))(*args, **kw) - def executemany(self, *args, **kw): - return wrap(super(Wrapper, self).executemany, - key(self.db, 'executemany'))(*args, **kw) + def execute(self, *args, **kwargs): + with statsd.timer(key(self.db, 'execute')): + return super(super_class, self).execute(*args, **kwargs) + def executemany(self, *args, **kwargs): + with statsd.timer(key(self.db, 'executemany')): + return super(super_class, self).execute(*args, **kwargs) + + super_class = super_class or Wrapper return Wrapper - def patch(): """ The CursorWrapper is a pretty small wrapper around the cursor. @@ -40,10 +52,10 @@ def patch(): Sadly if it's in debug mode, we get a different wrapper for version earlier than 1.6. """ - # So that it will work when DEBUG = True. - if django.VERSION < (1, 6): - util.CursorDebugWrapper = wrap_class(util.CursorDebugWrapper) + if django.VERSION > (1, 6): + util.CursorDebugWrapper = wrap_class(util.CursorDebugWrapper, + super_class=util.CursorDebugWrapper) + util.CursorWrapper = wrap_class(util.CursorWrapper) else: - util.CursorDebugWrapper.__getattr__ = __getattr__ - # So that it will work when DEBUG = False. - util.CursorWrapper.__getattr__ = __getattr__ + util.CursorDebugWrapper = wrap_class(util.CursorDebugWrapper) + util.CursorWrapper.__getattr__ = __getattr__