restructure listener and analyzer to simplify test implementation

Author: gruuya

README.md

@@ -1,19 +1,20 @@
 Django SQL Sniffer
 ==================
 A simple command line tool for analyzing Django ORM SQL execution from a running process.
-Minimally invasive and granular - no need to change logging config or restart the process.
+Minimally invasive and granular - no need to change the logging config or the source code 
+of the target process, and no need to restart the process.
 
 # Usage
 Install though pip
 ```
 pip install django-sql-sniffer
 ```
-Run the tool by passing it a process id which is to be analyzed
+Run the tool by passing it the target process id
 ```
 django-sql-sniffer -p 76441
 ```
-`Ctrl + C` to stop and show the query stats summary. Here's a short demo:
-![https://raw.githubusercontent.com/gruuya/django-sql-sniffer/master/demo.webp](demo.webp)
+`Ctrl + C` to stop and show the query stats summary. Take a look at a short demo:
+![demo](demo.webp)
 By default, stats summary shows queries sorted by max duration; the possible options include:
 - `-t` print queries in tail mode, i.e. as they are executed
 - `-c` sort stats summary by query count

django_sql_sniffer/__init__.py

@@ -1 +1 @@
-version = "1.0.2"
+version = "1.0.3.dev0"

django_sql_sniffer/analyzer.py

@@ -1,9 +1,8 @@
-import threading
 try:
     import sqlparse
 except ImportError:
     sqlparse = None
-from django_sql_sniffer import sniffer, version
+from django_sql_sniffer import version
 
 
 SQL_STATS_TEXT = """
@@ -26,17 +25,13 @@ def format_duration(duration):
     return "{0:0.9f}".format(duration)
 
 
-class SQLAnalyzer(threading.Thread):
-    def __init__(self, conn, verbose=False, tail=False, top=3, by_sum=False, by_count=False):
-        super().__init__(name=self.__class__.__name__)
-        self._conn = conn
+class SQLAnalyzer:
+    def __init__(self, tail=False, top=3, by_sum=False, by_count=False):
         self._executed_queries = dict()
         self._tail = tail
         self._top = top
         self._by_sum = by_sum
         self._by_count = by_count
-        self._running = False
-        self.logger = sniffer.configure_logger(__name__, verbose)
 
     def record_query(self, sql, duration):
         if sql in self._executed_queries:
@@ -50,6 +45,9 @@ def record_query(self, sql, duration):
                 sum=duration
             )
 
+        if self._tail:
+            self.print_query(sql, duration)
+
     def print_query(self, sql, duration):
         stats = self._executed_queries[sql]
         print("Count: ", stats["count"], "; Duration: ", format_duration(duration), "; Max Duration: ", format_duration(stats["max"]), "; Query:", sep="")
@@ -70,32 +68,3 @@ def print_summary(self, *a, **kw):
             print(format_sql(sql))
             print("-" * DELIMETER_LENGTH)
         print("=" * DELIMETER_LENGTH)
-
-    def start(self):
-        self.logger.debug("starting")
-        self._running = True
-        super().start()
-
-    def stop(self, *a, **kw):
-        self.logger.debug("stopping")
-        self._running = False
-
-    def run(self):
-        while self._running:
-            try:
-                if self._conn.poll(3):
-                    sql_packet = self._conn.recv()
-                    duration = sql_packet["duration"]
-                    sql = sql_packet["sql"]
-                    self.record_query(sql, duration)
-                    if self._tail:
-                        self.print_query(sql, duration)
-            except EOFError:
-                self.logger.info("sniffer disconnected, exiting")
-                self._running = False
-            except Exception as e:
-                self.logger.error(f"unexpected error: {str(e)}", exc_info=True)
-                self._running = False
-
-        self._conn.close()
-        self.logger.debug("done")

django_sql_sniffer/listener.py

@@ -1,12 +1,75 @@
 import argparse
 import signal
 import time
+import threading
 from multiprocessing.connection import Listener
 from django_sql_sniffer import analyzer, injector, sniffer
 
 
+class DjangoSQLListener(threading.Thread):
+    def __init__(self, cmdline_args):
+        super().__init__(name=self.__class__.__name__)
+        self.analyzer = analyzer.SQLAnalyzer(tail=cmdline_args.tail, top=cmdline_args.number, by_sum=cmdline_args.sum, by_count=cmdline_args.count)
+        self.logger = sniffer.configure_logger(__name__, cmdline_args.verbose)
+        self._target_pid = cmdline_args.pid
+        self._verbose = cmdline_args.verbose
+        self._listener = None
+        self._conn = None
+        self._running = False
+
+    def _create_socket(self):
+        self._listener = Listener(('localhost', 0))  # will force OS to assign a random available port
+        _, self.port = self._listener.address
+        self.logger.debug(f"listening on port {self.port}")
+
+    def _inject_sniffer(self):
+        with open(sniffer.__file__, "r") as source_file:
+            source_code = source_file.read()
+            code_to_inject = source_code.replace("PORT = 0", f"PORT = {self.port}")
+            code_to_inject = code_to_inject.replace("LOGGING_ENABLED = False", f"LOGGING_ENABLED = {self._verbose}")
+            code_to_inject += "sniffer.start()"
+        injector.inject(str(self._target_pid), code_to_inject, self._verbose)
+        self.logger.debug("SQL sniffer injected, waiting for a reply")
+
+    def _callback_wait(self):
+        self._conn = self._listener.accept()
+        self._listener.close()
+        self.logger.debug("reply received, sniffer active")
+
+    def start(self):
+        self.logger.debug("starting")
+        self._create_socket()
+        self._inject_sniffer()
+        self._callback_wait()
+        self._running = True
+        super().start()
+
+    def stop(self, *a, **kw):
+        self.logger.debug("stopping")
+        self._running = False
+
+    def run(self):
+        while self._running:
+            try:
+                if self._conn.poll(3):
+                    sql_packet = self._conn.recv()
+                    duration = sql_packet["duration"]
+                    sql = sql_packet["sql"]
+                    self.analyzer.record_query(sql, duration)
+            except EOFError:
+                self.logger.info("sniffer disconnected, exiting")
+                self._running = False
+            except Exception as e:
+                self.logger.error(f"unexpected error: {str(e)}", exc_info=True)
+                self._running = False
+
+        self.analyzer.print_summary()
+        self._conn.close()
+        injector.inject(str(self._target_pid), "sniffer.stop()", self._verbose)
+        self.logger.debug("done")
+
+
 def main():
-    # parse arguments
     parser = argparse.ArgumentParser(description="Analyze SQL queries originating from a running process")
     parser.add_argument("-p", "--pid", help="The id of the process executing Django SQL queries")
     parser.add_argument("-t", "--tail", action='store_true', help="Log queries as they are executed in tail mode")
@@ -15,39 +78,15 @@ def main():
     parser.add_argument("-c", "--count", action='store_true', help="Sort query summary by execution count")
     parser.add_argument("-n", "--number", type=int, default=3, help="Number of top queries and their stats to display in summary")
     args = parser.parse_args()
-    logger = sniffer.configure_logger(__name__, args.verbose)
-
-    # create socket and listen for a connection
-    listener = Listener(('localhost', 0))  # will force OS to assign a random available port
-    _, port = listener.address
-    logger.debug(f"listening on port {port}")
-
-    # inject sniffer monkey patch
-    with open(sniffer.__file__, "r") as source_file:
-        source_code = source_file.read()
-        code_to_inject = source_code.replace("PORT = 0", f"PORT = {port}")
-        code_to_inject = code_to_inject.replace("LOGGING_ENABLED = False", f"LOGGING_ENABLED = {args.verbose}")
-        code_to_inject += "sniffer.start()"
-    injector.inject(str(args.pid), code_to_inject, args.verbose)
-
-    # wait for callback
-    logger.debug("SQL sniffer injected, waiting for a reply")
-    conn = listener.accept()
-    listener.close()
-
-    # receive and analyze executed SQL
-    logger.debug("reply received, sniffer active")
-    sql_analyzer = analyzer.SQLAnalyzer(conn, verbose=args.verbose, tail=args.tail, top=args.number, by_sum=args.sum, by_count=args.count)
-    signal.signal(signal.SIGTERM, sql_analyzer.stop)
-    signal.signal(signal.SIGINT, sql_analyzer.stop)
-    if hasattr(signal, "SIGINFO"):
-        signal.signal(signal.SIGINFO, sql_analyzer.print_summary)  # print summary on ^T, without exiting the process
-    sql_analyzer.start()
 
-    while sql_analyzer.is_alive():
-        time.sleep(0.5)
-    injector.inject(str(args.pid), "sniffer.stop()", args.verbose)
-    sql_analyzer.print_summary()
+    listener = DjangoSQLListener(args)
+    listener.start()
+    signal.signal(signal.SIGTERM, listener.stop)
+    signal.signal(signal.SIGINT, listener.stop)
+    if hasattr(signal, "SIGINFO"):
+        signal.signal(signal.SIGINFO, listener.analyzer.print_summary)  # print summary on ^T, without exiting the process
+    while listener.is_alive():
+        time.sleep(1)
 
 
 if __name__ == "__main__":

setup.py

@@ -2,8 +2,10 @@
 import django_sql_sniffer
 
 
+URL = 'https://github.com/gruuya/django-sql-sniffer'
 with open("README.md", "r", encoding="utf-8") as fh:
     long_description = fh.read()
+    long_description = long_description.replace('![demo](demo.webp)', f'[demo]({URL})\\')
 
 
 setuptools.setup(
@@ -13,7 +15,7 @@
     long_description=long_description,
     long_description_content_type='text/markdown',
     keywords='django sql query remote process analysis',
-    url='https://github.com/gruuya/django-sql-sniffer',
+    url=URL,
     author='Marko Grujic',
     author_email='[email protected]',
     license='MIT',