1from __future__ import annotations
2
3import json
4
5import click
6
7from plain.cli import register_cli
8from plain.observer.models import Span, Trace
9
10
11@register_cli("observer")
12@click.group("observer")
13def observer_cli() -> None:
14 """Observability and tracing tools"""
15
16
17@observer_cli.command()
18@click.option("--force", is_flag=True, help="Skip confirmation prompt.")
19def clear(force: bool) -> None:
20 """Clear all observer data"""
21 query = Trace.query.all()
22 trace_count = query.count()
23
24 if trace_count == 0:
25 click.echo("No traces to clear.")
26 return
27
28 if not force:
29 confirm_msg = f"Are you sure you want to clear {trace_count} trace(s)? This will delete all observer data."
30 click.confirm(confirm_msg, abort=True)
31
32 deleted_count, _ = query.delete()
33 click.secho(f"✓ Cleared {deleted_count} traces and spans", fg="green")
34
35
36@observer_cli.command("traces")
37@click.option("--limit", default=20, help="Number of traces to show (default: 20)")
38@click.option("--user-id", help="Filter by user ID")
39@click.option("--request-id", help="Filter by request ID")
40@click.option("--session-id", help="Filter by session ID")
41@click.option("--json", "output_json", is_flag=True, help="Output as JSON")
42def trace_list(
43 limit: int,
44 user_id: str | None,
45 request_id: str | None,
46 session_id: str | None,
47 output_json: bool,
48) -> None:
49 """List recent traces"""
50 # Build query
51 query = Trace.query.all()
52
53 if user_id:
54 query = query.filter(user_id=user_id)
55 if request_id:
56 query = query.filter(request_id=request_id)
57 if session_id:
58 query = query.filter(session_id=session_id)
59
60 # Limit results
61 traces = list(query[:limit])
62
63 if not traces:
64 click.echo("No traces found.")
65 return
66
67 if output_json:
68 # Output as JSON array
69 output = []
70 for trace in traces:
71 output.append(
72 {
73 "trace_id": trace.trace_id,
74 "start_time": trace.start_time.isoformat(),
75 "end_time": trace.end_time.isoformat(),
76 "duration_ms": trace.duration_ms(),
77 "request_id": trace.request_id,
78 "user_id": trace.user_id,
79 "session_id": trace.session_id,
80 "root_span_name": trace.root_span_name,
81 "summary": trace.summary,
82 }
83 )
84 click.echo(json.dumps(output, indent=2))
85 else:
86 # Table-like output
87 click.secho(
88 f"Recent traces (showing {len(traces)} of {query.count()} total):",
89 fg="bright_blue",
90 bold=True,
91 )
92 click.echo()
93
94 # Headers
95 headers = [
96 "Trace ID",
97 "Start Time",
98 "Summary",
99 "Root Span",
100 "Request ID",
101 "User ID",
102 "Session ID",
103 ]
104 col_widths = [41, 21, 31, 31, 22, 11, 22]
105
106 # Print headers
107 header_line = ""
108 for header, width in zip(headers, col_widths):
109 header_line += header.ljust(width)
110 click.secho(header_line, bold=True)
111 click.echo("-" * sum(col_widths))
112
113 # Print traces
114 for trace in traces:
115 row = [
116 trace.trace_id[:37] + "..."
117 if len(trace.trace_id) > 40
118 else trace.trace_id,
119 trace.start_time.strftime("%Y-%m-%d %H:%M:%S"),
120 trace.summary[:27] + "..."
121 if len(trace.summary) > 30
122 else trace.summary,
123 trace.root_span_name[:27] + "..."
124 if len(trace.root_span_name) > 30
125 else trace.root_span_name,
126 trace.request_id[:18] + "..."
127 if len(trace.request_id) > 20
128 else trace.request_id,
129 trace.user_id[:10],
130 trace.session_id[:18] + "..."
131 if len(trace.session_id) > 20
132 else trace.session_id,
133 ]
134
135 row_line = ""
136 for value, width in zip(row, col_widths):
137 row_line += str(value).ljust(width)
138 click.echo(row_line)
139
140
141@observer_cli.command("trace")
142@click.argument("trace_id")
143@click.option("--json", "output_json", is_flag=True, help="Output as JSON")
144def trace_detail(trace_id: str, output_json: bool) -> None:
145 """Show detailed trace information"""
146 try:
147 trace = Trace.query.get(trace_id=trace_id)
148 except Trace.DoesNotExist:
149 click.secho(f"Error: Trace with ID '{trace_id}' not found", fg="red", err=True)
150 raise click.Abort()
151
152 if output_json:
153 click.echo(json.dumps(trace.as_dict(), indent=2))
154 else:
155 click.echo(format_trace_output(trace))
156
157
158@observer_cli.command("spans")
159@click.option("--trace-id", help="Filter by trace ID")
160@click.option("--limit", default=50, help="Number of spans to show (default: 50)")
161@click.option("--json", "output_json", is_flag=True, help="Output as JSON")
162def span_list(trace_id: str | None, limit: int, output_json: bool) -> None:
163 """List recent spans"""
164 # Build query
165 query = Span.query.all()
166
167 if trace_id:
168 query = query.filter(trace__trace_id=trace_id)
169
170 # Limit results
171 spans = list(query[:limit])
172
173 if not spans:
174 click.echo("No spans found.")
175 return
176
177 if output_json:
178 # Output as JSON array
179 output = []
180 for span in spans:
181 output.append(
182 {
183 "span_id": span.span_id,
184 "trace_id": span.trace.trace_id,
185 "name": span.name,
186 "kind": span.kind,
187 "parent_id": span.parent_id,
188 "start_time": span.start_time.isoformat(),
189 "end_time": span.end_time.isoformat(),
190 "duration_ms": span.duration_ms(),
191 "status": span.status,
192 }
193 )
194 click.echo(json.dumps(output, indent=2))
195 else:
196 # Table-like output
197 click.secho(
198 f"Recent spans (showing {len(spans)} of {query.count()} total):",
199 fg="bright_blue",
200 bold=True,
201 )
202 click.echo()
203
204 # Headers
205 headers = ["Span ID", "Trace ID", "Name", "Duration", "Kind", "Status"]
206 col_widths = [22, 22, 41, 12, 12, 16]
207
208 # Print headers
209 header_line = ""
210 for header, width in zip(headers, col_widths):
211 header_line += header.ljust(width)
212 click.secho(header_line, bold=True)
213 click.echo("-" * sum(col_widths))
214
215 # Print spans
216 for span in spans:
217 status_display = ""
218 if span.status:
219 if span.status in ["STATUS_CODE_OK", "OK"]:
220 status_display = "✓ OK"
221 elif span.status not in ["STATUS_CODE_UNSET", "UNSET"]:
222 status_display = f"✗ {span.status}"
223
224 row = [
225 span.span_id[:18] + "..." if len(span.span_id) > 20 else span.span_id,
226 span.trace.trace_id[:18] + "..."
227 if len(span.trace.trace_id) > 20
228 else span.trace.trace_id,
229 span.name[:37] + "..." if len(span.name) > 40 else span.name,
230 f"{span.duration_ms():.1f}ms",
231 span.kind[:10],
232 status_display[:15],
233 ]
234
235 # Build row with colored status
236 row_parts = []
237 for i, (value, width) in enumerate(zip(row, col_widths)):
238 if i == 5: # Status column
239 if span.status and span.status in ["STATUS_CODE_OK", "OK"]:
240 colored_value = click.style(str(value), fg="green")
241 # Need to account for the extra characters from coloring
242 padding = width - len(str(value))
243 row_parts.append(colored_value + " " * padding)
244 elif span.status and span.status not in [
245 "STATUS_CODE_UNSET",
246 "UNSET",
247 "",
248 ]:
249 colored_value = click.style(str(value), fg="red")
250 # Need to account for the extra characters from coloring
251 padding = width - len(str(value))
252 row_parts.append(colored_value + " " * padding)
253 else:
254 row_parts.append(str(value).ljust(width))
255 else:
256 row_parts.append(str(value).ljust(width))
257
258 click.echo("".join(row_parts))
259
260
261@observer_cli.command("span")
262@click.argument("span_id")
263@click.option("--json", "output_json", is_flag=True, help="Output as JSON")
264def span_detail(span_id: str, output_json: bool) -> None:
265 """Show detailed span information"""
266 try:
267 span = Span.query.select_related("trace").get(span_id=span_id)
268 except Span.DoesNotExist:
269 click.secho(f"Error: Span with ID '{span_id}' not found", fg="red", err=True)
270 raise click.Abort()
271
272 if output_json:
273 # Output as JSON
274 click.echo(json.dumps(span.span_data, indent=2))
275 else:
276 # Detailed output
277 label_width = 12
278 click.secho(
279 f"{'Span:':<{label_width}} {span.span_id}", fg="bright_blue", bold=True
280 )
281 click.echo(f"{'Trace:':<{label_width}} {span.trace.trace_id}")
282 click.echo(f"{'Name:':<{label_width}} {span.name}")
283 click.echo(f"{'Kind:':<{label_width}} {span.kind}")
284 click.echo(
285 f"{'Start:':<{label_width}} {span.start_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]} UTC"
286 )
287 click.echo(
288 f"{'End:':<{label_width}} {span.end_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]} UTC"
289 )
290 click.echo(f"{'Duration:':<{label_width}} {span.duration_ms():.2f}ms")
291
292 if span.parent_id:
293 click.echo(f"{'Parent ID:':<{label_width}} {span.parent_id}")
294
295 if span.status:
296 status_color = "green" if span.status in ["STATUS_CODE_OK", "OK"] else "red"
297 click.echo(f"{'Status:':<{label_width}} ", nl=False)
298 click.secho(span.status, fg=status_color)
299
300 # Show attributes
301 if span.attributes:
302 click.echo()
303 click.secho("Attributes:", fg="bright_blue", bold=True)
304 for key, value in span.attributes.items():
305 # Format value based on type
306 if isinstance(value, str) and len(value) > 100:
307 value = value[:97] + "..."
308 click.echo(f" {key}: {value}")
309
310 # Show SQL query if present
311 if span.sql_query:
312 click.echo()
313 click.secho("SQL Query:", fg="bright_blue", bold=True)
314 formatted_sql = span.get_formatted_sql()
315 if formatted_sql:
316 for line in formatted_sql.split("\n"):
317 click.echo(f" {line}")
318
319 # Show query parameters
320 if span.sql_query_params:
321 click.echo()
322 click.secho("Query Parameters:", fg="bright_blue", bold=True)
323 for param, value in span.sql_query_params.items():
324 click.echo(f" {param}: {value}")
325
326 # Show events
327 if span.events:
328 click.echo()
329 click.secho("Events:", fg="bright_blue", bold=True)
330 for event in span.events:
331 timestamp = span.format_event_timestamp(event.get("timestamp", ""))
332 click.echo(f" {event.get('name', 'unnamed')} at {timestamp}")
333 if event.get("attributes"):
334 for key, value in event["attributes"].items():
335 # Special handling for stack traces
336 if key == "exception.stacktrace" and isinstance(value, str):
337 click.echo(f" {key}:")
338 lines = value.split("\n")[:10] # Show first 10 lines
339 for line in lines:
340 click.echo(f" {line}")
341 if len(value.split("\n")) > 10:
342 click.echo(" ... (truncated)")
343 else:
344 click.echo(f" {key}: {value}")
345
346 # Show links
347 if span.links:
348 click.echo()
349 click.secho("Links:", fg="bright_blue", bold=True)
350 for link in span.links:
351 click.echo(
352 f" Trace: {link.get('context', {}).get('trace_id', 'unknown')}"
353 )
354 click.echo(
355 f" Span: {link.get('context', {}).get('span_id', 'unknown')}"
356 )
357 if link.get("attributes"):
358 for key, value in link["attributes"].items():
359 click.echo(f" {key}: {value}")
360
361
362def format_trace_output(trace: Trace) -> str:
363 """Format trace output for display - extracted for reuse."""
364 output_lines: list[str] = []
365
366 # Trace details with aligned labels
367 label_width = 12
368 start_time = trace.start_time
369 end_time = trace.end_time
370 output_lines.append(
371 click.style(
372 f"{'Trace:':<{label_width}} {trace.trace_id}", fg="bright_blue", bold=True
373 )
374 )
375 output_lines.append(
376 f"{'Start:':<{label_width}} {start_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]} UTC"
377 )
378 output_lines.append(
379 f"{'End:':<{label_width}} {end_time.strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]} UTC"
380 )
381 output_lines.append(f"{'Duration:':<{label_width}} {trace.duration_ms():.2f}ms")
382
383 if trace.summary:
384 output_lines.append(f"{'Summary:':<{label_width}} {trace.summary}")
385
386 if trace.request_id:
387 output_lines.append(f"{'Request ID:':<{label_width}} {trace.request_id}")
388 if trace.user_id:
389 output_lines.append(f"{'User ID:':<{label_width}} {trace.user_id}")
390 if trace.session_id:
391 output_lines.append(f"{'Session ID:':<{label_width}} {trace.session_id}")
392
393 output_lines.append("")
394 output_lines.append(click.style("Spans:", fg="bright_blue", bold=True))
395
396 # Get annotated spans with nesting levels
397 spans = trace.spans.query.all().annotate_spans() # type: ignore[attr-defined]
398
399 # Build parent-child relationships
400 span_dict = {span.span_id: span for span in spans}
401 children: dict[str, list[str]] = {}
402 for span in spans:
403 if span.parent_id:
404 children.setdefault(span.parent_id, []).append(span.span_id)
405
406 def format_span_tree(span: Span, level: int = 0) -> list[str]:
407 lines: list[str] = []
408 # Simple 4-space indentation
409 prefix = " " * level
410
411 # Span name with duration and status
412 duration = span.duration_ms()
413
414 # Determine status icon
415 status_icon = ""
416 if span.status:
417 if span.status in ["STATUS_CODE_OK", "OK"]:
418 status_icon = " ✓"
419 elif span.status not in ["STATUS_CODE_UNSET", "UNSET"]:
420 status_icon = " ✗"
421
422 # Color based on span kind, but red if error
423 if span.status and span.status not in [
424 "STATUS_CODE_OK",
425 "STATUS_CODE_UNSET",
426 "OK",
427 "UNSET",
428 ]:
429 color = "red"
430 else:
431 color_map = {
432 "SERVER": "green",
433 "CLIENT": "cyan",
434 "INTERNAL": "white",
435 "PRODUCER": "magenta",
436 "CONSUMER": "yellow",
437 }
438 color = color_map.get(span.kind, "white")
439
440 # Build span line
441 span_line = (
442 prefix
443 + click.style(span.name, fg=color, bold=True, underline=True)
444 + click.style(f" ({duration:.2f}ms){status_icon}", fg=color, bold=True)
445 + click.style(f" [{span.span_id}]", fg="bright_black")
446 )
447 lines.append(span_line)
448
449 # Show additional details with proper indentation
450 detail_prefix = " " * (level + 1)
451
452 # Show SQL queries
453 if span.sql_query:
454 lines.append(
455 f"{detail_prefix}SQL: {span.sql_query[:80]}{'...' if len(span.sql_query) > 80 else ''}"
456 )
457
458 # Show annotations (like duplicate queries)
459 for annotation in span.annotations:
460 severity_color = "yellow" if annotation["severity"] == "warning" else "red"
461 lines.append(
462 click.style(
463 f"{detail_prefix}⚠️ {annotation['message']}", fg=severity_color
464 )
465 )
466
467 # Show exceptions
468 if stacktrace := span.get_exception_stacktrace():
469 lines.append(click.style(f"{detail_prefix}❌ Exception occurred", fg="red"))
470 # Show first few lines of stacktrace
471 stack_lines = stacktrace.split("\n")[:3]
472 for line in stack_lines:
473 if line.strip():
474 lines.append(f"{detail_prefix} {line.strip()}")
475
476 # Format children recursively
477 if span.span_id in children:
478 child_ids = children[span.span_id]
479 for child_id in child_ids:
480 child_span = span_dict[child_id]
481 lines.extend(format_span_tree(child_span, level + 1))
482
483 return lines
484
485 # Start with root spans (spans without parents)
486 root_spans = [span for span in spans if not span.parent_id]
487 for root_span in root_spans:
488 output_lines.extend(format_span_tree(root_span, 0))
489
490 return "\n".join(output_lines)