clickhouse-tools/internal/chlog/format.go

package chlog

import (
	"encoding/json"
	"fmt"
	"io"
	"strings"
	"time"
)

const (
	colReset = "\x1b[0m"
	colDim   = "\x1b[2m"
	colCyan  = "\x1b[36m"
	colGreen = "\x1b[32m"
	colRed   = "\x1b[31m"
	colYel   = "\x1b[33m"
)

type Formatter func(w io.Writer, r Row) error

func NewFormatter(format string, color bool) (Formatter, error) {
	switch format {
	case "text":
		return textFormatter(color), nil
	case "json":
		return jsonFormatter, nil
	case "logfmt":
		return logfmtFormatter, nil
	default:
		return nil, fmt.Errorf("unknown format %q: want text, json or logfmt", format)
	}
}

// origin renders the row's source identity: ns/pod for k8s rows, host for vm.
func origin(r Row) string {
	if r.Pod != "" {
		return r.Namespace + "/" + r.Pod
	}
	if r.Namespace != "" {
		return r.Namespace
	}
	return r.Host
}

func severityColor(sev string) string {
	switch strings.ToLower(sev) {
	case "error", "err", "fatal", "critical", "crit":
		return colRed
	case "warn", "warning":
		return colYel
	default:
		return ""
	}
}

func textFormatter(color bool) Formatter {
	return func(w io.Writer, r Row) error {
		ts := r.Time().Format("2006-01-02T15:04:05.000Z")
		msg := r.Message
		if !color {
			_, err := fmt.Fprintf(w, "%s %s %s\n", ts, origin(r), msg)
			return err
		}
		if c := severityColor(r.Severity); c != "" {
			msg = c + msg + colReset
		}
		_, err := fmt.Fprintf(w, "%s%s%s %s%s%s %s\n",
			colDim, ts, colReset, colCyan, origin(r), colReset, msg)
		return err
	}
}

type jsonRow struct {
	Timestamp string            `json:"timestamp"`
	Host      string            `json:"host"`
	Source    string            `json:"source"`
	Namespace string            `json:"namespace,omitempty"`
	Pod       string            `json:"pod,omitempty"`
	Container string            `json:"container,omitempty"`
	Stream    string            `json:"stream,omitempty"`
	Severity  string            `json:"severity,omitempty"`
	Message   string            `json:"message"`
	Labels    map[string]string `json:"labels,omitempty"`
	Fields    map[string]string `json:"fields,omitempty"`
}

func jsonFormatter(w io.Writer, r Row) error {
	b, err := json.Marshal(jsonRow{
		Timestamp: r.Time().Format(time.RFC3339Nano),
		Host:      r.Host,
		Source:    r.Source,
		Namespace: r.Namespace,
		Pod:       r.Pod,
		Container: r.Container,
		Stream:    r.Stream,
		Severity:  r.Severity,
		Message:   r.Message,
		Labels:    r.Labels,
		Fields:    r.Fields,
	})
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(w, "%s\n", b)
	return err
}

func logfmtValue(s string) string {
	if s == "" {
		return `""`
	}
	if strings.ContainsAny(s, " \t\"=\n") {
		return fmt.Sprintf("%q", s)
	}
	return s
}

func logfmtFormatter(w io.Writer, r Row) error {
	var b strings.Builder
	pair := func(k, v string) {
		if v == "" {
			return
		}
		if b.Len() > 0 {
			b.WriteByte(' ')
		}
		b.WriteString(k)
		b.WriteByte('=')
		b.WriteString(logfmtValue(v))
	}
	pair("ts", r.Time().Format(time.RFC3339Nano))
	pair("source", r.Source)
	pair("host", r.Host)
	pair("ns", r.Namespace)
	pair("pod", r.Pod)
	pair("container", r.Container)
	pair("stream", r.Stream)
	pair("severity", r.Severity)
	pair("app", r.Labels["app"])
	if b.Len() > 0 {
		b.WriteByte(' ')
	}
	b.WriteString("msg=")
	b.WriteString(logfmtValue(r.Message))
	b.WriteByte('\n')
	_, err := io.WriteString(w, b.String())
	return err
}