clickhouse-tools/main.go

package main

import (
	"context"
	"errors"
	"fmt"
	"os"
	"os/signal"
	"path/filepath"
	"strings"
	"syscall"
	"time"

	"github.com/spf13/cobra"

	"git.unkin.net/unkin/clickhouse-tools/internal/chlog"
)

var version = "dev"

const guardWindow = 6 * time.Hour

type commonFlags struct {
	since     string
	until     string
	namespace string
	host      string
	pod       string
	container string
	app       string
	severity  string
	stream    string
	source    string
	limit     uint64
	format    string
}

func (cf *commonFlags) register(cmd *cobra.Command, defaultLimit uint64) {
	fl := cmd.Flags()
	fl.StringVar(&cf.since, "since", "1h", "start of time range: duration ago (15m, 1h, 2d) or RFC3339")
	fl.StringVar(&cf.until, "until", "", "end of time range: duration ago or RFC3339 (default now)")
	fl.StringVarP(&cf.namespace, "namespace", "n", "", "filter: k8s namespace")
	fl.StringVar(&cf.host, "host", "", "filter: host")
	fl.StringVar(&cf.pod, "pod", "", "filter: pod name")
	fl.StringVar(&cf.container, "container", "", "filter: container name")
	fl.StringVar(&cf.app, "app", "", "filter: labels['app']")
	fl.StringVar(&cf.severity, "severity", "", "filter: severity (case-insensitive)")
	fl.StringVar(&cf.stream, "stream", "", "filter: stream (stdout/stderr)")
	fl.StringVar(&cf.source, "source", "", "filter: source (k8s/vm)")
	fl.Uint64Var(&cf.limit, "limit", defaultLimit, "maximum rows to print (0 = unlimited)")
	fl.StringVar(&cf.format, "format", "text", "output format: text, json or logfmt")
}

func (cf *commonFlags) filter(now time.Time) (chlog.Filter, error) {
	since, err := chlog.ParseTimeSpec(cf.since, now)
	if err != nil {
		return chlog.Filter{}, fmt.Errorf("--since: %w", err)
	}
	until := now
	if cf.until != "" {
		until, err = chlog.ParseTimeSpec(cf.until, now)
		if err != nil {
			return chlog.Filter{}, fmt.Errorf("--until: %w", err)
		}
	}
	return chlog.Filter{
		Since:     since.UTC(),
		Until:     until.UTC(),
		Namespace: cf.namespace,
		Host:      cf.host,
		Pod:       cf.pod,
		Container: cf.container,
		App:       cf.app,
		Severity:  cf.severity,
		Stream:    cf.stream,
		Source:    cf.source,
		Limit:     cf.limit,
	}, nil
}

func stdoutIsTTY() bool {
	fi, err := os.Stdout.Stat()
	return err == nil && fi.Mode()&os.ModeCharDevice != 0
}

func (cf *commonFlags) formatter() (chlog.Formatter, error) {
	color := cf.format == "text" && stdoutIsTTY() && os.Getenv("NO_COLOR") == ""
	return chlog.NewFormatter(cf.format, color)
}

func newCatCmd(use string) *cobra.Command {
	cf := &commonFlags{}
	cmd := &cobra.Command{
		Use:   use,
		Short: "Print logs from the ClickHouse log store, oldest first",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			f, err := cf.filter(time.Now().UTC())
			if err != nil {
				return err
			}
			out, err := cf.formatter()
			if err != nil {
				return err
			}
			c := chlog.NewClient(chlog.ConfigFromEnv())
			_, err = chlog.Page(cmd.Context(), c, f, 0, func(r chlog.Row) error {
				return out(os.Stdout, r)
			})
			return err
		},
	}
	cf.register(cmd, 10000)
	return cmd
}

func newTailCmd(use string) *cobra.Command {
	cf := &commonFlags{}
	cmd := &cobra.Command{
		Use:   use,
		Short: "Follow logs from the ClickHouse log store",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, args []string) error {
			cf.limit = 0
			f, err := cf.filter(time.Now().UTC())
			if err != nil {
				return err
			}
			out, err := cf.formatter()
			if err != nil {
				return err
			}
			c := chlog.NewClient(chlog.ConfigFromEnv())
			err = chlog.Tail(cmd.Context(), c, f, nil, chlog.TailInterval, func(r chlog.Row) error {
				return out(os.Stdout, r)
			})
			if errors.Is(err, context.Canceled) {
				return nil
			}
			return err
		},
	}
	cf.register(cmd, 0)
	cmd.Flags().MarkHidden("until")
	cmd.Flags().MarkHidden("limit")
	return cmd
}

func newGrepCmd(use string) *cobra.Command {
	cf := &commonFlags{}
	var (
		regex      bool
		ignoreCase bool
		fields     []string
		force      bool
	)
	cmd := &cobra.Command{
		Use:   use + " <pattern>",
		Short: "Search log messages in the ClickHouse log store",
		Args:  cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			f, err := cf.filter(time.Now().UTC())
			if err != nil {
				return err
			}
			f.Pattern = args[0]
			f.Regex = regex
			f.IgnoreCase = ignoreCase
			f.Fields, err = parseFields(fields)
			if err != nil {
				return err
			}
			if !f.Selective() && f.Until.Sub(f.Since) > guardWindow && !force {
				return fmt.Errorf("unfiltered search over %s scans the whole table (no message index, ~281M rows/day); add --namespace/--host/--app, shrink --since to 6h or less, or pass --force",
					f.Until.Sub(f.Since).Round(time.Minute))
			}
			out, err := cf.formatter()
			if err != nil {
				return err
			}
			c := chlog.NewClient(chlog.ConfigFromEnv())
			_, err = chlog.Page(cmd.Context(), c, f, 0, func(r chlog.Row) error {
				return out(os.Stdout, r)
			})
			return err
		},
	}
	cf.register(cmd, 10000)
	fl := cmd.Flags()
	fl.BoolVar(&regex, "regex", false, "treat pattern as an RE2 regular expression")
	fl.BoolVarP(&ignoreCase, "ignore-case", "i", false, "case-insensitive match")
	fl.StringArrayVar(&fields, "fields", nil, "filter on structured fields: key=value (repeatable)")
	fl.BoolVar(&force, "force", false, "allow an unfiltered search wider than 6h")
	return cmd
}

func parseFields(kvs []string) (map[string]string, error) {
	if len(kvs) == 0 {
		return nil, nil
	}
	m := make(map[string]string, len(kvs))
	for _, kv := range kvs {
		k, v, ok := strings.Cut(kv, "=")
		if !ok || k == "" {
			return nil, fmt.Errorf("--fields %q: want key=value", kv)
		}
		m[k] = v
	}
	return m, nil
}

func newRootCmd() *cobra.Command {
	root := &cobra.Command{
		Use:           "chlog",
		Short:         "CLI for the ClickHouse log store (logs.raw)",
		Version:       version,
		SilenceUsage:  true,
		SilenceErrors: true,
	}
	root.AddCommand(newCatCmd("cat"), newTailCmd("tail"), newGrepCmd("grep"))
	return root
}

func entrypoint() *cobra.Command {
	var cmd *cobra.Command
	switch filepath.Base(os.Args[0]) {
	case "chcat":
		cmd = newCatCmd("chcat")
	case "chtail":
		cmd = newTailCmd("chtail")
	case "chgrep":
		cmd = newGrepCmd("chgrep")
	default:
		return newRootCmd()
	}
	cmd.Version = version
	cmd.SilenceUsage = true
	cmd.SilenceErrors = true
	return cmd
}

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()
	if err := entrypoint().ExecuteContext(ctx); err != nil {
		fmt.Fprintln(os.Stderr, "error:", err)
		os.Exit(1)
	}
}