Logging-System/cli/cli.go

// Package cli 实现 logctl：serve / tail / query。
package cli

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"strings"

	"git-ssh.u1680832.nyat.app/weichen/Logging-System.git/server"
)

// Main 执行命令行并返回退出码。
func Main(args []string) int { return Run(args, os.Stdout, os.Stderr) }

// Run 执行命令。
func Run(args []string, stdout, stderr io.Writer) int {
	if len(args) == 0 {
		fmt.Fprintln(stderr, usage)
		return 2
	}
	switch args[0] {
	case "serve":
		return cmdServe(args[1:], stderr)
	case "query":
		return cmdQuery(args[1:], stdout, stderr)
	case "tail":
		return cmdTail(args[1:], stdout, stderr)
	case "-h", "--help", "help":
		fmt.Fprintln(stdout, usage)
		return 0
	default:
		fmt.Fprintf(stderr, "未知命令 %q\n%s\n", args[0], usage)
		return 2
	}
}

const usage = `用法:
  logctl serve [--addr :8092] [--data 文件] [--token 密钥]
  logctl query --addr URL [--trace ID] [--level 级别] [--service 名] [--q 关键词] [--token 密钥]
  logctl tail --addr URL [--trace ID] [--token 密钥]`

func cmdServe(args []string, stderr io.Writer) int {
	fs := flag.NewFlagSet("serve", flag.ContinueOnError)
	fs.SetOutput(stderr)
	addr := fs.String("addr", server.DefaultAddr, "监听地址")
	data := fs.String("data", "", "JSONL 落盘路径，空则只放内存")
	token := fs.String("token", "", "可选写入令牌")
	if err := fs.Parse(args); err != nil {
		return 2
	}
	srv, err := server.New(*data, *token, 0)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	defer srv.Close()
	fmt.Fprintf(stderr, "logctl serve %s\n", *addr)
	if err := srv.ListenAndServe(*addr); err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	return 0
}

func cmdQuery(args []string, stdout, stderr io.Writer) int {
	fs := flag.NewFlagSet("query", flag.ContinueOnError)
	fs.SetOutput(stderr)
	addr := fs.String("addr", "http://127.0.0.1:8092", "采集器地址")
	traceID := fs.String("trace", "", "追踪标识")
	level := fs.String("level", "", "级别（精确匹配）")
	service := fs.String("service", "", "服务名")
	q := fs.String("q", "", "关键词")
	token := fs.String("token", "", "令牌")
	if err := fs.Parse(args); err != nil {
		return 2
	}
	u, err := url.Parse(strings.TrimRight(*addr, "/") + "/v1/logs")
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	params := u.Query()
	set := func(k, v string) {
		if v != "" {
			params.Set(k, v)
		}
	}
	set("trace_id", *traceID)
	set("level", *level)
	set("service", *service)
	set("q", *q)
	u.RawQuery = params.Encode()
	req, err := http.NewRequest(http.MethodGet, u.String(), nil)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	if *token != "" {
		req.Header.Set("Authorization", "Bearer "+*token)
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	defer res.Body.Close()
	body, _ := io.ReadAll(res.Body)
	if res.StatusCode != http.StatusOK {
		fmt.Fprintf(stderr, "query %s: %s\n", res.Status, body)
		return 1
	}
	var pretty bytes.Buffer
	if err := json.Indent(&pretty, body, "", "  "); err != nil {
		stdout.Write(body)
		return 0
	}
	fmt.Fprintln(stdout, pretty.String())
	return 0
}

func cmdTail(args []string, stdout, stderr io.Writer) int {
	fs := flag.NewFlagSet("tail", flag.ContinueOnError)
	fs.SetOutput(stderr)
	addr := fs.String("addr", "http://127.0.0.1:8092", "采集器地址")
	traceID := fs.String("trace", "", "只打印该追踪标识")
	token := fs.String("token", "", "令牌")
	if err := fs.Parse(args); err != nil {
		return 2
	}
	u := strings.TrimRight(*addr, "/") + "/v1/logs/tail"
	req, err := http.NewRequest(http.MethodGet, u, nil)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	if *token != "" {
		req.Header.Set("Authorization", "Bearer "+*token)
	}
	res, err := http.DefaultClient.Do(req)
	if err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	defer res.Body.Close()
	if res.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(res.Body)
		fmt.Fprintf(stderr, "tail %s: %s\n", res.Status, body)
		return 1
	}
	sc := bufio.NewScanner(res.Body)
	for sc.Scan() {
		line := sc.Text()
		if !strings.HasPrefix(line, "data: ") {
			continue
		}
		payload := strings.TrimPrefix(line, "data: ")
		if *traceID != "" && !strings.Contains(payload, *traceID) {
			continue
		}
		fmt.Fprintln(stdout, payload)
	}
	if err := sc.Err(); err != nil {
		fmt.Fprintln(stderr, err)
		return 1
	}
	return 0
}