Coverage for src/backend/InvenTree/InvenTree/tracing.py: 15%

1"""OpenTelemetry setup functions.""" 

2 

3import base64 

4import logging 

5from typing import Optional 

6 

7from opentelemetry import metrics, trace 

8from opentelemetry.instrumentation.django import DjangoInstrumentor 

9from opentelemetry.instrumentation.redis import RedisInstrumentor 

10from opentelemetry.instrumentation.requests import RequestsInstrumentor 

11from opentelemetry.instrumentation.sqlite3 import SQLite3Instrumentor 

12from opentelemetry.instrumentation.system_metrics import SystemMetricsInstrumentor 

13from opentelemetry.sdk import _logs as logs 

14from opentelemetry.sdk import resources 

15from opentelemetry.sdk._logs import export as logs_export 

16from opentelemetry.sdk.metrics import MeterProvider 

17from opentelemetry.sdk.metrics.export import ( 

18 ConsoleMetricExporter, 

19 PeriodicExportingMetricReader, 

20) 

21from opentelemetry.sdk.trace import TracerProvider 

22from opentelemetry.sdk.trace.export import BatchSpanProcessor, ConsoleSpanExporter 

23 

24import InvenTree.ready 

25from InvenTree.version import inventreeVersion 

26 

27TRACE_PROC = None 

28TRACE_PROV = None 

29 

30 

31def setup_tracing( 

32 endpoint: Optional[str] = None, 

33 headers: Optional[dict] = None, 

34 resources_input: Optional[dict] = None, 

35 console: bool = False, 

36 auth: Optional[dict] = None, 

37 is_http: bool = False, 

38 append_http: bool = True, 

39): # pragma: no cover 

40 """Set up tracing for the application in the current context. 

41 

42 Args: 

43 endpoint: The endpoint to send the traces to. 

44 headers: The headers to send with the traces. 

45 resources_input: The resources to send with the traces. 

46 console: Whether to output the traces to the console. 

47 auth: Dict with auth information 

48 is_http: Whether to use HTTP or gRPC for the exporter. 

49 append_http: Whether to append '/v1/traces' to the endpoint. 

50 """ 

51 if InvenTree.ready.isImportingData() or InvenTree.ready.isRunningMigrations(): 

52 return 

53 

54 if endpoint is None or headers is None: 

55 print( 

56 'Tracing endpoint or headers not specified - skipping tracing setup' 

57 ) # pragma: no cover 

58 return # pragma: no cover 

59 

60 # check if trace is already set up - if so, skip 

61 if trace.get_tracer_provider() is not None: 

62 return 

63 

64 # Logger configuration 

65 logger = logging.getLogger('inventree') 

66 

67 if resources_input is None: 

68 resources_input = {} 

69 if auth is None: 

70 auth = {} 

71 

72 # Setup the auth headers 

73 if 'basic' in auth: 

74 basic_auth = auth['basic'] 

75 if 'username' in basic_auth and 'password' in basic_auth: 

76 auth_raw = f'{basic_auth["username"]}:{basic_auth["password"]}' 

77 auth_token = base64.b64encode(auth_raw.encode('utf-8')).decode('utf-8') 

78 headers['Authorization'] = f'Basic {auth_token}' 

79 else: 

80 logger.warning('Basic auth is missing username or password') 

81 

82 # Clean up headers 

83 headers = {k: v for k, v in headers.items() if v is not None} 

84 

85 # Initialize the OTLP Resource 

86 service_name = 'Unknown' 

87 if InvenTree.ready.isInServerThread(): 

88 service_name = 'BACKEND' 

89 elif InvenTree.ready.isInWorkerThread(): 

90 service_name = 'WORKER' 

91 resource = resources.Resource( 

92 attributes={ 

93 resources.SERVICE_NAME: service_name, 

94 resources.SERVICE_NAMESPACE: 'INVENTREE', 

95 resources.SERVICE_VERSION: inventreeVersion(), 

96 **resources_input, 

97 } 

98 ) 

99 

100 # Import the OTLP exporters 

101 if is_http: 

102 from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter 

103 from opentelemetry.exporter.otlp.proto.http.metric_exporter import ( 

104 OTLPMetricExporter, 

105 ) 

106 from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( 

107 OTLPSpanExporter, 

108 ) 

109 else: 

110 from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter 

111 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import ( 

112 OTLPMetricExporter, 

113 ) 

114 from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( 

115 OTLPSpanExporter, 

116 ) 

117 

118 # Spans / Traces 

119 span_exporter = OTLPSpanExporter( 

120 headers=headers, 

121 endpoint=endpoint if not (is_http and append_http) else f'{endpoint}/v1/traces', 

122 ) 

123 trace_processor = BatchSpanProcessor(span_exporter) 

124 trace_provider = TracerProvider(resource=resource) 

125 trace.set_tracer_provider(trace_provider) 

126 trace_provider.add_span_processor(trace_processor) 

127 # For debugging purposes, export the traces to the console 

128 if console: 

129 trace_provider.add_span_processor(BatchSpanProcessor(ConsoleSpanExporter())) 

130 

131 # Metrics 

132 metric_periodic_reader = PeriodicExportingMetricReader( 

133 OTLPMetricExporter( 

134 headers=headers, 

135 endpoint=endpoint 

136 if not (is_http and append_http) 

137 else f'{endpoint}/v1/metrics', 

138 ) 

139 ) 

140 metric_readers = [metric_periodic_reader] 

141 

142 # For debugging purposes, export the metrics to the console 

143 if console: 

144 console_metric_exporter = ConsoleMetricExporter() 

145 console_metric_reader = PeriodicExportingMetricReader(console_metric_exporter) 

146 metric_readers.append(console_metric_reader) 

147 

148 meter_provider = MeterProvider(resource=resource, metric_readers=metric_readers) 

149 metrics.set_meter_provider(meter_provider) 

150 

151 # Logs 

152 log_exporter = OTLPLogExporter( 

153 headers=headers, 

154 endpoint=endpoint if not (is_http and append_http) else f'{endpoint}/v1/logs', 

155 ) 

156 log_provider = logs.LoggerProvider(resource=resource) 

157 log_provider.add_log_record_processor( 

158 logs_export.BatchLogRecordProcessor(log_exporter) 

159 ) 

160 handler = logs.LoggingHandler(level=logging.INFO, logger_provider=log_provider) 

161 logger = logging.getLogger('inventree') 

162 logger.addHandler(handler) 

163 

164 global TRACE_PROC, TRACE_PROV 

165 TRACE_PROC = trace_processor 

166 TRACE_PROV = trace_provider 

167 

168 

169def setup_instruments(db_engine: str): # pragma: no cover 

170 """Run auto-instrumentation for OpenTelemetry tracing.""" 

171 if not DjangoInstrumentor()._is_instrumented_by_opentelemetry: 

172 DjangoInstrumentor().instrument() 

173 RedisInstrumentor().instrument() 

174 RequestsInstrumentor().instrument() 

175 SystemMetricsInstrumentor().instrument() 

176 

177 db_engine = str(db_engine).lower().strip() 

178 

179 # DBs 

180 if 'sqlite' in db_engine: 

181 inst = SQLite3Instrumentor() 

182 if not inst._is_instrumented_by_opentelemetry: 

183 inst.instrument() 

184 elif 'postgresql' in db_engine: 

185 try: 

186 from opentelemetry.instrumentation.psycopg import PsycopgInstrumentor 

187 

188 inst = PsycopgInstrumentor() 

189 if not inst._is_instrumented_by_opentelemetry: 

190 inst.instrument(enable_commenter=False, commenter_options={}) 

191 except ModuleNotFoundError: 

192 pass 

193 elif 'mysql' in db_engine: 

194 try: 

195 from opentelemetry.instrumentation.pymysql import PyMySQLInstrumentor 

196 

197 inst = PyMySQLInstrumentor() 

198 if not inst._is_instrumented_by_opentelemetry: 

199 inst.instrument() 

200 except ModuleNotFoundError: 

201 pass