Coverage for mindsdb / utilities / hooks / profiling.py: 26%

1import os 

2import json 

3 

4import psycopg 

5 

6from mindsdb.utilities import log 

7 

8logger = log.getLogger(__name__) 

9 

10 

11MINDSDB_PROFILING_ENABLED = os.environ.get("MINDSDB_PROFILING_ENABLED") in ("1", "true") 

12MINDSDB_PROFILING_DB_HOST = os.environ.get("MINDSDB_PROFILING_DB_HOST") 

13MINDSDB_PROFILING_DB_USER = os.environ.get("MINDSDB_PROFILING_DB_USER") 

14MINDSDB_PROFILING_DB_PASSWORD = os.environ.get("MINDSDB_PROFILING_DB_PASSWORD") 

15 

16 

17def set_level(node, level, internal_id): 

18 internal_id["id"] += 1 

19 node["level"] = level 

20 node["value"] = node["stop_at"] - node["start_at"] 

21 node["value_thread"] = node["stop_at_thread"] - node["start_at_thread"] 

22 node["value_process"] = node["stop_at_process"] - node["start_at_process"] 

23 node["internal_id"] = internal_id["id"] 

24 

25 accum = 0 

26 for child_node in node["children"]: 

27 set_level(child_node, level + 1, internal_id) 

28 accum += child_node["value"] 

29 node["self"] = node["value"] - accum 

30 

31 

32def send_profiling_results(profiling_data: dict): 

33 if MINDSDB_PROFILING_ENABLED is False: 

34 return 

35 

36 profiling = profiling_data 

37 set_level(profiling["tree"], 0, {"id": 0}) 

38 

39 time_start_at = profiling["tree"]["time_start_at"] 

40 del profiling["tree"]["time_start_at"] 

41 

42 try: 

43 connection = psycopg.connect( 

44 host=MINDSDB_PROFILING_DB_HOST, 

45 port=5432, 

46 user=MINDSDB_PROFILING_DB_USER, 

47 password=MINDSDB_PROFILING_DB_PASSWORD, 

48 dbname="postgres", 

49 connect_timeout=5, 

50 ) 

51 except Exception: 

52 logger.warning("cant get acceess to profiling database") 

53 return 

54 cur = connection.cursor() 

55 cur.execute( 

56 """ 

57 insert into profiling 

58 (data, query, time, hostname, environment, api, total_time, company_id, instance_id) 

59 values 

60 (%s, %s, %s, %s, %s, %s, %s, %s, %s) 

61 """, 

62 ( 

63 json.dumps(profiling["tree"]), 

64 profiling.get("query", "?"), 

65 time_start_at, 

66 profiling["hostname"], 

67 profiling.get("environment", "?"), 

68 profiling.get("api", "?"), 

69 profiling["tree"]["value"], 

70 profiling["company_id"], 

71 profiling["instance_id"], 

72 ), 

73 ) 

74 

75 connection.commit() 

76 cur.close() 

77 connection.close()