Small review changes and hide SQL annotation behind a setting

Author: sewi-cpan

psqlextra/compiler.py

@@ -7,6 +7,7 @@
 
 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
@@ -24,12 +25,43 @@
 
 
 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()
-        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]
@@ -148,7 +180,6 @@ def as_sql(self, return_id=False):
             self._rewrite_insert(sql, params, return_id)
             for sql, params in super().as_sql()
         ]
-        print(f"InsertConflictCompiler {queries}")
 
         return queries
 
@@ -196,12 +227,11 @@ def _rewrite_insert(self, sql, params, return_id=False):
             self.qn(self.query.model._meta.pk.attname) if return_id else "*"
         )
 
-        if hasattr(self.query, "conflict_action"):
-            (sql, params) = self._rewrite_insert_on_conflict(
-                sql, params, self.query.conflict_action.value, returning
-            )
+        (sql, params) = self._rewrite_insert_on_conflict(
+            sql, params, self.query.conflict_action.value, returning
+        )
 
-        return (append_caller_to_sql(sql), params)
+        return append_caller_to_sql(sql), params
 
     def _rewrite_insert_on_conflict(
         self, sql, params, conflict_action: ConflictAction, returning

tests/test_append_caller_to_sql.py

@@ -1,7 +1,7 @@
 import pytest
 
 from django.db import connection, models
-from django.test.utils import CaptureQueriesContext
+from django.test.utils import CaptureQueriesContext, override_settings
 
 from psqlextra.compiler import append_caller_to_sql
 
@@ -22,27 +22,36 @@ 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(
         {

tests/test_on_conflict.py

@@ -3,7 +3,7 @@
 
 from django.core.exceptions import SuspiciousOperation
 from django.db import connection, models
-from django.test.utils import CaptureQueriesContext
+from django.test.utils import CaptureQueriesContext, override_settings
 from django.utils import timezone
 
 from psqlextra.fields import HStoreField
@@ -14,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."""
 

tests/test_query.py

@@ -1,6 +1,6 @@
 from django.db import connection, models
 from django.db.models import Case, F, Q, Value, When
-from django.test.utils import CaptureQueriesContext
+from django.test.utils import CaptureQueriesContext, override_settings
 
 from psqlextra.expressions import HStoreRef
 from psqlextra.fields import HStoreField
@@ -137,6 +137,7 @@ def test_query_hstore_value_update_escape():
     assert inst.title.get("en") == "console.log('test')"
 
 
+@override_settings(PSQLEXTRA_ANNOTATE_SQL=True)
 def test_query_comment():
     """Tests whether the query is commented."""
 

tests/test_view_models.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."""