scaylr-logs.go

package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/alecthomas/chroma"
	"github.com/alecthomas/chroma/formatters"
	"github.com/alecthomas/chroma/lexers"
	"github.com/alecthomas/chroma/styles"
	"github.com/fatih/color"
)

var QUERY_ENDPOINT = "https://www.scalyr.com/api/query"
var DEFAULT_REQUEST_FREQUENCY_SECONDS = 5
var DEFAULT_COUNT = 25

type QueryResp struct {
	Matches []QueryMatch `json: "matches"`
}

type QueryMatch struct {
	Message   string `json: "message"`
	Timestamp string `json: "timestamp"`
}

func (q *QueryMatch) ToLogItem() *LogItem {
	i, err := strconv.ParseInt(q.Timestamp, 10, 64)
	if err != nil {
		log.Fatalf("failed to parse int")
	}

	tm := time.Unix(0, i)

	return &LogItem{
		Message:   q.Message,
		Timestamp: tm,
	}
}

type LogItem struct {
	Message          string
	FormattedMessage bytes.Buffer
	Timestamp        time.Time
}

func (l *LogItem) Write(b []byte) (int, error) {
	l.FormattedMessage.WriteString(string(b))
	return len(b), nil
}

type Log struct {
	Items []LogItem
}

func (l *Log) MaybeAddAndPrint(item *LogItem) {
	if len(l.Items) > 0 {
		lastTime := l.Items[len(l.Items)-1].Timestamp
		if item.Timestamp.After(lastTime) {
			l.AddAndPrint(item)
		}
	} else {
		l.AddAndPrint(item)
	}
}

func (l *Log) AddAndPrint(item *LogItem) {
	PrintWithFormatting(item)
	l.Items = append(l.Items, *item)
}

func trueColorNoBG(w io.Writer, style *chroma.Style, it chroma.Iterator) error {
	for token := it(); token != nil; token = it() {
		entry := style.Get(token.Type)
		if !entry.IsZero() {
			out := ""
			if entry.Bold == chroma.Yes {
				out += "\033[1m"
			}
			if entry.Underline == chroma.Yes {
				out += "\033[4m"
			}
			if entry.Colour.IsSet() {
				out += fmt.Sprintf("\033[38;2;%d;%d;%dm", entry.Colour.Red(), entry.Colour.Green(), entry.Colour.Blue())
			}
			fmt.Fprint(w, out)
		}
		fmt.Fprint(w, token.Value)
		if !entry.IsZero() {
			fmt.Fprint(w, "\033[0m")
		}
	}
	return nil
}

var TTY16mNoBG = formatters.Register("terminal16m", chroma.FormatterFunc(trueColorNoBG))

func PrintWithFormatting(item *LogItem) {
	blue := color.New(color.FgBlue).SprintFunc()
	timestamp := blue(item.Timestamp.Format("2006-01-02 15:04:05"))

	lexer := lexers.Get("json")
	if lexer == nil {
		lexer = lexers.Fallback
	}
	// lexer = chroma.Coalesce(lexer)

	style := styles.Get("monokai")
	if style == nil {
		style = styles.Fallback
	}
	formatter := TTY16mNoBG
	if formatter == nil {
		formatter = formatters.Fallback
	}

	iterator, err := lexer.Tokenise(nil, item.Message)
	if err != nil {
		log.Fatalf("could not tokenize")
	}
	err = formatter.Format(item, style, iterator)

	fmt.Println(timestamp, item.FormattedMessage.String())
}

func main() {
	log := Log{}

	query := flag.String("query", "", "Query")
	requestFrequency := flag.Int("freq", DEFAULT_REQUEST_FREQUENCY_SECONDS, "Query")
	token := flag.String("token", os.Getenv("SCALYR_API_KEY"), "API Key with Read access")
	count := flag.String("count", strconv.Itoa(DEFAULT_COUNT), "How many entries to fetch on each request (default 25)")

	flag.Parse()

	body := makeRequest(*query, *token, *count)
	for _ = range time.Tick(time.Duration(*requestFrequency) * time.Second) {
		// body := makeRequest(`$app == "settler" NOT "rdkafka" NOT "health-status"`)
		body = makeRequest(*query, *token, *count)

		queryResponse := QueryResp{}
		_ = json.Unmarshal(body, &queryResponse)

		for _, match := range queryResponse.Matches {
			log.MaybeAddAndPrint(match.ToLogItem())
		}
	}
}

func makeRequest(query string, token string, count string) []byte {
	req := map[string]string{
		"token":     token,
		"queryType": "log",
		"filter":    query,
		"maxCount":  count,
	}

	reqJSON, err := json.Marshal(req)

	if err != nil {
		log.Fatalf("Could not marshal request to JSON")
	}

	resp, err := http.Post(QUERY_ENDPOINT, "application/json", bytes.NewBuffer(reqJSON))

	if err != nil {
		log.Fatalf("Failed to make HTTP request")
	}

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		log.Fatalf("Failed to read body")
	}

	return body
}