Append the caller data to each SQL query

This PR appends an SQL comment to each query containing the Python function/file/line
calling Django to run that command.

We experienced a busy database server blocked by one very expensive query. Another
situation caused too many connections to the database server. Identifying the troublemaker
within the source code wasn't easy.

This PR adds the caller of each query to `SELECT query FROM pg_stat_activity` as a SQL comment:
```
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 */
```
The comment contains:

1. The system process id (PID). `pg_stat_activity` has a column `client_addr`. The server IP and the process id let you track down the originating process, even if PgBouncer is used inbetween.

2. The function name of the last function or method calling Django.

3. The file name and path of the source file where the call to Django happened.

4. The line number within the source file.

Djangos internal commands show the process ID and name for easy identification of migrations and other maintenance operations.

[#181799346]

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,3 +1,7 @@
+import inspect
+import os
+import sys
+
 from collections.abc import Iterable
 from typing import Tuple, Union
 
@@ -6,13 +10,59 @@
 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):
+    try:
+        # Search for the first non-Django caller
+        stack = inspect.stack()
+        i = 0
+        for stack_frame in stack[1:]:
+            i += 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 +74,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,19 +123,32 @@ 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()
         ]
+        print(f"InsertConflictCompiler {queries}")
 
         return queries
 
@@ -132,9 +196,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.query.conflict_action.value, returning
-        )
+        if hasattr(self.query, "conflict_action"):
+            (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,72 @@
+import pytest
+
+from django.db import connection, models
+from django.test.utils import CaptureQueriesContext
+
+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")
+
+
+@pytest.mark.parametrize(
+    "entry_point",
+    [
+        MockedClass().mockedMethod,
+        psqlextraSimulated().callMockedClass,
+    ],
+)
+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
+
+
+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
+
+
+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
 from django.utils import timezone
 
 from psqlextra.fields import HStoreField
@@ -23,9 +24,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
 
 from psqlextra.expressions import HStoreRef
 from psqlextra.fields import HStoreField
@@ -134,3 +135,20 @@ def test_query_hstore_value_update_escape():
 
     inst = model.objects.all().first()
     assert inst.title.get("en") == "console.log('test')"
+
+
+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

@@ -26,7 +26,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(