1from __future__ import annotations
  2
  3from collections.abc import Sequence
  4
  5from plain import models
  6from plain.admin.views import (
  7    AdminModelDetailView,
  8    AdminModelListView,
  9    AdminViewset,
 10    register_viewset,
 11)
 12
 13from .models import Log, Span, Trace
 14
 15
 16@register_viewset
 17class TraceViewset(AdminViewset):
 18    class ListView(AdminModelListView):
 19        nav_section = "Observer"
 20        nav_icon = "activity"
 21        model = Trace
 22        fields = [
 23            "trace_id",
 24            "request_id",
 25            "session_id",
 26            "user_id",
 27            "start_time",
 28        ]
 29        allow_global_search = False
 30        actions = ["Delete"]
 31
 32        def perform_action(self, action: str, target_ids: Sequence[int]) -> None:
 33            if action == "Delete":
 34                Trace.query.filter(id__in=target_ids).delete()
 35
 36    class DetailView(AdminModelDetailView):
 37        model = Trace
 38
 39
 40@register_viewset
 41class SpanViewset(AdminViewset):
 42    class ListView(AdminModelListView):
 43        nav_section = "Observer"
 44        nav_icon = "activity"
 45        model = Span
 46        fields = [
 47            "name",
 48            "kind",
 49            "status",
 50            "span_id",
 51            "parent_id",
 52            "start_time",
 53        ]
 54        queryset_order = ["-id"]
 55        allow_global_search = False
 56        presets = ["Parents only"]
 57        search_fields = ["name", "span_id", "parent_id"]
 58        actions = ["Delete"]
 59
 60        def perform_action(self, action: str, target_ids: Sequence[int]) -> None:
 61            if action == "Delete":
 62                Span.query.filter(id__in=target_ids).delete()
 63
 64        def get_objects(self) -> models.QuerySet:
 65            return (
 66                super()
 67                .get_objects()
 68                .only(
 69                    "name",
 70                    "kind",
 71                    "span_id",
 72                    "parent_id",
 73                    "start_time",
 74                )
 75            )
 76
 77        def get_initial_queryset(self) -> models.QuerySet:
 78            queryset = super().get_initial_queryset()
 79            if self.preset == "Parents only":
 80                queryset = queryset.filter(parent_id="")
 81            return queryset
 82
 83    class DetailView(AdminModelDetailView):
 84        model = Span
 85
 86
 87@register_viewset
 88class LogViewset(AdminViewset):
 89    class ListView(AdminModelListView):
 90        nav_section = "Observer"
 91        nav_icon = "activity"
 92        model = Log
 93        fields = [
 94            "timestamp",
 95            "level",
 96            "message",
 97            "trace",
 98            "span",
 99        ]
100        queryset_order = ["-timestamp"]
101        allow_global_search = False
102        search_fields = ["message", "level"]
103        filters = ["level"]
104        actions = ["Delete selected", "Delete all"]
105
106        def perform_action(self, action: str, target_ids: Sequence[int]) -> None:
107            if action == "Delete selected":
108                Log.query.filter(id__in=target_ids).delete()
109            elif action == "Delete all":
110                Log.query.all().delete()
111
112        def get_objects(self) -> models.QuerySet:
113            return (
114                super()
115                .get_objects()
116                .select_related("trace", "span")
117                .only(
118                    "timestamp",
119                    "level",
120                    "message",
121                    "span__span_id",
122                    "trace__trace_id",
123                )
124            )
125
126    class DetailView(AdminModelDetailView):
127        model = Log