Skip to content

Commit a3e8cc0

Browse files
committed
Fix sql recording for async views
1 parent e951347 commit a3e8cc0

File tree

1 file changed

+13
-25
lines changed

1 file changed

+13
-25
lines changed

debug_toolbar/panels/sql/tracking.py

Lines changed: 13 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1+
import contextvars
12
import datetime
23
import json
3-
from threading import local
44
from time import time
55

66
from django.utils.encoding import force_str
@@ -13,30 +13,12 @@
1313
except ImportError:
1414
PostgresJson = None
1515

16+
recording = contextvars.ContextVar("debug-toolbar-recording", default=True)
17+
1618

1719
class SQLQueryTriggered(Exception):
1820
"""Thrown when template panel triggers a query"""
1921

20-
pass
21-
22-
23-
class ThreadLocalState(local):
24-
def __init__(self):
25-
self.enabled = True
26-
27-
@property
28-
def Wrapper(self):
29-
if self.enabled:
30-
return NormalCursorWrapper
31-
return ExceptionCursorWrapper
32-
33-
def recording(self, v):
34-
self.enabled = v
35-
36-
37-
state = ThreadLocalState()
38-
recording = state.recording # export function
39-
4022

4123
def wrap_cursor(connection, panel):
4224
if not hasattr(connection, "_djdt_cursor"):
@@ -50,16 +32,22 @@ def cursor(*args, **kwargs):
5032
# See:
5133
# https://github.com/jazzband/django-debug-toolbar/pull/615
5234
# https://github.com/jazzband/django-debug-toolbar/pull/896
53-
return state.Wrapper(
54-
connection._djdt_cursor(*args, **kwargs), connection, panel
55-
)
35+
if recording.get():
36+
wrapper = NormalCursorWrapper
37+
else:
38+
wrapper = ExceptionCursorWrapper
39+
return wrapper(connection._djdt_cursor(*args, **kwargs), connection, panel)
5640

5741
def chunked_cursor(*args, **kwargs):
5842
# prevent double wrapping
5943
# solves https://github.com/jazzband/django-debug-toolbar/issues/1239
6044
cursor = connection._djdt_chunked_cursor(*args, **kwargs)
6145
if not isinstance(cursor, BaseCursorWrapper):
62-
return state.Wrapper(cursor, connection, panel)
46+
if recording.get():
47+
wrapper = NormalCursorWrapper
48+
else:
49+
wrapper = ExceptionCursorWrapper
50+
return wrapper(cursor, connection, panel)
6351
return cursor
6452

6553
connection.cursor = cursor

0 commit comments

Comments
 (0)