Redesigned architecture, added the ability to add your own output handlers

Author: Friskes

README.md

@@ -104,3 +104,23 @@ for _ in CaptureQueries(advanced_verb=True, queries=True, explain=True):
 >>>
 >>> Tests count: 1  |  Total queries count: 2  |  Total execution time: 0.22s  |  Median time one test is: 0.109s  |  Vendor: sqlite
 ```
+
+### Customization of the display
+> To customize the display of SQL queries, you can import a list with handlers and remove handlers from it or expand it with your own handlers.
+
+```python
+from capture_db_queries import settings, IHandler
+
+# NOTE: The handler must comply with the specified interface.
+class SomeHandler(IHandler):
+    def handle(self, queries_log):
+        for query in queries_log:
+            query.sql = "Hello World!"
+        return queries_log
+
+settings.PRINTER_HANDLERS.remove("capture_db_queries.handlers.ColorizeSqlHandler")
+settings.PRINTER_HANDLERS.append("path.to.your.handler.SomeHandler")
+```
+
+## TODO:
+1. Add support for other ORM's, SQLAlchemy, etc.

pyproject.toml

@@ -114,6 +114,7 @@ disable_error_code = [
     "operator",
     "assignment",
     "union-attr",
+    "comparison-overlap",
 ]
 
 
@@ -204,6 +205,7 @@ ignore = [
     "F403",    # (не ругаться на использование from ... import *)
     # https://docs.astral.sh/ruff/rules/#pyupgrade-up
     "UP031",   # (не ругаться на форматирование с помощью %s)
+    "UP035",   # (не ругаться на импорт из typing_extensions)
     "UP036",   # (не ругаться на использование sys.version_info если текущая версия не подпадает под условие)
     # https://docs.astral.sh/ruff/rules/#flake8-logging-format-g
     "G004",    # (не ругаться на использование f-строк для сообщения лога)

src/capture_db_queries/__init__.py

@@ -1,3 +1,5 @@
+# ruff: noqa: F401, E402
+
 import sys
 import warnings
 
@@ -8,7 +10,8 @@
 )
 
 if sys.version_info >= (3, 12):
-    from . import _sqlite3_adapters_and_converters  # noqa: F401
+    from . import _sqlite3_adapters_and_converters
 
-from ._logging import switch_logger, switch_trace  # noqa: F401, E402
-from .decorators import CaptureQueries, ExtCaptureQueriesContext, capture_queries  # noqa: F401, E402
+from ._logging import switch_logger, switch_trace
+from .decorators import CaptureQueries, ExtCaptureQueriesContext, capture_queries
+from .handlers import IHandler

src/capture_db_queries/decorators.py

@@ -1,12 +1,9 @@
 from __future__ import annotations
 
-import abc
 import functools
 import statistics
 import time
 import warnings
-from collections import deque
-from contextlib import contextmanager
 from functools import wraps
 from typing import TYPE_CHECKING, Any
 from unittest.util import safe_repr
@@ -17,26 +14,119 @@
     reset_queries,
 )
 from django.test.utils import CaptureQueriesContext
-from typing_extensions import Self  # noqa: UP035
+from typing_extensions import Self
 
 from ._logging import log
 from .dtos import IterationPrintDTO, SeveralPrintDTO, SinglePrintDTO
 from .printers import AbcPrinter, PrinterSql
 from .wrappers import BaseExecutionWrapper, ExplainExecutionWrapper
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Generator
+    from collections.abc import Callable
     from types import TracebackType
 
     from django.db.backends.base.base import BaseDatabaseWrapper
-    from django.db.backends.utils import CursorWrapper
 
-    from .types import DecoratedCallable, QueriesLog, Query, T
+    from .dtos import Query
+    from .types import DecoratedCallable, T
 
 __all__ = ('CaptureQueries', 'capture_queries', 'ExtCaptureQueriesContext')
 
 
-class AbcCapture(abc.ABC):
+class CaptureQueries:
+    """
+    #### Class to simplify the search for slow and suboptimal sql queries\
+    to the database in django projects.
+
+    ---
+
+    #### About class
+    - Class allows you to track any sql queries that are executed inside the body of a loop,
+    a decorated function, or a context manager,
+    the body can be executed a specified number of times to get the average query execution time.
+
+    - The class allows you to display formatted output data
+    containing brief information about the current iteration of the measurement,
+    display sql queries and explain information on them,
+    as well as summary information containing data on all measurements.
+
+    `Do not use the class inside the business logic of your application,
+    this will greatly slow down the execution of the queries,
+    the class is intended only for the test environment.`
+
+    ---
+
+    #### - Optional parameters:
+        - `assert_q_count`: The expected number of database queries during all `number_runs`, otherwise an exception will be raised: "AssertionError: `N` not less than or equal to `N` queries".
+        - `number_runs`: The number of runs of the decorated function or test for loop.
+        - `verbose`: Displaying the final results of test measurements within all `number_runs`.
+        - `advanced_verb`: Displaying the result of each test measurement.
+        - `auto_call_func`: Autorun of the decorated function. (without passing arguments to the function, since the launch takes place inside the class).
+        - `queries`: Displaying colored and formatted SQL queries to the database.
+        - `explain`: Displaying explain information about each query. (has no effect on the original query).
+        - `explain_opts`: Parameters for explain. (for more information about the parameters for explain, see the documentation for your DBMS).
+        - `connection`: Connecting to your database, by default: django.db.connection
+
+    ---
+
+    #### Usage examples::
+
+        for ctx in CaptureQueries(number_runs=2, advanced_verb=True):
+            response = self.client.get(url)
+
+        >>> Test №1 | Queries count: 10 | Execution time: 0.04s
+        >>> Test №2 | Queries count: 10 | Execution time: 0.04s
+        >>> Tests count: 2  |  Total queries count: 20  |  Total execution time: 0.08s  |  Median time one test is: 0.041s  |  Vendor: sqlite
+
+        # OR
+
+        @CaptureQueries(number_runs=2, advanced_verb=True)
+        def test_request():
+            response = self.client.get(url)
+
+        >>> Test №1 | Queries count: 10 | Execution time: 0.04s
+        >>> Test №2 | Queries count: 10 | Execution time: 0.04s
+        >>> Tests count: 2  |  Total queries count: 20  |  Total execution time: 0.08s  |  Median time one test is: 0.041s  |  Vendor: sqlite
+
+        # OR
+
+        # NOTE: The with context manager does not support multi-launch number_runs > 1
+        with CaptureQueries(number_runs=1, advanced_verb=True) as ctx:
+            response = self.client.get(url)
+
+        >>> Queries count: 10  |  Execution time: 0.04s  |  Vendor: sqlite
+
+    ---
+
+    #### Example of output when using queries and explain::
+
+        for _ in CaptureQueries(advanced_verb=True, queries=True, explain=True):
+            list(Reporter.objects.filter(pk=1))
+            list(Article.objects.filter(pk=1))
+
+        >>> Test №1 | Queries count: 2 | Execution time: 0.22s
+        >>>
+        >>>
+        >>> №[1] time=[0.109] explain=['2 0 0 SEARCH TABLE tests_reporter USING INTEGER PRIMARY KEY (rowid=?)']
+        >>> SELECT "tests_reporter"."id",
+        >>>     "tests_reporter"."full_name"
+        >>> FROM "tests_reporter"
+        >>> WHERE "tests_reporter"."id" = 1
+        >>>
+        >>>
+        >>> №[2] time=[0.109] explain=['2 0 0 SEARCH TABLE tests_article USING INTEGER PRIMARY KEY (rowid=?)']
+        >>> SELECT "tests_article"."id",
+        >>>     "tests_article"."pub_date",
+        >>>     "tests_article"."headline",
+        >>>     "tests_article"."content",
+        >>>     "tests_article"."reporter_id"
+        >>> FROM "tests_article"
+        >>> WHERE "tests_article"."id" = 1
+        >>>
+        >>>
+        >>> Tests count: 1  |  Total queries count: 2  |  Total execution time: 0.22s  |  Median time one test is: 0.109s  |  Vendor: sqlite
+    """  # noqa: E501
+
     def __init__(
         self,
         assert_q_count: int | None = None,
@@ -73,12 +163,8 @@ def __init__(
             self.connection.vendor, assert_q_count, verbose, advanced_verb, queries
         )
 
-        self.queries_log: QueriesLog = deque(maxlen=self.connection.queries_limit)
-
-        self.wrapper = self.wrapper_cls(
-            connection=self.connection, queries_log=self.queries_log, explain_opts=explain_opts or {}
-        )
-        self.wrapper_ctx_manager = self.__wrap_reqs_in_wrapper()
+        self.wrapper = self.wrapper_cls(connection=self.connection, explain_opts=explain_opts or {})
+        self.wrapper_ctx_manager = self.wrapper.wrap_reqs_in_wrapper()
 
     @property
     def printer_cls(self) -> type[AbcPrinter]:
@@ -92,15 +178,6 @@ def wrapper_cls(self) -> type[BaseExecutionWrapper]:
             return ExplainExecutionWrapper
         return BaseExecutionWrapper
 
-    @contextmanager
-    def __wrap_reqs_in_wrapper(self) -> Generator[None, None, CursorWrapper]:
-        """Wraps all database requests in a wrapper."""
-        log.debug('')
-
-        # https://docs.djangoproject.com/en/5.1/topics/db/instrumentation/#connection-execute-wrapper
-        with self.connection.execute_wrapper(self.wrapper):
-            yield
-
     def __enter__(self) -> Self:
         log.debug('')
 
@@ -133,7 +210,7 @@ def __exit__(
         self.printer.print_single_sql(
             SinglePrintDTO(
                 queries_count=queries_count,
-                queries_log=self.queries_log,
+                queries_log=self.wrapper.queries_log,
                 execution_time_per_iter=self.wrapper.timer.execution_time_per_iter,
             )
         )
@@ -169,7 +246,7 @@ def __next__(self) -> Self:
         self.printer.print_several_sql(
             SeveralPrintDTO(
                 queries_count=queries_count,
-                queries_log=self.queries_log,
+                queries_log=self.wrapper.queries_log,
                 current_iteration=self.current_iteration,
                 all_execution_times=self.wrapper.timer.all_execution_times,
             )
@@ -197,115 +274,16 @@ def wrapped(*args: Any, **kwargs: Any) -> Any:
     def __len__(self) -> int:
         log.debug('')
 
-        return len(self.queries_log)
+        return len(self.wrapper.queries_log)
 
     def __getitem__(self, index: int) -> Query | None:
         log.debug('')
 
         try:
-            return self.queries_log[index]
+            return self.wrapper.queries_log[index]
         except IndexError:
             return None
 
-    @abc.abstractmethod
-    def _assert_queries_count(self, queries_count: int) -> None:
-        raise NotImplementedError
-
-
-class CaptureQueries(AbcCapture):
-    """
-    #### Class to simplify the search for slow and suboptimal sql queries\
-    to the database in django projects.
-
-    ---
-
-    #### About class
-    - Class allows you to track any sql queries that are executed inside the body of a loop,
-    a decorated function, or a context manager,
-    the body can be executed a specified number of times to get the average query execution time.
-
-    - The class allows you to display formatted output data
-    containing brief information about the current iteration of the measurement,
-    display sql queries and explain information on them,
-    as well as summary information containing data on all measurements.
-
-    `Do not use the class inside the business logic of your application,
-    this will greatly slow down the execution of the queries,
-    the class is intended only for the test environment.`
-
-    ---
-
-    #### - Optional parameters:
-        - `assert_q_count`: The expected number of database queries during all `number_runs`, otherwise an exception will be raised: "AssertionError: `N` not less than or equal to `N` queries".
-        - `number_runs`: The number of runs of the decorated function or test for loop.
-        - `verbose`: Displaying the final results of test measurements within all `number_runs`.
-        - `advanced_verb`: Displaying the result of each test measurement.
-        - `auto_call_func`: Autorun of the decorated function. (without passing arguments to the function, since the launch takes place inside the class).
-        - `queries`: Displaying colored and formatted SQL queries to the database.
-        - `explain`: Displaying explain information about each query. (has no effect on the original query).
-        - `explain_opts`: Parameters for explain. (for more information about the parameters for explain, see the documentation for your DBMS).
-        - `connection`: Connecting to your database, by default: django.db.connection
-
-    ---
-
-    #### Usage examples::
-
-        for ctx in CaptureQueries(number_runs=2, advanced_verb=True):
-            response = self.client.get(url)
-
-        >>> Test №1 | Queries count: 10 | Execution time: 0.04s
-        >>> Test №2 | Queries count: 10 | Execution time: 0.04s
-        >>> Tests count: 2  |  Total queries count: 20  |  Total execution time: 0.08s  |  Median time one test is: 0.041s  |  Vendor: sqlite
-
-        # OR
-
-        @CaptureQueries(number_runs=2, advanced_verb=True)
-        def test_request():
-            response = self.client.get(url)
-
-        >>> Test №1 | Queries count: 10 | Execution time: 0.04s
-        >>> Test №2 | Queries count: 10 | Execution time: 0.04s
-        >>> Tests count: 2  |  Total queries count: 20  |  Total execution time: 0.08s  |  Median time one test is: 0.041s  |  Vendor: sqlite
-
-        # OR
-
-        # NOTE: The with context manager does not support multi-launch number_runs > 1
-        with CaptureQueries(number_runs=1, advanced_verb=True) as ctx:
-            response = self.client.get(url)
-
-        >>> Queries count: 10  |  Execution time: 0.04s  |  Vendor: sqlite
-
-    ---
-
-    #### Example of output when using queries and explain::
-
-        for _ in CaptureQueries(advanced_verb=True, queries=True, explain=True):
-            list(Reporter.objects.filter(pk=1))
-            list(Article.objects.filter(pk=1))
-
-        >>> Test №1 | Queries count: 2 | Execution time: 0.22s
-        >>>
-        >>>
-        >>> №[1] time=[0.109] explain=['2 0 0 SEARCH TABLE tests_reporter USING INTEGER PRIMARY KEY (rowid=?)']
-        >>> SELECT "tests_reporter"."id",
-        >>>     "tests_reporter"."full_name"
-        >>> FROM "tests_reporter"
-        >>> WHERE "tests_reporter"."id" = 1
-        >>>
-        >>>
-        >>> №[2] time=[0.109] explain=['2 0 0 SEARCH TABLE tests_article USING INTEGER PRIMARY KEY (rowid=?)']
-        >>> SELECT "tests_article"."id",
-        >>>     "tests_article"."pub_date",
-        >>>     "tests_article"."headline",
-        >>>     "tests_article"."content",
-        >>>     "tests_article"."reporter_id"
-        >>> FROM "tests_article"
-        >>> WHERE "tests_article"."id" = 1
-        >>>
-        >>>
-        >>> Tests count: 1  |  Total queries count: 2  |  Total execution time: 0.22s  |  Median time one test is: 0.109s  |  Vendor: sqlite
-    """  # noqa: E501
-
     def _assert_queries_count(self, queries_count: int) -> None:
         log.debug('')
 

src/capture_db_queries/dtos.py

@@ -8,6 +8,17 @@
     from .types import QueriesLog
 
 
+@dataclass
+class Query:
+    sql: str
+    time: float
+
+
+@dataclass
+class ExpQuery(Query):
+    explain: str
+
+
 @dataclass(frozen=True)
 class BasePrintDTO:
     queries_count: int