<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/storage/remote"
)

// Escaping for strings for Crate.io SQL.
var escaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"", "'", "\\'")

// Escape a string for use in SQL.
func escape(s, quote string) string {
	return quote + escaper.Replace(s) + quote
}

type crateRequest struct {
	Stmt     string          `json:"stmt"`
	BulkArgs [][]interface{} `json:"bulk_args,omitempty"`
}

type crateResponse struct {
	Cols []string        `json:"cols,omitempty"`
	Rows [][]interface{} `json:"rows,omitempty"`
}

// Convert a read query into a Crate SQL query.
func queryToSQL(q *remote.Query) string {
	selectors := make([]string, 0, len(q.Matchers)+2)
	for _, m := range q.Matchers {
		switch m.Type {
		case remote.MatchType_EQUAL:
			if m.Value == "" {
				// Empty labels are recorded as NULL.
				// In PromQL, empty labels and missing labels are the same thing.
				selectors = append(selectors, fmt.Sprintf("(%s IS NULL)", escape("l"+m.Name, "\"")))
			} else {
				selectors = append(selectors, fmt.Sprintf("(%s = %s)", escape("l"+m.Name, "\""), escape(m.Value, "'")))
			}
		case remote.MatchType_NOT_EQUAL:
			if m.Value == "" {
				selectors = append(selectors, fmt.Sprintf("(%s IS NOT NULL)", escape("l"+m.Name, "\"")))
			} else {
				selectors = append(selectors, fmt.Sprintf("(%s != %s)", escape("l"+m.Name, "\""), escape(m.Value, "'")))
			}
		case remote.MatchType_REGEX_MATCH:
			// XXX Handle that null == empty string.
			// Crate regexes are not RE2, so there may be small semantic differences here.
			selectors = append(selectors, fmt.Sprintf("(%s ~ %s)", escape("l"+m.Name, "\""), escape(m.Value, "'")))
		case remote.MatchType_REGEX_NO_MATCH:
			selectors = append(selectors, fmt.Sprintf("(%s !~ %s)", escape("l"+m.Name, "\""), escape(m.Value, "'")))
		}
	}
	selectors = append(selectors, fmt.Sprintf("(timestamp <= %d)", q.EndTimestampMs))
	selectors = append(selectors, fmt.Sprintf("(timestamp >= %d)", q.StartTimestampMs))

	return fmt.Sprintf("SELECT * from metrics WHERE %s ORDER BY timestamp", strings.Join(selectors, " AND "))
}

func runQuery(q *remote.Query) []*remote.TimeSeries {
	resp := []*remote.TimeSeries{}

	query := queryToSQL(q)

	request := crateRequest{Stmt: query}
	jsonRequest, err := json.Marshal(request)
	if err != nil {
		return resp
	}

	result, err := http.Post("http://localhost:4200/_sql", "application/json", bytes.NewReader(jsonRequest))
	if err != nil {
		return resp
	}
	defer result.Body.Close()
	if result.StatusCode != http.StatusOK {
		return resp
	}
	var data crateResponse
	decoder := json.NewDecoder(result.Body)
	decoder.UseNumber()
	if err = decoder.Decode(&data); err != nil {
		return resp
	}

	timeseries := map[string]*remote.TimeSeries{}
	for _, row := range data.Rows {
		metric := model.Metric{}
		var v float64
		var t int64
		for i, value := range row {
			column := data.Cols[i]
			if column[0] == 'l' && value != nil {
				metric[model.LabelName(column[1:])] = model.LabelValue(value.(string))
			} else if column == "timestamp" {
				t, _ = value.(json.Number).Int64()
			} else if column == "valueRaw" {
				val, _ := value.(json.Number).Int64()
				v = math.Float64frombits(uint64(val))
			}
		}
		ts, ok := timeseries[metric.String()]
		if !ok {
			ts = &remote.TimeSeries{}
			for k, v := range metric {
				ts.Labels = append(ts.Labels, &remote.LabelPair{Name: string(k), Value: string(v)})
			}
			timeseries[metric.String()] = ts
		}
		ts.Samples = append(ts.Samples, &remote.Sample{Value: v, TimestampMs: t})
	}

	for _, ts := range timeseries {
		resp = append(resp, ts)
	}
	return resp
}

func main() {
	http.HandleFunc("/write", func(w http.ResponseWriter, r *http.Request) {
		compressed, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		reqBuf, err := snappy.Decode(nil, compressed)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		var req remote.WriteRequest
		if err := proto.Unmarshal(reqBuf, &req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		// Build a list of every label name used.
		labelsUsed := map[string]struct{}{}
		for _, ts := range req.Timeseries {
			for _, l := range ts.Labels {
				labelsUsed[l.Name] = struct{}{}
			}
		}
		labels := make([]string, 0, len(labelsUsed))
		escapedLabels := make([]string, 0, len(labelsUsed))

		for l := range labelsUsed {
			labels = append(labels, l)
			escapedLabels = append(escapedLabels, escape("l"+l, "\""))
		}

		request := crateRequest{
			BulkArgs: make([][]interface{}, 0, len(req.Timeseries)),
		}
		placeholders := strings.Repeat("?, ", len(labels))
		columns := strings.Join(escapedLabels, ", ")
		request.Stmt = fmt.Sprintf("INSERT INTO metrics (%s, \"value\", \"valueRaw\", \"timestamp\") VALUES (%s ?, ?, ?)", columns, placeholders)

		for _, ts := range req.Timeseries {
			metric := make(map[string]string, len(ts.Labels))
			for _, l := range ts.Labels {
				metric[l.Name] = l.Value
			}

			for _, s := range ts.Samples {
				args := make([]interface{}, 0, len(labels)+2)
				for _, l := range labels {
					if metric[l] == "" {
						args = append(args, nil)<commit_msg>Refactor escaping.<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"math"
	"net/http"
	"strings"

	"github.com/golang/protobuf/proto"
	"github.com/golang/snappy"
	"github.com/prometheus/common/model"
	"github.com/prometheus/prometheus/storage/remote"
)

// Escaping for strings for Crate.io SQL.
var escaper = strings.NewReplacer("\\", "\\\\", "\"", "\\\"", "'", "\\'")

// Escape a labelname for use in SQL as a column name.
func escapeLabelName(s string) string {
	return "\"l" + escaper.Replace(s) + "\""
}

// Escape a labelvalue for use in SQL as a string value.
func escapeLabelValue(s string) string {
	return "'" + escaper.Replace(s) + "'"
}

type crateRequest struct {
	Stmt     string          `json:"stmt"`
	BulkArgs [][]interface{} `json:"bulk_args,omitempty"`
}

type crateResponse struct {
	Cols []string        `json:"cols,omitempty"`
	Rows [][]interface{} `json:"rows,omitempty"`
}

// Convert a read query into a Crate SQL query.
func queryToSQL(q *remote.Query) string {
	selectors := make([]string, 0, len(q.Matchers)+2)
	for _, m := range q.Matchers {
		switch m.Type {
		case remote.MatchType_EQUAL:
			if m.Value == "" {
				// Empty labels are recorded as NULL.
				// In PromQL, empty labels and missing labels are the same thing.
				selectors = append(selectors, fmt.Sprintf("(%s IS NULL)", escapeLabelName(m.Name)))
			} else {
				selectors = append(selectors, fmt.Sprintf("(%s = %s)", escapeLabelName(m.Name), escapeLabelValue(m.Value)))
			}
		case remote.MatchType_NOT_EQUAL:
			if m.Value == "" {
				selectors = append(selectors, fmt.Sprintf("(%s IS NOT NULL)", escapeLabelName(m.Name)))
			} else {
				selectors = append(selectors, fmt.Sprintf("(%s != %s)", escapeLabelName(m.Name), escapeLabelValue(m.Value)))
			}
		case remote.MatchType_REGEX_MATCH:
			// XXX Handle that null == empty string.
			// Crate regexes are not RE2, so there may be small semantic differences here.
			selectors = append(selectors, fmt.Sprintf("(%s ~ %s)", escapeLabelName(m.Name), escapeLabelValue(m.Value)))
		case remote.MatchType_REGEX_NO_MATCH:
			selectors = append(selectors, fmt.Sprintf("(%s !~ %s)", escapeLabelName(m.Name), escapeLabelValue(m.Value)))
		}
	}
	selectors = append(selectors, fmt.Sprintf("(timestamp <= %d)", q.EndTimestampMs))
	selectors = append(selectors, fmt.Sprintf("(timestamp >= %d)", q.StartTimestampMs))

	return fmt.Sprintf("SELECT * from metrics WHERE %s ORDER BY timestamp", strings.Join(selectors, " AND "))
}

func runQuery(q *remote.Query) []*remote.TimeSeries {
	resp := []*remote.TimeSeries{}

	query := queryToSQL(q)

	request := crateRequest{Stmt: query}
	jsonRequest, err := json.Marshal(request)
	if err != nil {
		return resp
	}

	result, err := http.Post("http://localhost:4200/_sql", "application/json", bytes.NewReader(jsonRequest))
	if err != nil {
		return resp
	}
	defer result.Body.Close()
	if result.StatusCode != http.StatusOK {
		return resp
	}
	var data crateResponse
	decoder := json.NewDecoder(result.Body)
	decoder.UseNumber()
	if err = decoder.Decode(&data); err != nil {
		return resp
	}

	timeseries := map[string]*remote.TimeSeries{}
	for _, row := range data.Rows {
		metric := model.Metric{}
		var v float64
		var t int64
		for i, value := range row {
			column := data.Cols[i]
			if column[0] == 'l' && value != nil {
				metric[model.LabelName(column[1:])] = model.LabelValue(value.(string))
			} else if column == "timestamp" {
				t, _ = value.(json.Number).Int64()
			} else if column == "valueRaw" {
				val, _ := value.(json.Number).Int64()
				v = math.Float64frombits(uint64(val))
			}
		}
		ts, ok := timeseries[metric.String()]
		if !ok {
			ts = &remote.TimeSeries{}
			for k, v := range metric {
				ts.Labels = append(ts.Labels, &remote.LabelPair{Name: string(k), Value: string(v)})
			}
			timeseries[metric.String()] = ts
		}
		ts.Samples = append(ts.Samples, &remote.Sample{Value: v, TimestampMs: t})
	}

	for _, ts := range timeseries {
		resp = append(resp, ts)
	}
	return resp
}

func main() {
	http.HandleFunc("/write", func(w http.ResponseWriter, r *http.Request) {
		compressed, err := ioutil.ReadAll(r.Body)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		reqBuf, err := snappy.Decode(nil, compressed)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		var req remote.WriteRequest
		if err := proto.Unmarshal(reqBuf, &req); err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		// Build a list of every label name used.
		labelsUsed := map[string]struct{}{}
		for _, ts := range req.Timeseries {
			for _, l := range ts.Labels {
				labelsUsed[l.Name] = struct{}{}
			}
		}
		labels := make([]string, 0, len(labelsUsed))
		escapedLabels := make([]string, 0, len(labelsUsed))

		for l := range labelsUsed {
			labels = append(labels, l)
			escapedLabels = append(escapedLabels, escapeLabelName(l))
		}

		request := crateRequest{
			BulkArgs: make([][]interface{}, 0, len(req.Timeseries)),
		}
		placeholders := strings.Repeat("?, ", len(labels))
		columns := strings.Join(escapedLabels, ", ")
		request.Stmt = fmt.Sprintf("INSERT INTO metrics (%s, \"value\", \"valueRaw\", \"timestamp\") VALUES (%s ?, ?, ?)", columns, placeholders)

		for _, ts := range req.Timeseries {
			metric := make(map[string]string, len(ts.Labels))
			for _, l := range ts.Labels {
				metric[l.Name] = l.Value
			}

			for _, s := range ts.Samples {
				args := make([]interface{}, 0, len(labels)+2)
				for _, l := range labels {
					if metric[l] == "" {
						args = append(args, nil)
{'fraction_non_alphanumeric': 0.173071706302559, 'fraction_numerical': 0.004702052626819784, 'mean_word_length': 3.6983857264231097, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3155671', 'n_tokens_mistral': 4182, 'n_tokens_neox': 3747, 'n_words': 911}