Coverage for django_query_capture/capture.py: 100%

1""" 

2This module is a code that expands and holds data by hooking whenever a query occurs in django. 

3""" 

4 

5import typing 

6 

7import inspect 

8import time 

9from contextlib import ContextDecorator, ExitStack 

10from distutils.sysconfig import get_python_lib 

11 

12from django.db import connection 

13from django.db.backends.dummy.base import DatabaseWrapper 

14from django.db.backends.utils import CursorWrapper 

15 

16 

17class CapturedQueryContext(typing.TypedDict): 

18 """ 

19 It is a `data class` that captures the data that appears as Context when you capture Query in django. 

20 """ 

21 

22 connection: DatabaseWrapper 

23 cursor: CursorWrapper 

24 

25 

26class CapturedQuery(typing.TypedDict): 

27 """ 

28 A `data class` that adds the time and place of occurrence to the data that comes out when you capture Query in django. 

29 """ 

30 

31 sql: str 

32 raw_sql: str 

33 raw_params: str 

34 many: bool 

35 duration: float 

36 file_name: str 

37 function_name: str 

38 line_no: int 

39 context: CapturedQueryContext 

40 

41 

42class native_query_capture(ContextDecorator): 

43 """ 

44 This is the `ContextDecorator` that extends django's `connection.execute_wrapper`.<br> 

45 measure the time of the query and guess where the query occurred.<br> 

46 the main attribute is [self.captured_queries][capture.CapturedQuery], [native_query_capture][capture.native_query_capture] returns data from some captured_queries. 

47 """ 

48 

49 def __init__(self): 

50 """ 

51 `self._exit_stack`: `ExitStack` was used to wrap `connection.execute_wrapper`.<br> 

52 `self.captured_queries`: Used to store captured queries and expanded data. 

53 """ 

54 self._exit_stack = ExitStack().__enter__() 

55 self.captured_queries: typing.List[CapturedQuery] = [] 

56 

57 def __enter__(self) -> "native_query_capture": 

58 """ 

59 Use exit_stack to perform `connection.execute_wrapper.__enter__`. 

60 

61 Returns: 

62 Returns yourself with the property of [self.captured_queries][capture.CapturedQuery] so that you can check captured queries in real time. 

63 """ 

64 self._exit_stack.enter_context(connection.execute_wrapper(self._save_queries)) 

65 return self 

66 

67 def __exit__(self, exc_type, exc_value, traceback): 

68 """ 

69 Close exit_stack to call `connection.execute_wrapper.__exit___`. 

70 """ 

71 self._exit_stack.close() 

72 

73 def __len__(self) -> int: 

74 """ 

75 Returns: 

76 Returns the length of [self.captured_queries][capture.CapturedQuery]. 

77 """ 

78 return len(self.captured_queries) 

79 

80 def _save_queries(self, execute, sql, params, many, context): 

81 """ 

82 https://docs.djangoproject.com/en/3.2/topics/db/instrumentation/ 

83 

84 It is a function used as a callback for execute_wrapper and receives a factor of `connection.execute_wrapper`.<br> 

85 measure the time of the query with the data provided by `connection.execute_wrapper`, track and store the query-generated CallStack. 

86 

87 Args: 

88 execute: a callable, which should be invoked with the rest of the parameters in order to execute the query. 

89 sql: a str, the SQL query to be sent to the database. 

90 params: a list/tuple of parameter values for the SQL command, or a list/tuple of lists/tuples if the wrapped call is executemany(). 

91 many: a bool indicating whether the ultimately invoked call is execute() or executemany() (and whether params is expected to be a sequence of values, or a sequence of sequences of values). 

92 context: a dictionary with further data about the context of invocation. This includes the connection and cursor. 

93 

94 Returns: 

95 Returns the result of the exit for the basic operation of `connection.execute_wrapper`. 

96 """ 

97 python_library_directory = get_python_lib() 

98 called_by = [ 

99 stack 

100 for stack in inspect.stack() 

101 if not stack.filename.startswith(python_library_directory) 

102 ][0] 

103 file_name = called_by.filename 

104 function_name = called_by.function 

105 line_no = called_by.lineno 

106 start_timestamp = time.monotonic() 

107 result = execute(sql, params, many, context) 

108 duration = time.monotonic() - start_timestamp 

109 self.captured_queries.append( 

110 { 

111 "sql": sql % tuple(params) if params else sql, 

112 "raw_sql": sql, 

113 "raw_params": params, 

114 "many": many, 

115 "duration": duration, 

116 "file_name": file_name, 

117 "function_name": function_name, 

118 "line_no": line_no, 

119 "context": context, 

120 } 

121 ) 

122 

123 return result