redesigned logic and added tests

Author: Friskes

src/capture_db_queries/wrappers.py

@@ -2,10 +2,12 @@
 
 import json
 import time
+import traceback
 from typing import TYPE_CHECKING, Any, NamedTuple
 
 from django.utils.regex_helper import _lazy_re_compile
 
+from ._logging import log
 from .timers import ContextTimer
 
 if TYPE_CHECKING:
@@ -19,33 +21,48 @@
 
 
 class BaseExecutionWrapper:
-    def __init__(self, queries_log: QueriesLog, *args: Any, **kwargs: Any) -> None:
-        self.timer = ContextTimer(time.perf_counter)
+    def __init__(
+        self, connection: BaseDatabaseWrapper, queries_log: QueriesLog, *args: Any, **kwargs: Any
+    ) -> None:
+        log.debug('')
+
+        self.connection = connection
         self.queries_log = queries_log
+        self.timer = ContextTimer(time.perf_counter)
+
+        # Wrappers for specific databases are stored at addresses:
+        # django.db.backends.sqlite3.operations.DatabaseOperations
+        # django.db.backends.postgresql.operations.DatabaseOperations
+        self.db_operations: BaseDatabaseOperations = self.connection.ops
 
     def __call__(
         self,
-        execute: Callable[..., CursorWrapper | None],
+        execute: Callable[..., CursorWrapper],
         sql: str,
         params: tuple[Any, ...],
         many: bool,
         context: dict[str, Any],
     ) -> CursorWrapper | None:
         """
-        Выполняет оригинальный запрос
+        Executes the original SQL request.
         """
         with self.timer as timer:
             try:
                 result = execute(sql, params, many, context)
             except Exception as exc:
-                print('Что-то пошло не так:', exc)
+                print('Something went wrong:', exc)
                 return None
 
-        # from django.db.backends.utils import CursorDebugWrapper(connection.cursor(), connection)
+        if not many:
+            # Get filled SQL with params
+            sql = self.db_operations.last_executed_query(context['cursor'], sql, params)
+
         query: Query = {'sql': sql, 'time': timer.execution_time}
         query = self.update_query(query)
         self.queries_log.append(query)
 
+        log.trace('Location of SQL Call:\n%s', ''.join(traceback.format_stack()))
+
         return result
 
     def update_query(self, query: Query) -> Query:
@@ -54,12 +71,13 @@ def update_query(self, query: Query) -> Query:
 
 class ExplainExecutionWrapper(BaseExecutionWrapper):
     """
-    Класс для вызова EXPLAIN на каждом SELECT-запросе.
-    С сохранением полной функциональности метода QuerySet.explain().
+    A class for calling EXPLAIN before each original SELECT request.
+    While maintaining the full functionality of the Query Set.explain() method.
+
+    The EXPLAIN call is not fixed in any way and does not affect the measurement results.
 
-    https://docs.djangoproject.com/en/5.1/topics/db/instrumentation/#connection-execute-wrapper
     https://docs.djangoproject.com/en/5.1/ref/models/querysets/#explain
-    from django_extensions.management.debug_cursor import monkey_patch_cursordebugwrapper
+    https://www.postgresql.org/docs/current/sql-explain.html
     """
 
     # Inspired from
@@ -72,20 +90,18 @@ class ExplainInfo(NamedTuple):
 
     def __init__(
         self,
-        queries_log: QueriesLog,
         connection: BaseDatabaseWrapper,
+        queries_log: QueriesLog,
         explain_opts: dict[str, Any],
         *args: Any,
         **kwargs: Any,
     ) -> None:
-        # # https://www.postgresql.org/docs/current/sql-explain.html
-        super().__init__(queries_log, *args, **kwargs)
-        self.connection = connection
+        super().__init__(connection, queries_log, *args, **kwargs)
         self.explain_info = self.build_explain_info(**explain_opts)
 
     def __call__(
         self,
-        execute: Callable[..., CursorWrapper | None],
+        execute: Callable[..., CursorWrapper],
         sql: str,
         params: tuple[Any, ...],
         many: bool,
@@ -100,36 +116,31 @@ def update_query(self, query: Query) -> Query:
         return query
 
     def _execute_explain(self, sql: str, params: tuple[Any, ...], many: bool) -> str | None:
-        # Проверяем, является ли запрос SELECT-запросом
+        # Checking whether the request is a SELECT request
         if not sql.strip().lower().startswith('select'):
             return None
 
-        # Обёртки для конкретных бд хранятся по адресам:
-        # django.db.backends.sqlite3.operations.DatabaseOperations
-        # django.db.backends.postgresql.operations.DatabaseOperations
-        db_operations: BaseDatabaseOperations = self.connection.ops
-
-        explain_query = db_operations.explain_query_prefix(
+        explain_query = self.db_operations.explain_query_prefix(
             self.explain_info.format, **self.explain_info.options
         )
         explain_query = f'{explain_query} {sql}'
 
         try:
             raw_explain = self.__execute(explain_query, params, many)
         except Exception as exc:
-            print('Что-то пошло не так:', exc)
+            print('Something went wrong:', exc)
             return None
         else:
             return '\n'.join(self.format_explain(raw_explain))
 
     def __execute(self, explain_query: str, params: tuple[Any, ...], many: bool) -> Explain:
         with self.connection.cursor() as cursor:
-            # нельзя вызывать execute или executemany
-            # который вызывает внутри себя обёртку,
-            # потому что это приведёт к дублированию вызова,
-            # и вызовет __call__ текущей обёртки
+            # you cannot call execute or executemany
+            # which invokes a wrapper inside itself,
+            # because it will result in duplicate call,
+            # and will call __call__ of the current wrapper
 
-            # данные запросы идут в обход обёрток
+            # these requests bypass wrappers
             if many:
                 cursor._executemany(explain_query, params)
             else:
@@ -139,15 +150,17 @@ def __execute(self, explain_query: str, params: tuple[Any, ...], many: bool) ->
 
     def build_explain_info(self, *, format: str | None = None, **options: dict[str, Any]) -> ExplainInfo:  # noqa: A002
         """
-        Runs an EXPLAIN on the SQL query this QuerySet would perform, and
-        returns the results.
+        Validates explain options and build ExplainInfo object.
         """
         for option_name in options:
             if not self.EXPLAIN_OPTIONS_PATTERN.fullmatch(option_name) or '--' in option_name:
                 raise ValueError(f'Invalid option name: {option_name!r}.')
         return self.ExplainInfo(format, options)
 
     def format_explain(self, result: Explain) -> Generator[str, None, None]:
+        """
+        Splits the explain tuple into its components and collects the final explain string from them.
+        """
         nested_result = [list(result)]
         # Some backends return 1 item tuples with strings, and others return
         # tuples with integers and strings. Flatten them out into strings.

tests/test_decorators.py

@@ -41,14 +41,15 @@ def _select(reporter_id: UUID, article_id: UUID) -> None:
 
 
 class BasicTestsFor3ChoicesOfCaptureQueries:
-    """Обязательно должен быть передан аргумент -s при запуске pytest, для вывода output"""
+    """"""
 
     def setup_method(self, method: Callable[..., Any]) -> None:
         self.reporter, self.article = request_to_db()
 
     def call_capture_queries(self, **kwargs: Any) -> CaptureQueries:
         raise NotImplementedError
 
+    # @pytest.mark.usefixtures('_debug_true')
     def test_basic_logic(self) -> None:
         obj = self.call_capture_queries()
 
@@ -74,14 +75,14 @@ def test_basic_logic(self) -> None:
         data = obj.queries_log[0]['sql']
         assert obj.queries_log[0]['sql'] == (
             'SELECT "tests_reporter"."id", "tests_reporter"."full_name" '
-            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s'
+            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s' % self.reporter.pk
         ), data
 
         data = obj.queries_log[1]['sql']
         assert obj.queries_log[1]['sql'] == (
             '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" = %s'
+            'FROM "tests_article" WHERE "tests_article"."id" = %s' % self.article.pk
         ), data
 
         with pytest.raises(KeyError, match='explain'):
@@ -110,14 +111,14 @@ def test_param__explain(self) -> None:
         data = obj.queries_log[0]['sql']
         assert data == (
             'SELECT "tests_reporter"."id", "tests_reporter"."full_name" '
-            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s'
+            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s' % self.reporter.pk
         ), data
 
         data = obj.queries_log[1]['sql']
         assert data == (
             '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" = %s'
+            'FROM "tests_article" WHERE "tests_article"."id" = %s' % self.article.pk
         ), data
 
         # data = obj.queries_log[0]['explain']
@@ -134,6 +135,7 @@ def test_param__connection(self) -> None:
         class FakeConnection:
             vendor = 'fake_vendor'
             queries_limit = 4
+            ops = connection.ops
 
             @contextmanager
             def execute_wrapper(
@@ -177,6 +179,31 @@ def test_param__number_runs(self) -> None:
         gt, lt = 0.08, 0.13
         assert gt < data < lt, f'{gt} < {data} < {lt}'
 
+    def test_execute_raw_sql(self) -> None:
+        reporter_raw_sql = (
+            'SELECT "tests_reporter"."id", "tests_reporter"."full_name" '
+            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s' % self.reporter.pk
+        )
+        article_raw_sql = (
+            '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" = %s' % self.article.pk
+        )
+        obj = CaptureQueries()
+        for _ in obj:
+            list(Reporter.objects.raw(reporter_raw_sql))
+            list(Article.objects.raw(article_raw_sql))
+
+        data = obj.queries_log[0]['sql']
+        assert data == (reporter_raw_sql), data
+
+        data = obj.queries_log[1]['sql']
+        assert data == (article_raw_sql), data
+
+    def test_without_requests(self) -> None:
+        for _ in CaptureQueries(advanced_verb=True):
+            pass  # no have requests
+
 
 @pytest.mark.django_db(transaction=True)
 class TestDecoratorCaptureQueries(BasicTestsFor3ChoicesOfCaptureQueries):
@@ -230,6 +257,34 @@ def test_param__auto_call_func(self) -> None:
         gt, lt = 0.08, 0.13
         assert gt < data < lt, f'{gt} < {data} < {lt}'
 
+    def test_execute_raw_sql(self) -> None:
+        reporter_raw_sql = (
+            'SELECT "tests_reporter"."id", "tests_reporter"."full_name" '
+            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s' % self.reporter.pk
+        )
+        article_raw_sql = (
+            '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" = %s' % self.article.pk
+        )
+        obj = CaptureQueries(auto_call_func=True)
+
+        @obj
+        def _() -> None:
+            list(Reporter.objects.raw(reporter_raw_sql))
+            list(Article.objects.raw(article_raw_sql))
+
+        data = obj.queries_log[0]['sql']
+        assert data == (reporter_raw_sql), data
+
+        data = obj.queries_log[1]['sql']
+        assert data == (article_raw_sql), data
+
+    def test_without_requests(self) -> None:
+        @CaptureQueries(advanced_verb=True, auto_call_func=True)
+        def func() -> None:
+            pass  # no have requests
+
 
 @pytest.mark.django_db(transaction=True)
 class TestContextManagerCaptureQueries(BasicTestsFor3ChoicesOfCaptureQueries):
@@ -238,26 +293,50 @@ class TestContextManagerCaptureQueries(BasicTestsFor3ChoicesOfCaptureQueries):
     def call_capture_queries(self, **kwargs: Any) -> CaptureQueries:
         with slow_down_execute(0.1):  # noqa: SIM117
             with CaptureQueries(**kwargs) as obj:
-                obj.current_iteration = 1  # XXX(Ars): Временный хак для тестов, позже надо переработать
+                obj.current_iteration = 1  # XXX(Ars): A temporary hack only for tests, reworked later
                 _select(self.reporter.pk, self.article.pk)
         return obj
 
-    # @pytest.mark.filterwarnings("ignore::UserWarning")  # warn не отображается, и не вызывает ошибки
-    # @pytest.mark.filterwarnings('default::UserWarning')  # warn отображается, и не вызывает ошибки
+    # @pytest.mark.filterwarnings("ignore::UserWarning")  # warn not show, and not raise exc
+    # @pytest.mark.filterwarnings('default::UserWarning')  # warn show, and not raise exc
     def test_param__number_runs(self) -> None:
         with pytest.raises(
             UserWarning,
             match=(
-                'При использовании: CaptureQueries как контекстного менеджера, '
-                'параметр number_runs > 1 не используеться.'
+                'When using: CaptureQueries as a context manager,'
+                ' the number_runs > 1 parameter is not used.'
             ),
         ):
             self.call_capture_queries(number_runs=3)
 
+    def test_execute_raw_sql(self) -> None:
+        reporter_raw_sql = (
+            'SELECT "tests_reporter"."id", "tests_reporter"."full_name" '
+            'FROM "tests_reporter" WHERE "tests_reporter"."id" = %s' % self.reporter.pk
+        )
+        article_raw_sql = (
+            '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" = %s' % self.article.pk
+        )
+        with CaptureQueries() as obj:
+            list(Reporter.objects.raw(reporter_raw_sql))
+            list(Article.objects.raw(article_raw_sql))
+
+        data = obj.queries_log[0]['sql']
+        assert data == (reporter_raw_sql), data
+
+        data = obj.queries_log[1]['sql']
+        assert data == (article_raw_sql), data
+
+    def test_without_requests(self) -> None:
+        with CaptureQueries(advanced_verb=True):
+            pass  # no have requests
+
 
 @pytest.mark.django_db(transaction=True)
 class TestOutputCaptureQueries:
-    """Обязательно должен быть передан аргумент -s при запуске pytest, для вывода output"""
+    """The -s argument must be passed when running py test to output output"""
 
     # @pytest.mark.usefixtures('_debug_true')
     def test_capture_queries_loop(self, intercept_output: StringIO) -> None:
@@ -271,7 +350,7 @@ def test_capture_queries_loop(self, intercept_output: StringIO) -> None:
         assert re.match(
             f'\n\nTests count: 100  |  Total queries count: 200  |  Total execution time: {ANYNUM}s  |  Median time one test is: {ANYNUM}s\n',  # noqa: E501
             output,
-        ), 'incorrect output'
+        ), f'incorrect output = {output}'
 
     # @pytest.mark.usefixtures('_debug_true')
     def test_capture_queries_decorator(self, intercept_output: StringIO) -> None:
@@ -286,7 +365,7 @@ def _() -> None:
         assert re.match(
             f'\n\nTests count: 100  |  Total queries count: 200  |  Total execution time: {ANYNUM}s  |  Median time one test is: {ANYNUM}s\n',  # noqa: E501
             output,
-        ), 'incorrect output'
+        ), f'incorrect output = {output}'
 
     # @pytest.mark.usefixtures('_debug_true')
     def test_capture_queries_context_manager(self, intercept_output: StringIO) -> None:
@@ -296,9 +375,9 @@ def test_capture_queries_context_manager(self, intercept_output: StringIO) -> No
         output = intercept_output.getvalue()
 
         assert re.match(
-            f'Queries count: 2  |  Execution time: {ANYNUM}s',
+            f'\nQueries count: 2  |  Execution time: {ANYNUM}s  |  Vendor: sqlite\n\n',
             output,
-        ), 'incorrect output'
+        ), f'incorrect output = {output}'
 
 
 @pytest.mark.django_db(transaction=True)