1from functools import cached_property
  2
  3from plain.admin.toolbar import ToolbarPanel, register_toolbar_panel
  4from plain.admin.views import (
  5    AdminModelDetailView,
  6    AdminModelListView,
  7    AdminViewset,
  8    register_viewset,
  9)
 10
 11from .core import Observer
 12from .models import Span, Trace
 13
 14
 15@register_viewset
 16class TraceViewset(AdminViewset):
 17    class ListView(AdminModelListView):
 18        nav_section = "Observer"
 19        model = Trace
 20        fields = [
 21            "trace_id",
 22            "request_id",
 23            "session_id",
 24            "user_id",
 25            "start_time",
 26        ]
 27        allow_global_search = False
 28        # Actually want a button to delete ALL! not possible yet
 29        # actions = ["Delete"]
 30
 31        # def perform_action(self, action: str, target_pks: list):
 32        #     if action == "Delete":
 33        #         Trace.objects.filter(id__in=target_pks).delete()
 34
 35    class DetailView(AdminModelDetailView):
 36        model = Trace
 37        template_name = "admin/observer/trace_detail.html"
 38
 39        def get_template_context(self):
 40            context = super().get_template_context()
 41            trace_id = self.url_kwargs["pk"]
 42            context["trace"] = Trace.objects.get(pk=trace_id)
 43            context["show_delete_button"] = False
 44            return context
 45
 46
 47@register_viewset
 48class SpanViewset(AdminViewset):
 49    class ListView(AdminModelListView):
 50        nav_section = "Observer"
 51        model = Span
 52        fields = [
 53            "name",
 54            "kind",
 55            "status",
 56            "span_id",
 57            "parent_id",
 58            "start_time",
 59        ]
 60        queryset_order = ["-pk"]
 61        allow_global_search = False
 62        displays = ["Parents only"]
 63        search_fields = ["name", "span_id", "parent_id"]
 64
 65        def get_objects(self):
 66            return (
 67                super()
 68                .get_objects()
 69                .only(
 70                    "name",
 71                    "kind",
 72                    "span_id",
 73                    "parent_id",
 74                    "start_time",
 75                )
 76            )
 77
 78        def get_initial_queryset(self):
 79            queryset = super().get_initial_queryset()
 80            if self.display == "Parents only":
 81                queryset = queryset.filter(parent_id="")
 82            return queryset
 83
 84    class DetailView(AdminModelDetailView):
 85        model = Span
 86
 87
 88@register_toolbar_panel
 89class ObserverToolbarPanel(ToolbarPanel):
 90    name = "Observer"
 91    template_name = "toolbar/observer.html"
 92    button_template_name = "toolbar/observer_button.html"
 93
 94    @cached_property
 95    def observer(self):
 96        """Get the Observer instance for this request."""
 97        return Observer(self.request)
 98
 99    def get_template_context(self):
100        context = super().get_template_context()
101        context["observer"] = self.observer
102        return context