Merge pull request #174 from SectorLabs/181799346-caller-sql-comment

Append the caller data to each SQL query

Author: sewi-cpan

manage.py

@@ -1,11 +1,25 @@
-from importlib import import_module
+from psqlextra.compiler import (
+    PostgresAggregateCompiler,
+    PostgresCompiler,
+    PostgresDeleteCompiler,
+    PostgresInsertCompiler,
+    PostgresUpdateCompiler,
+)
 
 from . import base_impl
 
 
 class PostgresOperations(base_impl.operations()):
     """Simple operations specific to PostgreSQL."""
 
+    compiler_map = {
+        "SQLCompiler": PostgresCompiler,
+        "SQLInsertCompiler": PostgresInsertCompiler,
+        "SQLUpdateCompiler": PostgresUpdateCompiler,
+        "SQLDeleteCompiler": PostgresDeleteCompiler,
+        "SQLAggregateCompiler": PostgresAggregateCompiler,
+    }
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
@@ -14,14 +28,9 @@ def __init__(self, *args, **kwargs):
     def compiler(self, compiler_name: str):
         """Gets the SQL compiler with the specified name."""
 
-        # first let django try to find the compiler
-        try:
-            return super().compiler(compiler_name)
-        except AttributeError:
-            pass
-
-        # django can't find it, look in our own module
-        if self._compiler_cache is None:
-            self._compiler_cache = import_module("psqlextra.compiler")
+        postgres_compiler = self.compiler_map.get(compiler_name)
+        if postgres_compiler:
+            return postgres_compiler
 
-        return getattr(self._compiler_cache, compiler_name)
+        # Let Django try to find the compiler. Better run without caller comment than break
+        return super().compiler(compiler_name)

psqlextra/backend/operations.py

@@ -1,18 +1,100 @@
+import inspect
+import os
+import sys
+
 from collections.abc import Iterable
 from typing import Tuple, Union
 
 import django
 
+from django.conf import settings
 from django.core.exceptions import SuspiciousOperation
 from django.db.models import Expression, Model, Q
 from django.db.models.fields.related import RelatedField
-from django.db.models.sql.compiler import SQLInsertCompiler, SQLUpdateCompiler
+from django.db.models.sql.compiler import (
+    SQLAggregateCompiler,
+    SQLCompiler,
+    SQLDeleteCompiler,
+    SQLInsertCompiler,
+    SQLUpdateCompiler,
+)
 from django.db.utils import ProgrammingError
 
 from .expressions import HStoreValue
 from .types import ConflictAction
 
 
+def append_caller_to_sql(sql):
+    """Append the caller to SQL queries.
+
+    Adds the calling file and function as an SQL comment to each query.
+    Examples:
+     INSERT INTO "tests_47ee19d1" ("id", "title")
+     VALUES (1, 'Test')
+     RETURNING "tests_47ee19d1"."id"
+     /* 998020 test_append_caller_to_sql_crud .../django-postgres-extra/tests/test_append_caller_to_sql.py 55 */
+
+     SELECT "tests_47ee19d1"."id", "tests_47ee19d1"."title"
+     FROM "tests_47ee19d1"
+     WHERE "tests_47ee19d1"."id" = 1
+     LIMIT 1
+     /* 998020 test_append_caller_to_sql_crud .../django-postgres-extra/tests/test_append_caller_to_sql.py 69 */
+
+     UPDATE "tests_47ee19d1"
+     SET "title" = 'success'
+     WHERE "tests_47ee19d1"."id" = 1
+     /* 998020 test_append_caller_to_sql_crud .../django-postgres-extra/tests/test_append_caller_to_sql.py 64 */
+
+     DELETE FROM "tests_47ee19d1"
+     WHERE "tests_47ee19d1"."id" IN (1)
+     /* 998020 test_append_caller_to_sql_crud .../django-postgres-extra/tests/test_append_caller_to_sql.py 74 */
+
+    Slow and blocking queries could be easily tracked down to their originator
+    within the source code using the "pg_stat_activity" table.
+
+    Enable "PSQLEXTRA_ANNOTATE_SQL" within the database settings to enable this feature.
+    """
+
+    if not getattr(settings, "PSQLEXTRA_ANNOTATE_SQL", None):
+        return sql
+
+    try:
+        # Search for the first non-Django caller
+        stack = inspect.stack()
+        for stack_frame in stack[1:]:
+            frame_filename = stack_frame[1]
+            frame_line = stack_frame[2]
+            frame_function = stack_frame[3]
+            if "/django/" in frame_filename or "/psqlextra/" in frame_filename:
+                continue
+
+            return f"{sql} /* {os.getpid()} {frame_function} {frame_filename} {frame_line} */"
+
+        # Django internal commands (like migrations) end up here
+        return f"{sql} /* {os.getpid()} {sys.argv[0]} */"
+    except Exception:
+        # Don't break anything because this convinence function runs into an unexpected situation
+        return sql
+
+
+class PostgresCompiler(SQLCompiler):
+    def as_sql(self):
+        sql, params = super().as_sql()
+        return append_caller_to_sql(sql), params
+
+
+class PostgresDeleteCompiler(SQLDeleteCompiler):
+    def as_sql(self):
+        sql, params = super().as_sql()
+        return append_caller_to_sql(sql), params
+
+
+class PostgresAggregateCompiler(SQLAggregateCompiler):
+    def as_sql(self):
+        sql, params = super().as_sql()
+        return append_caller_to_sql(sql), params
+
+
 class PostgresUpdateCompiler(SQLUpdateCompiler):
     """Compiler for SQL UPDATE statements that allows us to use expressions
     inside HStore values.
@@ -24,7 +106,8 @@ class PostgresUpdateCompiler(SQLUpdateCompiler):
 
     def as_sql(self):
         self._prepare_query_values()
-        return super().as_sql()
+        sql, params = super().as_sql()
+        return append_caller_to_sql(sql), params
 
     def _prepare_query_values(self):
         """Extra prep on query values by converting dictionaries into.
@@ -72,15 +155,27 @@ def _does_dict_contain_expression(data: dict) -> bool:
 class PostgresInsertCompiler(SQLInsertCompiler):
     """Compiler for SQL INSERT statements."""
 
-    def __init__(self, *args, **kwargs):
-        """Initializes a new instance of :see:PostgresInsertCompiler."""
+    def as_sql(self, return_id=False):
+        """Builds the SQL INSERT statement."""
+        queries = [
+            (append_caller_to_sql(sql), params)
+            for sql, params in super().as_sql()
+        ]
+
+        return queries
 
+
+class PostgresInsertOnConflictCompiler(SQLInsertCompiler):
+    """Compiler for SQL INSERT statements."""
+
+    def __init__(self, *args, **kwargs):
+        """Initializes a new instance of
+        :see:PostgresInsertOnConflictCompiler."""
         super().__init__(*args, **kwargs)
         self.qn = self.connection.ops.quote_name
 
     def as_sql(self, return_id=False):
         """Builds the SQL INSERT statement."""
-
         queries = [
             self._rewrite_insert(sql, params, return_id)
             for sql, params in super().as_sql()
@@ -132,10 +227,12 @@ def _rewrite_insert(self, sql, params, return_id=False):
             self.qn(self.query.model._meta.pk.attname) if return_id else "*"
         )
 
-        return self._rewrite_insert_on_conflict(
+        (sql, params) = self._rewrite_insert_on_conflict(
             sql, params, self.query.conflict_action.value, returning
         )
 
+        return append_caller_to_sql(sql), params
+
     def _rewrite_insert_on_conflict(
         self, sql, params, conflict_action: ConflictAction, returning
     ):

psqlextra/compiler.py

@@ -8,7 +8,7 @@
 from django.db.models import sql
 from django.db.models.constants import LOOKUP_SEP
 
-from .compiler import PostgresInsertCompiler, PostgresUpdateCompiler
+from .compiler import PostgresInsertOnConflictCompiler, PostgresUpdateCompiler
 from .expressions import HStoreColumn
 from .fields import HStoreField
 from .types import ConflictAction
@@ -179,7 +179,7 @@ def values(self, objs: List, insert_fields: List, update_fields: List = []):
     def get_compiler(self, using=None, connection=None):
         if using:
             connection = connections[using]
-        return PostgresInsertCompiler(self, connection, using)
+        return PostgresInsertOnConflictCompiler(self, connection, using)
 
 
 class PostgresUpdateQuery(sql.UpdateQuery):

psqlextra/sql.py

@@ -0,0 +1,81 @@
+import pytest
+
+from django.db import connection, models
+from django.test.utils import CaptureQueriesContext, override_settings
+
+from psqlextra.compiler import append_caller_to_sql
+
+from .fake_model import get_fake_model
+
+
+class psqlextraSimulated:
+    def callMockedClass(self):
+        return MockedClass().mockedMethod()
+
+
+class MockedClass:
+    def mockedMethod(self):
+        return append_caller_to_sql("sql")
+
+
+def mockedFunction():
+    return append_caller_to_sql("sql")
+
+
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=False)
+def test_disable_append_caller_to_sql():
+    commented_sql = mockedFunction()
+    assert commented_sql == "sql"
+
+
+@pytest.mark.parametrize(
+    "entry_point",
+    [
+        MockedClass().mockedMethod,
+        psqlextraSimulated().callMockedClass,
+    ],
+)
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
+def test_append_caller_to_sql_class(entry_point):
+    commented_sql = entry_point()
+    assert commented_sql.startswith("sql /* ")
+    assert "mockedMethod" in commented_sql
+    assert __file__ in commented_sql
+
+
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
+def test_append_caller_to_sql_function():
+    commented_sql = mockedFunction()
+    assert commented_sql.startswith("sql /* ")
+    assert "mockedFunction" in commented_sql
+    assert __file__ in commented_sql
+
+
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
+def test_append_caller_to_sql_crud():
+    model = get_fake_model(
+        {
+            "title": models.CharField(max_length=255, null=True),
+        }
+    )
+
+    obj = None
+    with CaptureQueriesContext(connection) as queries:
+        obj = model.objects.create(
+            id=1,
+            title="Test",
+        )
+        assert "test_append_caller_to_sql_crud " in queries[0]["sql"]
+
+    obj.title = "success"
+    with CaptureQueriesContext(connection) as queries:
+        obj.save()
+        assert "test_append_caller_to_sql_crud " in queries[0]["sql"]
+
+    with CaptureQueriesContext(connection) as queries:
+        assert model.objects.filter(id=obj.id)[0].id == obj.id
+        assert "test_append_caller_to_sql_crud " in queries[0]["sql"]
+
+    with CaptureQueriesContext(connection) as queries:
+        obj.delete()
+        assert "test_append_caller_to_sql_crud " in queries[0]["sql"]

tests/test_append_caller_to_sql.py

@@ -3,6 +3,7 @@
 
 from django.core.exceptions import SuspiciousOperation
 from django.db import connection, models
+from django.test.utils import CaptureQueriesContext, override_settings
 from django.utils import timezone
 
 from psqlextra.fields import HStoreField
@@ -13,6 +14,7 @@
 
 
 @pytest.mark.parametrize("conflict_action", ConflictAction.all())
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
 def test_on_conflict(conflict_action):
     """Tests whether simple inserts work correctly."""
 
@@ -23,9 +25,11 @@ def test_on_conflict(conflict_action):
         }
     )
 
-    obj = model.objects.on_conflict(
-        [("title", "key1")], conflict_action
-    ).insert_and_get(title={"key1": "beer"}, cookies="cheers")
+    with CaptureQueriesContext(connection) as queries:
+        obj = model.objects.on_conflict(
+            [("title", "key1")], conflict_action
+        ).insert_and_get(title={"key1": "beer"}, cookies="cheers")
+        assert " test_on_conflict " in queries[0]["sql"]
 
     model.objects.on_conflict(
         [("title", "key1")], conflict_action

tests/test_on_conflict.py

@@ -1,5 +1,6 @@
-from django.db import models
+from django.db import connection, models
 from django.db.models import Case, F, Q, Value, When
+from django.test.utils import CaptureQueriesContext, override_settings
 
 from psqlextra.expressions import HStoreRef
 from psqlextra.fields import HStoreField
@@ -134,3 +135,21 @@ def test_query_hstore_value_update_escape():
 
     inst = model.objects.all().first()
     assert inst.title.get("en") == "console.log('test')"
+
+
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
+def test_query_comment():
+    """Tests whether the query is commented."""
+
+    model = get_fake_model(
+        {
+            "name": models.CharField(max_length=10),
+            "value": models.IntegerField(),
+        }
+    )
+
+    with CaptureQueriesContext(connection) as queries:
+        qs = model.objects.all()
+        assert " test_query_comment " in str(qs.query)
+        list(qs)
+        assert " test_query_comment " in queries[0]["sql"]

tests/test_query.py

@@ -2,6 +2,7 @@
 
 from django.core.exceptions import ImproperlyConfigured
 from django.db import models
+from django.test.utils import override_settings
 
 from psqlextra.models import PostgresMaterializedViewModel, PostgresViewModel
 
@@ -11,6 +12,7 @@
 @pytest.mark.parametrize(
     "model_base", [PostgresViewModel, PostgresMaterializedViewModel]
 )
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
 def test_view_model_meta_query_set(model_base):
     """Tests whether you can set a :see:QuerySet to be used as the underlying
     query for a view."""
@@ -26,7 +28,8 @@ def test_view_model_meta_query_set(model_base):
     expected_sql = 'SELECT "{0}"."id", "{0}"."name" FROM "{0}"'.format(
         model._meta.db_table
     )
-    assert view_model._view_meta.query == (expected_sql, tuple())
+    assert view_model._view_meta.query[0].startswith(expected_sql + " /* ")
+    assert view_model._view_meta.query[1] == tuple()
 
 
 @pytest.mark.parametrize(