Coverage for src/couchers/profiler.py: 52%

1import io 

2import logging 

3import pstats 

4from contextvars import ContextVar 

5from cProfile import Profile 

6from threading import get_ident 

7from time import perf_counter_ns 

8 

9from couchers.utils import now 

10 

11logger = logging.getLogger(__name__) 

12 

13do_profile_ctx = ContextVar("do_profile_ctx", default=False) 

14sql_queries = ContextVar("sql_queries", default=[]) 

15 

16 

17def add_sql_statement(statement, params, start, start_ns, end, end_ns): 

18 if do_profile_ctx.get(): 

19 sql_queries.get().append((statement, params, start, start_ns, end, end_ns)) 

20 

21 

22def format_ns(ns): 

23 return f"{ns/1e6:0.2f} ms" 

24 

25 

26class CouchersProfiler: 

27 def __init__(self, do_profile=False): 

28 self._do_profile = do_profile 

29 do_profile_ctx.set(do_profile) 

30 sql_queries.set([]) 

31 

32 def __enter__(self): 

33 self.report = None 

34 if self._do_profile: 

35 self._pr = Profile() 

36 self._pr.enable() 

37 self._start = now() 

38 self._start_ns = perf_counter_ns() 

39 return self 

40 

41 def __exit__(self, exc_type, exc_obj, exc_tb): 

42 if self._do_profile: 

43 self._end_ns = perf_counter_ns() 

44 self._end = now() 

45 self._pr.disable() 

46 s = io.StringIO() 

47 ps = pstats.Stats(self._pr, stream=s).sort_stats("cumulative") 

48 ps.print_stats() 

49 

50 def sql_time(item): 

51 statement, params, start, start_ns, end, end_ns = item 

52 return end_ns - start_ns 

53 

54 queries = sql_queries.get() 

55 total_sql_ns = sum(map(sql_time, queries)) 

56 

57 def format_sql(item): 

58 statement, params, start, start_ns, end, end_ns = item 

59 return f"* SQL query from {format_ns(start_ns-self._start_ns)} to {format_ns(end_ns-self._start_ns)} (duration: {format_ns(end_ns-start_ns)}), aka {start} to {end}\n\n{statement}\n\nParams: {params}" 

60 

61 sql_query_info = f"==== SQL queries ====\n\nTotal {len(queries)} queries\n\n" + "\n\n\n".join( 

62 map(format_sql, queries) 

63 ) 

64 total_ns = self._end_ns - self._start_ns 

65 basics = f"==== Run info ====\n\nStart: {self._start}\nEnd: {self._end}\nDuration: {format_ns(total_ns)}\nTotal SQL: {format_ns(total_sql_ns)}\nSQL %: {100.*total_sql_ns/total_ns:0.2f}\nSQL queries: {len(queries)}\nThread: {get_ident()}" 

66 cprofile = f"==== cProfile ====\n\n{s.getvalue()}" 

67 

68 self.report = "\n\n\n\n".join([basics, cprofile, sql_query_info]) 

69 sql_queries.set([]) 

70 do_profile_ctx.set(False)