diff --git a/django_statsd/patches/db.py b/django_statsd/patches/db.py index fd6c8e1..3f8ff97 100644 --- a/django_statsd/patches/db.py +++ b/django_statsd/patches/db.py @@ -7,6 +7,7 @@ from django_statsd.clients import statsd def key(db, attr): return 'db.%s.%s.%s' % (db.client.executable_name, db.alias, attr) + def pre_django_1_6_cursorwrapper_getattr(self, attr): """ The CursorWrapper is a pretty small wrapper around the cursor. @@ -22,33 +23,39 @@ def pre_django_1_6_cursorwrapper_getattr(self, attr): return wrap(getattr(self.cursor, attr), key(self.db, attr)) return getattr(self.cursor, attr) + +def patched_execute(orig_execute, self, *args, **kwargs): + with statsd.timer(key(self.db, 'execute')): + return orig_execute(self, *args, **kwargs) + + +def patched_executemany(orig_executemany, self, *args, **kwargs): + with statsd.timer(key(self.db, 'executemany')): + return orig_executemany(self, *args, **kwargs) + + +def patched_callproc(orig_callproc, self, *args, **kwargs): + with statsd.timer(key(self.db, 'callproc')): + return orig_callproc(self, *args, **kwargs) + + def patch(): """ - 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 for version earlier than 1.6. + 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 for version + earlier than 1.6. """ - def execute(orig_execute, self, *args, **kwargs): - with statsd.timer(key(self.db, 'execute')): - return orig_execute(self, *args, **kwargs) - - def executemany(orig_executemany, self, *args, **kwargs): - with statsd.timer(key(self.db, 'executemany')): - return orig_executemany(self, *args, **kwargs) - - def callproc(orig_callproc, self, *args, **kwargs): - with statsd.timer(key(self.db, 'callproc')): - return orig_callproc(self, *args, **kwargs) - if django.VERSION > (1, 6): # In 1.6+ util.CursorDebugWrapper just makes calls to CursorWrapper # As such, we only need to instrument CursorWrapper. # Instrumenting both will result in duplicated metrics - patch_method(util.CursorWrapper, 'execute')(execute) - patch_method(util.CursorWrapper, 'executemany')(executemany) - patch_method(util.CursorWrapper, 'callproc')(callproc) + patch_method(util.CursorWrapper, 'execute')(patched_execute) + patch_method(util.CursorWrapper, 'executemany')(patched_executemany) + patch_method(util.CursorWrapper, 'callproc')(patched_callproc) else: util.CursorWrapper.__getattr__ = pre_django_1_6_cursorwrapper_getattr - patch_method(util.CursorDebugWrapper, 'execute')(execute) - patch_method(util.CursorDebugWrapper, 'executemany')(executemany) + patch_method(util.CursorDebugWrapper, 'execute')(patched_execute) + patch_method( + util.CursorDebugWrapper, 'executemany')(patched_executemany) diff --git a/django_statsd/tests.py b/django_statsd/tests.py index d6df3b4..8bf3ced 100644 --- a/django_statsd/tests.py +++ b/django_statsd/tests.py @@ -5,7 +5,9 @@ import sys from django.conf import settings from nose.exc import SkipTest from nose import tools as nose_tools +from unittest import skipUnless +from django import VERSION from django.core.urlresolvers import reverse from django.http import HttpResponse, HttpResponseForbidden from django.test import TestCase @@ -15,7 +17,7 @@ from django.utils import unittest import mock from nose.tools import eq_ -from django_statsd.clients import get_client +from django_statsd.clients import get_client, statsd from django_statsd.patches import utils from django_statsd import middleware @@ -464,3 +466,90 @@ class TestPatchMethod(TestCase): self.assertEqual(obj.badfn(1, d=2), ((1,), {'d': 2})) self.assertEqual(obj.badfn(1, d=2), ((1,), {'d': 2})) self.assertEqual(obj.badfn(1, 2, c=1, d=2), ((1, 2), {'c': 1, 'd': 2})) + + +class TestCursorWrapperPatching(TestCase): + + def test_patched_callproc_calls_timer(self): + from django_statsd.patches.db import patched_callproc + with mock.patch.object(statsd, 'timer') as timer: + db = mock.Mock(executable_name='name', alias='alias') + instance = mock.Mock(db=db) + patched_callproc(lambda *args, **kwargs: None, instance) + self.assertEqual(timer.call_count, 1) + + def test_patched_execute_calls_timer(self): + from django_statsd.patches.db import patched_execute + with mock.patch.object(statsd, 'timer') as timer: + db = mock.Mock(executable_name='name', alias='alias') + instance = mock.Mock(db=db) + patched_execute(lambda *args, **kwargs: None, instance) + self.assertEqual(timer.call_count, 1) + + def test_patched_executemany_calls_timer(self): + from django_statsd.patches.db import patched_executemany + with mock.patch.object(statsd, 'timer') as timer: + db = mock.Mock(executable_name='name', alias='alias') + instance = mock.Mock(db=db) + patched_executemany(lambda *args, **kwargs: None, instance) + self.assertEqual(timer.call_count, 1) + + @mock.patch( + 'django_statsd.patches.db.pre_django_1_6_cursorwrapper_getattr') + @mock.patch('django_statsd.patches.db.patched_executemany') + @mock.patch('django_statsd.patches.db.patched_execute') + @mock.patch('django.db.backends.util.CursorDebugWrapper') + @skipUnless(VERSION < (1, 6, 0), "CursorWrapper Patching for Django<1.6") + def test_cursorwrapper_patching(self, + CursorDebugWrapper, + execute, + executemany, + _getattr): + try: + from django.db.backends import util + + # We need to patch CursorWrapper like this because setting + # __getattr__ on Mock instances raises AttributeError. + class CursorWrapper(object): + pass + + _CursorWrapper = util.CursorWrapper + util.CursorWrapper = CursorWrapper + + from django_statsd.patches.db import patch + execute.__name__ = 'execute' + executemany.__name__ = 'executemany' + _getattr.__name__ = '_getattr' + execute.return_value = 'execute' + executemany.return_value = 'executemany' + _getattr.return_value = 'getattr' + patch() + + self.assertEqual(CursorDebugWrapper.execute(), 'execute') + self.assertEqual(CursorDebugWrapper.executemany(), 'executemany') + self.assertEqual(CursorWrapper.__getattr__(), 'getattr') + finally: + util.CursorWrapper = _CursorWrapper + + @mock.patch('django_statsd.patches.db.patched_callproc') + @mock.patch('django_statsd.patches.db.patched_executemany') + @mock.patch('django_statsd.patches.db.patched_execute') + @mock.patch('django.db.backends.util.CursorWrapper') + @skipUnless(VERSION >= (1, 6, 0), "CursorWrapper Patching for Django>=1.6") + def test_cursorwrapper_patching16(self, + CursorWrapper, + execute, + executemany, + callproc): + from django_statsd.patches.db import patch + execute.__name__ = 'execute' + executemany.__name__ = 'executemany' + callproc.__name__ = 'callproc' + execute.return_value = 'execute' + executemany.return_value = 'executemany' + callproc.return_value = 'callproc' + patch() + + self.assertEqual(CursorWrapper.execute(), 'execute') + self.assertEqual(CursorWrapper.executemany(), 'executemany') + self.assertEqual(CursorWrapper.callproc(), 'callproc')