agentmon/cmd/query-api/main.go

package main

import (
	"database/sql"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"agentmon/internal/httpx"
	"agentmon/internal/store/postgres"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
)

func main() {
	addr := envDefault("AGENTMON_QUERY_ADDR", ":8081")
	dsn := os.Getenv("DATABASE_URL")
	if dsn == "" {
		log.Fatalf("DATABASE_URL is required")
	}

	db, err := postgres.Open(dsn)
	if err != nil {
		log.Fatalf("failed to open DB: %v", err)
	}
	defer func() { _ = db.Close() }()

	r := chi.NewRouter()
	r.Use(middleware.RequestID)
	r.Use(middleware.RealIP)
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)

	r.Get("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})

	r.Get("/v1/events", func(w http.ResponseWriter, r *http.Request) {
		limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
		events, err := db.ListRecentEvents(r.Context(), limit)
		if err != nil {
			httpx.WriteJSON(w, http.StatusInternalServerError, map[string]any{"error": "db_error"})
			return
		}
		httpx.WriteJSON(w, http.StatusOK, map[string]any{"events": events})
	})

	r.Get("/v1/sessions", func(w http.ResponseWriter, r *http.Request) {
		q := r.URL.Query()
		f := postgres.SessionsFilter{
			Framework: q.Get("framework"),
			Host:      q.Get("host"),
		}

		if v := q.Get("limit"); v != "" {
			f.Limit, _ = strconv.Atoi(v)
		}

		if v := q.Get("from"); v != "" {
			if t, err := time.Parse(time.RFC3339, v); err == nil {
				f.From = &t
			} else if t, err := time.Parse("2006-01-02", v); err == nil {
				f.From = &t
			}
		}

		if v := q.Get("to"); v != "" {
			if t, err := time.Parse(time.RFC3339, v); err == nil {
				f.To = &t
			} else if t, err := time.Parse("2006-01-02", v); err == nil {
				end := t.Add(24*time.Hour - time.Nanosecond)
				f.To = &end
			}
		}

		if v := q.Get("cursor"); v != "" {
			if t, err := time.Parse(time.RFC3339Nano, v); err == nil {
				f.Cursor = &t
			}
		}

		sessions, nextCursor, err := db.ListSessions(r.Context(), f)
		if err != nil {
			httpx.WriteJSON(w, http.StatusInternalServerError, map[string]any{"error": "db_error"})
			return
		}

		resp := map[string]any{"sessions": sessions}
		if nextCursor != nil {
			resp["next_cursor"] = nextCursor.Format(time.RFC3339Nano)
		}
		httpx.WriteJSON(w, http.StatusOK, resp)
	})

	r.Get("/v1/sessions/{sessionID}", func(w http.ResponseWriter, r *http.Request) {
		sessionID := chi.URLParam(r, "sessionID")
		session, runs, err := db.GetSessionWithRuns(r.Context(), sessionID)
		if err == sql.ErrNoRows {
			httpx.WriteJSON(w, http.StatusNotFound, map[string]any{"error": "not_found"})
			return
		}
		if err != nil {
			httpx.WriteJSON(w, http.StatusInternalServerError, map[string]any{"error": "db_error"})
			return
		}
		httpx.WriteJSON(w, http.StatusOK, map[string]any{"session": session, "runs": runs})
	})

	r.Get("/v1/runs/{runID}", func(w http.ResponseWriter, r *http.Request) {
		runID := chi.URLParam(r, "runID")
		run, spans, err := db.GetRunWithSpans(r.Context(), runID)
		if err == sql.ErrNoRows {
			httpx.WriteJSON(w, http.StatusNotFound, map[string]any{"error": "not_found"})
			return
		}
		if err != nil {
			httpx.WriteJSON(w, http.StatusInternalServerError, map[string]any{"error": "db_error"})
			return
		}
		httpx.WriteJSON(w, http.StatusOK, map[string]any{"run": run, "spans": spans})
	})

	log.Printf("query-api listening on %s", addr)
	log.Fatal(http.ListenAndServe(addr, r))
}

func envDefault(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}