From e448ac83af81cde86963efa85d568eed50e59889 Mon Sep 17 00:00:00 2001
From: Jonathan Slenders <jonathan@slenders.be>
Date: Wed, 3 Sep 2014 15:07:06 +0200
Subject: [PATCH] Log update/select/insert queries separately.

---
 django_statsd/patches/db.py | 21 +++++++++-------
 django_statsd/tests.py      | 50 ++++++++++++++++++++++++++-----------
 2 files changed, 47 insertions(+), 24 deletions(-)

diff --git a/django_statsd/patches/db.py b/django_statsd/patches/db.py
index 3f8ff97..df491af 100644
--- a/django_statsd/patches/db.py
+++ b/django_statsd/patches/db.py
@@ -24,19 +24,22 @@ def pre_django_1_6_cursorwrapper_getattr(self, 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 _get_query_type(query):
+    return (query.split(None, 1) or ['__empty__'])[0].lower()
 
 
-def patched_executemany(orig_executemany, self, *args, **kwargs):
-    with statsd.timer(key(self.db, 'executemany')):
-        return orig_executemany(self, *args, **kwargs)
+def patched_execute(orig_execute, self, query, *args, **kwargs):
+    with statsd.timer(key(self.db, 'execute.%s' % _get_query_type(query))):
+        return orig_execute(self, query, *args, **kwargs)
+
+def patched_executemany(orig_executemany, self, query, *args, **kwargs):
+    with statsd.timer(key(self.db, 'executemany.%s' % _get_query_type(query))):
+        return orig_executemany(self, query, *args, **kwargs)
 
 
-def patched_callproc(orig_callproc, self, *args, **kwargs):
-    with statsd.timer(key(self.db, 'callproc')):
-        return orig_callproc(self, *args, **kwargs)
+def patched_callproc(orig_callproc, self, query, *args, **kwargs):
+    with statsd.timer(key(self.db, 'callproc.%s' % _get_query_type(query))):
+        return orig_callproc(self, query, *args, **kwargs)
 
 
 def patch():
diff --git a/django_statsd/tests.py b/django_statsd/tests.py
index f4a8dad..8366122 100644
--- a/django_statsd/tests.py
+++ b/django_statsd/tests.py
@@ -474,27 +474,47 @@ class TestPatchMethod(TestCase):
 
 
 class TestCursorWrapperPatching(TestCase):
+    example_queries = {
+        'select': 'select * from something;',
+        'insert': 'insert (1, 2) into something;',
+        'update': 'update something set a=1;',
+    }
 
     def test_patched_callproc_calls_timer(self):
-        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)
+        for operation, query in self.example_queries.items():
+            with mock.patch.object(statsd, 'timer') as timer:
+                client = mock.Mock(executable_name='client_executable_name')
+                db = mock.Mock(executable_name='name', alias='alias', client=client)
+                instance = mock.Mock(db=db)
+
+                patched_callproc(lambda *args, **kwargs: None, instance, query)
+
+                self.assertEqual(timer.call_count, 1)
+                self.assertEqual(timer.call_args[0][0], 'db.client_executable_name.alias.callproc.%s' % operation)
 
     def test_patched_execute_calls_timer(self):
-        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)
+        for operation, query in self.example_queries.items():
+            with mock.patch.object(statsd, 'timer') as timer:
+                client = mock.Mock(executable_name='client_executable_name')
+                db = mock.Mock(executable_name='name', alias='alias', client=client)
+                instance = mock.Mock(db=db)
+
+                patched_execute(lambda *args, **kwargs: None, instance, query)
+
+                self.assertEqual(timer.call_count, 1)
+                self.assertEqual(timer.call_args[0][0], 'db.client_executable_name.alias.execute.%s' % operation)
 
     def test_patched_executemany_calls_timer(self):
-        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)
+        for operation, query in self.example_queries.items():
+            with mock.patch.object(statsd, 'timer') as timer:
+                client = mock.Mock(executable_name='client_executable_name')
+                db = mock.Mock(executable_name='name', alias='alias', client=client)
+                instance = mock.Mock(db=db)
+
+                patched_executemany(lambda *args, **kwargs: None, instance, query)
+
+                self.assertEqual(timer.call_count, 1)
+                self.assertEqual(timer.call_args[0][0], 'db.client_executable_name.alias.executemany.%s' % operation)
 
     @mock.patch(
         'django_statsd.patches.db.pre_django_1_6_cursorwrapper_getattr')