unibots/shell/logger/query.go

package logger

import (
	"bufio"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"
)

// ReadLogs returns all log entries for agentID between from and to (inclusive).
func ReadLogs(baseDir, agentID string, from, to time.Time) ([]json.RawMessage, error) {
	var result []json.RawMessage
	for d := from; !d.After(to); d = d.AddDate(0, 0, 1) {
		entries, err := ReadDayLogs(baseDir, agentID, d)
		if err != nil {
			continue // skip missing days
		}
		result = append(result, entries...)
	}
	return result, nil
}

// ReadDayLogs returns all log entries for a specific day.
func ReadDayLogs(baseDir, agentID string, date time.Time) ([]json.RawMessage, error) {
	dir := filepath.Join(baseDir, agentID)
	prefix := date.Format("2006-01-02")

	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, fmt.Errorf("read dir %s: %w", dir, err)
	}

	var result []json.RawMessage
	for _, e := range entries {
		name := e.Name()
		if !strings.HasPrefix(name, prefix) || !isLogFile(name) {
			continue
		}
		lines, err := readLogFile(filepath.Join(dir, name))
		if err != nil {
			continue
		}
		result = append(result, lines...)
	}
	return result, nil
}

// SearchLogs returns log entries where field equals value, within the date range.
func SearchLogs(baseDir, agentID string, field, value string, from, to time.Time) ([]json.RawMessage, error) {
	all, err := ReadLogs(baseDir, agentID, from, to)
	if err != nil {
		return nil, err
	}

	var matched []json.RawMessage
	for _, raw := range all {
		var m map[string]any
		if json.Unmarshal(raw, &m) != nil {
			continue
		}
		if v, ok := m[field]; ok && fmt.Sprint(v) == value {
			matched = append(matched, raw)
		}
	}
	return matched, nil
}

// ListAgents returns the agent IDs that have log directories under baseDir.
func ListAgents(baseDir string) ([]string, error) {
	entries, err := os.ReadDir(baseDir)
	if err != nil {
		return nil, fmt.Errorf("read base dir %s: %w", baseDir, err)
	}
	var ids []string
	for _, e := range entries {
		if e.IsDir() {
			ids = append(ids, e.Name())
		}
	}
	sort.Strings(ids)
	return ids, nil
}

// ListDates returns the dates for which logs exist for the given agent.
func ListDates(baseDir, agentID string) ([]time.Time, error) {
	dir := filepath.Join(baseDir, agentID)
	entries, err := os.ReadDir(dir)
	if err != nil {
		return nil, fmt.Errorf("read dir %s: %w", dir, err)
	}

	seen := make(map[string]bool)
	var dates []time.Time
	for _, e := range entries {
		if e.IsDir() || !isLogFile(e.Name()) {
			continue
		}
		d := parseDateFromFilename(e.Name())
		if d.IsZero() {
			continue
		}
		key := d.Format("2006-01-02")
		if !seen[key] {
			seen[key] = true
			dates = append(dates, d)
		}
	}
	sort.Slice(dates, func(i, j int) bool { return dates[i].Before(dates[j]) })
	return dates, nil
}

// readLogFile reads all JSONL lines from a file (.jsonl or .jsonl.gz).
func readLogFile(path string) ([]json.RawMessage, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer f.Close()

	var r io.Reader = f
	if strings.HasSuffix(path, ".gz") {
		gz, err := gzip.NewReader(f)
		if err != nil {
			return nil, err
		}
		defer gz.Close()
		r = gz
	}

	var lines []json.RawMessage
	scanner := bufio.NewScanner(r)
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)
	for scanner.Scan() {
		line := scanner.Bytes()
		if len(line) == 0 {
			continue
		}
		cp := make([]byte, len(line))
		copy(cp, line)
		lines = append(lines, json.RawMessage(cp))
	}
	return lines, scanner.Err()
}