github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/integrations/loki/loki.go

/*
Copyright 2023.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package loki

import (
	"encoding/json"
	"fmt"
	"regexp"
	"strconv"
	"strings"

	"github.com/golang/snappy"
	"github.com/siglens/siglens/pkg/ast"
	"github.com/siglens/siglens/pkg/ast/pipesearch"
	dtu "github.com/siglens/siglens/pkg/common/dtypeutils"
	"github.com/siglens/siglens/pkg/es/writer"
	lokilog "github.com/siglens/siglens/pkg/integrations/loki/log"
	rutils "github.com/siglens/siglens/pkg/readerUtils"
	"github.com/siglens/siglens/pkg/segment"
	"github.com/siglens/siglens/pkg/segment/query/metadata"
	"github.com/siglens/siglens/pkg/segment/reader/record"
	"github.com/siglens/siglens/pkg/segment/search"
	"github.com/siglens/siglens/pkg/segment/structs"
	segwriter "github.com/siglens/siglens/pkg/segment/writer"
	"github.com/siglens/siglens/pkg/utils"
	vtable "github.com/siglens/siglens/pkg/virtualtable"
	log "github.com/sirupsen/logrus"
	"github.com/valyala/fasthttp"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/proto"
)

const (
	ContentJson        = "application/json; charset=utf-8"
	LOKIINDEX          = "loki-index"
	TimeStamp          = "timestamp"
	Index              = "_index"
	DefaultLimit       = 100
	MsToNanoConversion = 1_000_000
)

func parseLabels(labelsString string) map[string]string {
	labelsString = strings.Trim(labelsString, "{}")
	labelPairs := strings.Split(labelsString, ", ")

	labels := make(map[string]string)

	for _, pair := range labelPairs {
		parts := strings.Split(pair, "=")
		if len(parts) == 2 {
			key := strings.Trim(parts[0], "\"")
			value := strings.Trim(parts[1], "\"")
			labels[key] = value
		}
	}

	return labels
}

// Format of loki logs generated by promtail:
//
//	{
//		"streams": [
//		  {
//			"labels": "{filename=\"test.log\",job=\"test\"}",
//			"entries": [
//			  {
//				"timestamp": "2021-03-31T18:00:00.000Z",
//				"line": "test log line"
//			  }
//			]
//		  }
//		]
//	}
func ProcessLokiLogsIngestRequest(ctx *fasthttp.RequestCtx, myid uint64) {

	responsebody := make(map[string]interface{})
	buf, err := snappy.Decode(nil, ctx.PostBody())
	if err != nil {
		log.Errorf("ProcessLokiLogsIngestRequest: error decompressing request body, err: %v", err)
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = "Error decompressing request body"
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	logLine := lokilog.PushRequest{}

	err = proto.Unmarshal(buf, &logLine)
	if err != nil {
		log.Errorf("ProcessLokiLogsIngestRequest: Unable to unmarshal request body, err=%v", err)
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = "Unable to unmarshal request body"
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	logJson := protojson.Format(&logLine)

	tsNow := utils.GetCurrentTimeInMs()
	indexNameIn := LOKIINDEX
	if !vtable.IsVirtualTablePresent(&indexNameIn, myid) {
		log.Errorf("ProcessLokiLogsIngestRequest: Index name %v does not exist. Adding virtual table name and mapping.", indexNameIn)
		body := logJson
		err = vtable.AddVirtualTable(&indexNameIn, myid)
		if err != nil {
			ctx.SetStatusCode(fasthttp.StatusServiceUnavailable)
			responsebody["error"] = "Failed to add virtual table for index"
			utils.WriteJsonResponse(ctx, responsebody)
			return
		}
		err = vtable.AddMappingFromADoc(&indexNameIn, &body, myid)
		if err != nil {
			ctx.SetStatusCode(fasthttp.StatusServiceUnavailable)
			responsebody["error"] = "Failed to add mapping from a doc for index"
			utils.WriteJsonResponse(ctx, responsebody)
			return
		}
	}

	localIndexMap := make(map[string]string)

	var jsonData map[string][]map[string]interface{}
	err = json.Unmarshal([]byte(logJson), &jsonData)
	if err != nil {
		log.Errorf("ProcessLokiLogsIngestRequest: Unable to unmarshal request body, err=%v", err)
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = "Unable to unmarshal request body"
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	streams := jsonData["streams"]
	allIngestData := make(map[string]interface{})

	for _, stream := range streams {
		labels := stream["labels"].(string)
		ingestCommonFields := parseLabels(labels)

		// Note: We might not need separate filename and job fields in the future
		allIngestData["filename"] = ingestCommonFields["filename"]
		allIngestData["job"] = ingestCommonFields["job"]
		allIngestData["labels"] = labels

		entries, ok := stream["entries"].([]interface{})
		if !ok {
			log.Errorf("ProcessLokiLogsIngestRequest: Unable to convert entries to []interface{}")
			ctx.SetStatusCode(fasthttp.StatusBadRequest)
			responsebody["error"] = "Unable to convert entries to []interface{}"
			utils.WriteJsonResponse(ctx, responsebody)
			return
		}

		if len(entries) > 0 {
			for _, entry := range entries {
				entryMap, ok := entry.(map[string]interface{})
				if !ok {
					log.Errorf("ProcessLokiLogsIngestRequest: Unable to convert entry to map[string]interface{}")
					ctx.SetStatusCode(fasthttp.StatusBadRequest)
					responsebody["error"] = "Unable to convert entry to map[string]interface{}"
					utils.WriteJsonResponse(ctx, responsebody)
					return
				}
				timestamp, ok := entryMap["timestamp"].(string)
				if !ok {
					log.Errorf("ProcessLokiLogsIngestRequest: Unable to convert timestamp to string")
					ctx.SetStatusCode(fasthttp.StatusBadRequest)
					responsebody["error"] = "Unable to convert timestamp to string"
					utils.WriteJsonResponse(ctx, responsebody)
					return
				}
				line, ok := entryMap["line"].(string)
				if !ok {
					log.Errorf("ProcessLokiLogsIngestRequest: Unable to convert line to string")
					ctx.SetStatusCode(fasthttp.StatusBadRequest)
					responsebody["error"] = "Unable to convert line to string"
					utils.WriteJsonResponse(ctx, responsebody)
					return
				}

				allIngestData["timestamp"] = timestamp
				allIngestData["line"] = line

				test, err := json.Marshal(allIngestData)
				if err != nil {
					log.Errorf("ProcessLokiLogsIngestRequest: Unable to marshal data, err=%v", err)
					ctx.SetStatusCode(fasthttp.StatusBadRequest)
					responsebody["error"] = "Unable to marshal request body"
					utils.WriteJsonResponse(ctx, responsebody)
					return
				}

				err = writer.ProcessIndexRequest([]byte(test), tsNow, indexNameIn, uint64(len(test)), false, localIndexMap, myid)
				if err != nil {
					ctx.SetStatusCode(fasthttp.StatusServiceUnavailable)
					responsebody["error"] = "Failed to add entry to in mem buffer"
					utils.WriteJsonResponse(ctx, responsebody)
					return
				}
			}
		}
	}

	responsebody["status"] = "Success"
	utils.WriteJsonResponse(ctx, responsebody)
	ctx.SetStatusCode(fasthttp.StatusOK)
}

func ProcessLokiLabelRequest(ctx *fasthttp.RequestCtx) {
	indexName := []string{LOKIINDEX}
	responsebody := make(map[string]interface{})
	colNames := remove(metadata.GetAllColNames(indexName), "line")
	responsebody["data"] = colNames
	responsebody["status"] = "Success"
	utils.WriteJsonResponse(ctx, responsebody)
	ctx.SetStatusCode(fasthttp.StatusOK)
}

func ProcessLokiLabelValuesRequest(ctx *fasthttp.RequestCtx, myid uint64) {
	responsebody := make(map[string]interface{})

	labelName := utils.ExtractParamAsString(ctx.UserValue("labelName"))
	indexName := LOKIINDEX
	qid := rutils.GetNextQid()
	colVals, err := ast.GetColValues(labelName, indexName, qid, myid)

	if err != nil {
		ctx.SetStatusCode(fasthttp.StatusUnauthorized)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	responsebody["data"] = colVals
	responsebody["status"] = "Success"
	utils.WriteJsonResponse(ctx, responsebody)
	ctx.SetStatusCode(fasthttp.StatusOK)
}

func ProcessQueryRequest(ctx *fasthttp.RequestCtx, myid uint64) {
	query := removeUnimplementedMethods(string(ctx.QueryArgs().Peek("query")))
	if query == "" {
		log.Errorf(" ProcessQueryRequest: received empty search request body ")
		responsebody := make(map[string]interface{})
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = "received empty search request body"
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	limit, err := strconv.ParseUint(string(ctx.QueryArgs().Peek("limit")), 10, 64)
	if err != nil {
		responsebody := make(map[string]interface{})
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	qid := rutils.GetNextQid()

	ti := structs.InitTableInfo(LOKIINDEX, myid, false)
	simpleNode, aggs, err := pipesearch.ParseRequest(query, 0, 0, qid, "Log QL", LOKIINDEX)

	if aggs != nil && aggs.GroupByRequest != nil {
		aggs.GroupByRequest.GroupByColumns = remove(aggs.GroupByRequest.GroupByColumns, "line")
	}

	if err != nil {
		responsebody := make(map[string]interface{})
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	segment.LogASTNode("logql query parser", simpleNode, qid)
	segment.LogQueryAggsNode("logql aggs parser", aggs, qid)
	startTime := utils.GetCurrentTimeInMs()
	qc := structs.InitQueryContextWithTableInfo(ti, limit, 0, myid, false)
	queryResult := segment.ExecuteQuery(simpleNode, aggs, qid, qc)

	allJsons, allCols, err := record.GetJsonFromAllRrc(queryResult.AllRecords, false, qid, queryResult.SegEncToKey, aggs)

	if len(queryResult.MeasureResults) > 0 {
		lokiMetricsResponse := getMetricsResponse(queryResult)
		utils.WriteJsonResponse(ctx, lokiMetricsResponse)
		ctx.SetStatusCode(fasthttp.StatusOK)
		return
	}

	lokiQueryResponse := LokiQueryResponse{}

	if len(allJsons) > 0 {
		lokiQueryResponse.Data = Data{ResultType: "streams"}
		lokiQueryResponse.Data.Result = make([]StreamValue, 0)
		for _, row := range allJsons {
			queryResultLine := StreamValue{Values: make([][]string, 0)}
			line, ok := row["line"].(string)
			if !ok {
				responsebody := make(map[string]interface{})
				log.Errorf("ProcessLokiLogsIngestRequest: Unable to convert line to string")
				ctx.SetStatusCode(fasthttp.StatusBadRequest)
				responsebody["error"] = "ProcessLokiLogsIngestRequest: Unable to convert line to string"
				utils.WriteJsonResponse(ctx, responsebody)
			}
			valuesRow := make([]string, 0)
			timeStamp, ok := row["timestamp"].(uint64)
			if !ok {
				responsebody := make(map[string]interface{})
				log.Errorf("ProcessLokiLogsIngestRequest: Unable to convert line to string")
				ctx.SetStatusCode(fasthttp.StatusBadRequest)
				responsebody["error"] = "ProcessLokiLogsIngestRequest: Unable to convert line to string"
				utils.WriteJsonResponse(ctx, responsebody)
			}
			valuesRow = append(valuesRow, fmt.Sprintf("%v", timeStamp*MsToNanoConversion), line)
			queryResultLine.Values = append(queryResultLine.Values, valuesRow)

			newRow := make(map[string]interface{}, 0)
			labelsKeys := remove(allCols, "line")
			for _, label := range labelsKeys {
				if label != TimeStamp && label != Index {
					newRow[label] = row[label]
				}
			}

			queryResultLine.Stream = newRow
			lokiQueryResponse.Data.Result = append(lokiQueryResponse.Data.Result, queryResultLine)
		}
	} else {
		lokiQueryResponse.Data.Result = make([]StreamValue, 0)
	}

	if err != nil {
		log.Errorf(" ProcessMetricsSearchRequest: received empty search request body ")
		responsebody := make(map[string]interface{})
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = "received empty search request body"
		utils.WriteJsonResponse(ctx, responsebody)
		_, err = ctx.WriteString(err.Error())
		if err != nil {
			log.Errorf("qid=%v, ProcessMetricsSearchRequest: could not write error message err=%v", qid, err)
		}
		log.Errorf("qid=%v, ProcessMetricsSearchRequest: failed to decode search request body! Err=%+v", qid, err)
		return
	}

	lokiQueryResponse.Data.Stats = getQueryStats(queryResult, startTime, myid)

	utils.WriteJsonResponse(ctx, lokiQueryResponse)
	ctx.SetStatusCode(fasthttp.StatusOK)

}

func ProcessIndexStatsRequest(ctx *fasthttp.RequestCtx, myid uint64) {
	query := removeUnimplementedMethods(string(ctx.QueryArgs().Peek("query")))

	qid := rutils.GetNextQid()

	ti := structs.InitTableInfo(LOKIINDEX, myid, false)
	simpleNode, aggs, err := pipesearch.ParseQuery(query, qid, "Log QL")
	if err != nil {
		writeEmptyIndexStatsResponse(ctx)
		return
	}

	segment.LogASTNode("logql query parser", simpleNode, qid)
	segment.LogQueryAggsNode("logql aggs parser", aggs, qid)

	simpleNode.TimeRange = rutils.GetESDefaultQueryTimeRange()
	qc := structs.InitQueryContextWithTableInfo(ti, rutils.DefaultBucketCount, 0, myid, false)

	queryResult := segment.ExecuteQuery(simpleNode, aggs, qid, qc)
	allJsons, allCols, err := record.GetJsonFromAllRrc(queryResult.AllRecords, false, qid, queryResult.SegEncToKey, aggs)
	if err != nil {
		writeEmptyIndexStatsResponse(ctx)
		return
	}

	responsebody := make(map[string]interface{})
	responsebody["streams"] = len(allCols)
	responsebody["chunks"] = getChunkCount(queryResult)
	responsebody["entries"] = len(allJsons)
	byteCount := 0
	for _, row := range allJsons {
		lineString, ok := row["line"].(string)
		if !ok {
			writeEmptyIndexStatsResponse(ctx)
			return
		}
		byteCount += len([]byte(lineString))
	}
	responsebody["bytes"] = byteCount
	utils.WriteJsonResponse(ctx, responsebody)
	ctx.SetStatusCode(fasthttp.StatusOK)

}

func ProcessLokiSeriesRequest(ctx *fasthttp.RequestCtx, myid uint64) {
	responsebody := make(map[string]interface{})
	query := string(ctx.QueryArgs().Peek("match[]"))
	startEpoch, err := strconv.ParseUint(string(ctx.QueryArgs().Peek("start")), 10, 64)
	if err != nil {
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}
	endEpoch, err := strconv.ParseUint(string(ctx.QueryArgs().Peek("end")), 10, 64)
	if err != nil {
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}
	responsebody["status"] = "success"

	qid := rutils.GetNextQid()

	ti := structs.InitTableInfo(LOKIINDEX, myid, false)
	simpleNode, aggs, err := pipesearch.ParseQuery(query, qid, "Log QL")
	if err != nil {
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}

	segment.LogASTNode("logql query parser", simpleNode, qid)
	segment.LogQueryAggsNode("logql aggs parser", aggs, qid)

	simpleNode.TimeRange = &dtu.TimeRange{
		StartEpochMs: startEpoch / MsToNanoConversion,
		EndEpochMs:   endEpoch / MsToNanoConversion,
	}
	qc := structs.InitQueryContextWithTableInfo(ti, DefaultLimit, 0, myid, false)
	queryResult := segment.ExecuteQuery(simpleNode, aggs, qid, qc)
	allJsons, _, err := record.GetJsonFromAllRrc(queryResult.AllRecords, false, qid, queryResult.SegEncToKey, aggs)
	if err != nil {
		ctx.SetStatusCode(fasthttp.StatusBadRequest)
		responsebody["error"] = err.Error()
		utils.WriteJsonResponse(ctx, responsebody)
		return
	}
	responsebody["data"] = allJsons
	utils.WriteJsonResponse(ctx, responsebody)
	ctx.SetStatusCode(fasthttp.StatusOK)
}

func writeEmptyIndexStatsResponse(ctx *fasthttp.RequestCtx) {
	responsebody := make(map[string]interface{})
	responsebody["streams"] = 0
	responsebody["chunks"] = 0
	responsebody["entries"] = 0
	responsebody["bytes"] = 0
	utils.WriteJsonResponse(ctx, responsebody)
	ctx.SetStatusCode(fasthttp.StatusOK)
}

func remove(slice []string, stringToRemove string) []string {
	for i, v := range slice {
		if v == stringToRemove {
			return append(slice[:i], slice[i+1:]...)
		}
	}
	return slice
}

// needs to be modified as changes are made to logql parsing. Logfmt, json, label/line expression
// and stream selectors are supported
func removeUnimplementedMethods(queryString string) string {

	// Define the regular expression pattern
	pattern := `sum\s+by\s+\(level\)\s+\(count_over_time\(`
	regex := regexp.MustCompile(pattern)
	matchIndex := regex.FindStringIndex(queryString)
	if matchIndex != nil {
		extractedString := queryString[matchIndex[1]:]
		return extractedString
	} else {
		return queryString
	}
}

func getMetricsResponse(queryResult *structs.NodeResult) LokiMetricsResponse {
	lokiMetricsResponse := LokiMetricsResponse{}
	lokiMetricsResponse.Status = "success"
	lokiMetricsResponse.Data = MetricsData{ResultType: "vector"}
	lokiMetricsResponse.Data.MetricResult = make([]MetricValue, 0)
	if queryResult.MeasureResults != nil {
		metricResult := make([]MetricValue, 0)
		for _, bucket := range queryResult.MeasureResults {
			groupByCols := queryResult.GroupByCols
			newMetricVal := MetricValue{}
			newMetricVal.Stream = make(map[string]interface{})
			for index, colName := range groupByCols {
				newMetricVal.Stream[colName] = bucket.GroupByValues[index]
			}
			valResult := make([]interface{}, 0)
			valResult = append(valResult, 1689919818.158, queryResult.MeasureFunctions[0])
			newMetricVal.Values = valResult
			metricResult = append(metricResult, newMetricVal)
		}
		lokiMetricsResponse.Data.MetricResult = metricResult
	}

	lokiMetricsResponse.Data.Stats = MetricStats{}
	return lokiMetricsResponse
}
func getQueryStats(queryResult *structs.NodeResult, startTime uint64, myid uint64) Stats {
	lokiQueryStats := Stats{}
	if queryResult == nil {
		return lokiQueryStats
	}

	bytesReceivedCount, recordCount, onDiskBytesCount := segwriter.GetVTableCounts(LOKIINDEX, myid)
	unrotatedByteCount, unrotatedEventCount, unrotatedOnDiskBytesCount := segwriter.GetUnrotatedVTableCounts(LOKIINDEX, myid)

	chunkCount := getChunkCount(queryResult)

	ingesterStats := Ingester{}
	ingesterStats.CompressedBytes = int(onDiskBytesCount + unrotatedOnDiskBytesCount)
	ingesterStats.DecompressedBytes = int(bytesReceivedCount + unrotatedByteCount)
	ingesterStats.DecompressedLines = unrotatedEventCount
	ingesterStats.TotalReached = 1 //single node
	ingesterStats.TotalLinesSent = len(queryResult.AllRecords)
	ingesterStats.TotalChunksMatched = chunkCount
	ingesterStats.TotalBatches = chunkCount * search.BLOCK_BATCH_SIZE
	ingesterStats.HeadChunkBytes = int(onDiskBytesCount)
	ingesterStats.HeadChunkLines = int(recordCount)

	lokiQueryStats.Ingester = ingesterStats

	storeStats := Store{}
	storeStats.DecompressedBytes = int(unrotatedOnDiskBytesCount) + int(onDiskBytesCount)
	summaryStats := Summary{}

	summaryStats.TotalBytesProcessed = int(bytesReceivedCount) + int(unrotatedOnDiskBytesCount)
	summaryStats.ExecTime = float64(utils.GetCurrentTimeInMs() - startTime)
	summaryStats.TotalLinesProcessed = len(queryResult.AllRecords)

	return lokiQueryStats
}

func getChunkCount(queryResult *structs.NodeResult) int {
	uniqueChunks := make(map[uint16]bool)
	for _, record := range queryResult.AllRecords {
		uniqueChunks[record.BlockNum] = true
	}
	return len(uniqueChunks)
}