fix(django/sql): Run cursor.mogrify as fallback to regular formatting (getsentry#325)

untitaker · web-flow · commit 46f8527e4f7b · 2019-04-06T01:11:47.000+02:00
* fix(django/sql): Run cursor.mogrify as fallback to regular formatting

* fix: Fix imports
diff --git a/sentry_sdk/integrations/django/__init__.py b/sentry_sdk/integrations/django/__init__.py
@@ -322,21 +322,38 @@ def record_sql(sql, params, cursor=None):
     if hub.get_integration(DjangoIntegration) is None:
         return
 
-    with capture_internal_exceptions():
-        if cursor and hasattr(cursor, "mogrify"):  # psycopg2
-            real_sql = cursor.mogrify(sql, params)
-            with capture_internal_exceptions():
+    real_sql = None
+    real_params = None
+
+    try:
+        # Prefer our own SQL formatting logic because it's the only one that
+        # has proper value trimming.
+        real_sql, real_params = format_sql(sql, params)
+        if real_sql:
+            real_sql = format_and_strip(real_sql, real_params)
+    except Exception:
+        pass
+
+    if not real_sql and cursor and hasattr(cursor, "mogrify"):
+        # If formatting failed and we're using psycopg2, it could be that we're
+        # looking at a query that uses Composed objects. Use psycopg2's mogrify
+        # function to format the query. We lose per-parameter trimming but gain
+        # accuracy in formatting.
+        #
+        # This is intentionally the second choice because we assume Composed
+        # queries are not widely used, while per-parameter trimming is
+        # generally highly desirable.
+        try:
+            if cursor and hasattr(cursor, "mogrify"):
+                real_sql = cursor.mogrify(sql, params)
                 if isinstance(real_sql, bytes):
                     real_sql = real_sql.decode(cursor.connection.encoding)
-        else:
-            real_sql, real_params = format_sql(sql, params)
-
-            if real_params:
-                try:
-                    real_sql = format_and_strip(real_sql, real_params)
-                except Exception:
-                    pass
-        hub.add_breadcrumb(message=real_sql, category="query")
+        except Exception:
+            pass
+
+    if real_sql:
+        with capture_internal_exceptions():
+            hub.add_breadcrumb(message=real_sql, category="query")
 
 
 def install_sql_hook():
diff --git a/tests/integrations/django/test_basic.py b/tests/integrations/django/test_basic.py
@@ -226,7 +226,7 @@ def test_sql_psycopg2_string_composition(sentry_init, capture_events, query):
     if "postgres" not in connections:
         pytest.skip("postgres tests disabled")
 
-    import psycopg2
+    import psycopg2.sql
 
     sql = connections["postgres"].cursor()
 
@@ -249,7 +249,7 @@ def test_sql_psycopg2_placeholders(sentry_init, capture_events):
     if "postgres" not in connections:
         pytest.skip("postgres tests disabled")
 
-    import psycopg2
+    import psycopg2.sql
 
     sql = connections["postgres"].cursor()
 
