fix(django): handle composed queries in postgres (#2463) (#2513)

* fix(django): handle composed queries in postgres

* fix patching

* nit

* black

* ducktype

Co-authored-by: Brett Langdon <[email protected]>
(cherry picked from commit af326b6)

Co-authored-by: Tahir H. Butt <[email protected]>
Co-authored-by: Brett Langdon <[email protected]>

Author: 3 people

ddtrace/contrib/django/patch.py

@@ -19,6 +19,16 @@
 from ddtrace.constants import SPAN_MEASURED_KEY
 from ddtrace.contrib import dbapi
 from ddtrace.contrib import func_name
+
+
+try:
+    from psycopg2._psycopg import cursor as psycopg_cursor_cls
+
+    from ddtrace.contrib.psycopg.patch import Psycopg2TracedCursor
+except ImportError:
+    psycopg_cursor_cls = None
+    Psycopg2TracedCursor = None
+
 from ddtrace.ext import SpanTypes
 from ddtrace.ext import http
 from ddtrace.ext import sql as sqlx
@@ -82,7 +92,15 @@ def cursor(django, pin, func, instance, args, kwargs):
             "django.db.alias": alias,
         }
         pin = Pin(service, tags=tags, tracer=pin.tracer, app=prefix)
-        return dbapi.TracedCursor(func(*args, **kwargs), pin, config.django)
+        cursor = func(*args, **kwargs)
+        traced_cursor_cls = dbapi.TracedCursor
+        if (
+            Psycopg2TracedCursor is not None
+            and hasattr(cursor, "cursor")
+            and isinstance(cursor.cursor, psycopg_cursor_cls)
+        ):
+            traced_cursor_cls = Psycopg2TracedCursor
+        return traced_cursor_cls(cursor, pin, config.django)
 
     if not isinstance(conn.cursor, wrapt.ObjectProxy):
         conn.cursor = wrapt.FunctionWrapper(conn.cursor, trace_utils.with_traced_module(cursor)(django))

docs/spelling_wordlist.txt

@@ -28,6 +28,7 @@ CGroup
 cgroups
 cherrypy
 compat
+composable
 config
 coroutine
 coroutines
@@ -104,6 +105,7 @@ respawn
 runnable
 runtime
 sanic
+sql
 sqlalchemy
 sqlite
 starlette

releasenotes/notes/fix-django-db-postgres-7aa62a1c3aff647d.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    django: fix a bug when postgres query is composable sql object.

tests/contrib/django/test_django_snapshots.py

@@ -68,3 +68,38 @@ def test_safe_string_encoding(client):
 )
 def test_404_exceptions(client):
     assert client.get("/404-view/").status_code == 404
+
+
+@pytest.fixture()
+def psycopg2_patched(transactional_db):
+    from django.db import connections
+
+    from ddtrace.contrib.psycopg.patch import patch
+    from ddtrace.contrib.psycopg.patch import unpatch
+
+    patch()
+
+    # # force recreate connection to ensure psycopg2 patching has occurred
+    del connections["postgres"]
+    connections["postgres"].close()
+    connections["postgres"].connect()
+
+    yield
+
+    unpatch()
+
+
+@snapshot(ignores=["meta.out.host"])
+@pytest.mark.django_db
+def test_psycopg_query_default(client, psycopg2_patched):
+    """Execute a psycopg2 query on a Django database wrapper"""
+    from django.db import connections
+    from psycopg2.sql import SQL
+
+    query = SQL("""select 'one' as x""")
+    conn = connections["postgres"]
+    with conn.cursor() as cur:
+        cur.execute(query)
+        rows = cur.fetchall()
+        assert len(rows) == 1, rows
+        assert rows[0][0] == "one"

tests/snapshots/tests.contrib.django.test_django_snapshots.test_psycopg_query_default.snap

@@ -0,0 +1,38 @@
+[[{"name" "postgres.query"
+   "service" "postgresdb"
+   "resource" "select 'one' as x"
+   "type" "sql"
+   "error" 0
+   "span_id" 0
+   "trace_id" 0
+   "parent_id" nil
+   "start" 1623177658188007300
+   "duration" 1828500
+   "meta" {"runtime-id" "98e7426d5d6d44d1a5d76d1322d69cac"
+           "django.db.vendor" "postgresql"
+           "django.db.alias" "postgres"}
+   "metrics" {"_dd.agent_psr" 1.0
+              "_dd.measured" 1
+              "_sampling_priority_v1" 1
+              "system.pid" 236
+              "db.rowcount" 1
+              "sql.rows" 1
+              "_dd.tracer_kr" 1.0}}
+      {"name" "postgres.query"
+       "service" "postgres"
+       "resource" "select 'one' as x"
+       "type" "sql"
+       "error" 0
+       "span_id" 1
+       "trace_id" 0
+       "parent_id" 0
+       "start" 1623177658188405100
+       "duration" 1265000
+       "meta" {"out.host" "127.0.0.1"
+               "db.name" "test_postgres"
+               "db.user" "postgres"
+               "db.application" "None"}
+       "metrics" {"_dd.measured" 1
+                  "out.port" 5432
+                  "db.rowcount" 1
+                  "sql.rows" 1}}]]