Repository: adamchainz/django-mysql

Issue:
fingerprint function that wraps pt-fingerprint
Simple subprocess wrapping, would help with fingerprinting queries for analysis. Maybe useful in #18 ?



Discussion:

Resolving patch:
diff --git a/django_mysql/utils.py b/django_mysql/utils.py
--- a/django_mysql/utils.py
+++ b/django_mysql/utils.py
@@ -1,9 +1,14 @@
 # -*- coding:utf-8 -*-
 from __future__ import division
 
+import os
+import pty
 import time
 from contextlib import contextmanager
-from subprocess import PIPE, call
+from subprocess import PIPE, Popen, call
+from threading import Lock, Thread
+
+from django.utils.six.moves.queue import Empty, Queue
 
 
 class WeightedAverageRate(object):
@@ -99,6 +104,108 @@ def settings_to_cmd_args(settings_dict):
     return args
 
 
+programs_memo = {}
+
+
 def have_program(program_name):
-    status = call(['which', program_name], stdout=PIPE)
-    return (status == 0)
+    global programs_memo
+    if program_name not in programs_memo:
+        status = call(['which', program_name], stdout=PIPE)
+        programs_memo[program_name] = (status == 0)
+
+    return programs_memo[program_name]
+
+
+def pt_fingerprint(query):
+    """
+    Takes a query (in a string) and returns its 'fingerprint'
+    """
+    if not have_program('pt-fingerprint'):  # pragma: no cover
+        raise OSError("pt-fingerprint doesn't appear to be installed")
+
+    thread = PTFingerprintThread.get_thread()
+    thread.in_queue.put(query)
+    return thread.out_queue.get()
+
+
+class PTFingerprintThread(Thread):
+    """
+    Class for a singleton background thread to pass queries to pt-fingerprint
+    and get their fingerprints back. This is done because the process launch
+    time is relatively expensive and it's useful to be able to fingerprinting
+    queries quickly.
+
+    The get_thread() class method returns the singleton thread - either
+    instantiating it or returning the existing one.
+
+    The thread launches pt-fingerprint with subprocess and then takes queries
+    from an input queue, passes them the subprocess and returns the fingerprint
+    to an output queue. If it receives no queries in PROCESS_LIFETIME seconds,
+    it closes the subprocess and itself - so you don't have processes hanging
+    around.
+    """
+
+    the_thread = None
+    life_lock = Lock()
+
+    PROCESS_LIFETIME = 60.0  # seconds
+
+    @classmethod
+    def get_thread(cls):
+        with cls.life_lock:
+            if cls.the_thread is None:
+                in_queue = Queue()
+                out_queue = Queue()
+                thread = cls(in_queue, out_queue)
+                thread.daemon = True
+                thread.in_queue = in_queue
+                thread.out_queue = out_queue
+                thread.start()
+                cls.the_thread = thread
+
+        return cls.the_thread
+
+    def __init__(self, in_queue, out_queue, **kwargs):
+        self.in_queue = in_queue
+        self.out_queue = out_queue
+        super(PTFingerprintThread, self).__init__(**kwargs)
+
+    def run(self):
+        global fingerprint_thread
+        master, slave = pty.openpty()
+        proc = Popen(
+            ['pt-fingerprint'],
+            stdin=PIPE,
+            stdout=slave,
+            close_fds=True
+        )
+        stdin = proc.stdin
+        stdout = os.fdopen(master)
+
+        while True:
+            try:
+                query = self.in_queue.get(timeout=self.PROCESS_LIFETIME)
+            except Empty:
+                self.life_lock.acquire()
+                # We timed out, but there was something put into the queue
+                # since
+                if (
+                    self.__class__.the_thread is self and
+                    self.in_queue.qsize()
+                ):  # pragma: no cover
+                    self.life_lock.release()
+                    break
+                # Die
+                break
+
+            stdin.write(query.encode('utf-8'))
+            if not query.endswith(';'):
+                stdin.write(';'.encode('ascii'))
+            stdin.write('\n'.encode('ascii'))
+            stdin.flush()
+            fingerprint = stdout.readline()
+            self.out_queue.put(fingerprint.strip())
+
+        stdin.close()
+        self.__class__.the_thread = None
+        self.life_lock.release()


Test patch:
diff --git a/tests/django_mysql_tests/test_utils.py b/tests/django_mysql_tests/test_utils.py
--- a/tests/django_mysql_tests/test_utils.py
+++ b/tests/django_mysql_tests/test_utils.py
@@ -1,7 +1,14 @@
 # -*- coding:utf-8 -*-
+from __future__ import unicode_literals
+
+from time import sleep
+from unittest import skipUnless
+
 from django.test import TestCase
 
-from django_mysql.utils import WeightedAverageRate
+from django_mysql.utils import (
+    have_program, pt_fingerprint, PTFingerprintThread, WeightedAverageRate
+)
 
 
 class WeightedAverageRateTests(TestCase):
@@ -37,3 +44,52 @@ def test_good_guess(self):
         self.assertEqual(rate.update(50, 0.5), 50)
         self.assertEqual(rate.update(50, 0.5), 50)
         self.assertEqual(rate.update(50, 0.5), 50)
+
+
+@skipUnless(have_program('pt-fingerprint'),
+            "pt-fingerprint must be installed")
+class PTFingerprintTests(TestCase):
+
+    def test_basic(self):
+        self.assertEqual(pt_fingerprint('SELECT 5'), 'select ?')
+        self.assertEqual(pt_fingerprint('SELECT 5;'), 'select ?')
+
+    def test_long(self):
+        query = """
+            SELECT
+                CONCAT(customer.last_name, ', ', customer.first_name)
+                    AS customer,
+                address.phone,
+                film.title
+            FROM rental
+                INNER JOIN customer
+                    ON rental.customer_id = customer.customer_id
+                INNER JOIN address
+                    ON customer.address_id = address.address_id
+                INNER JOIN inventory
+                    ON rental.inventory_id = inventory.inventory_id
+                INNER JOIN film
+                    ON inventory.film_id = film.film_id
+            WHERE
+                rental.return_date IS NULL AND
+                rental_date + INTERVAL film.rental_duration DAY <
+                    CURRENT_DATE()
+            LIMIT 5"""
+        self.assertEqual(
+            pt_fingerprint(query),
+            "select concat(customer.last_name, ?, customer.first_name) as "
+            "customer, address.phone, film.title from rental inner join "
+            "customer on rental.customer_id = customer.customer_id inner join "
+            "address on customer.address_id = address.address_id inner join "
+            "inventory on rental.inventory_id = inventory.inventory_id inner "
+            "join film on inventory.film_id = film.film_id where "
+            "rental.return_date is ? and rental_date ? interval "
+            "film.rental_duration day < current_date() limit ?"
+        )
+
+    def test_the_thread_shuts_on_time_out(self):
+        PTFingerprintThread.PROCESS_LIFETIME = 0.1
+        pt_fingerprint("select 123")
+        sleep(0.2)
+        self.assertIsNone(PTFingerprintThread.the_thread)
+        PTFingerprintThread.PROCESS_LIFETIME = 60