<commit_before>package clickhouse

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/lomik/graphite-clickhouse/helper/version"
	"github.com/lomik/zapwriter"

	"go.uber.org/zap"
)

var ErrUvarintRead = errors.New("ReadUvarint: Malformed array")
var ErrUvarintOverflow = errors.New("ReadUvarint: varint overflows a 64-bit integer")
var ErrClickHouseResponse = errors.New("Malformed response from clickhouse")

type Options struct {
	Timeout        time.Duration
	ConnectTimeout time.Duration
}

func formatSQL(q string) string {
	s := strings.Split(q, "\n")
	for i := 0; i < len(s); i++ {
		s[i] = strings.TrimSpace(s[i])
	}

	return strings.Join(s, " ")
}

func Escape(s string) string {
	s = strings.Replace(s, `\`, `\\`, -1)
	s = strings.Replace(s, `'`, `\'`, -1)
	return s
}

func Query(ctx context.Context, dsn string, query string, table string, opts Options) ([]byte, error) {
	return Post(ctx, dsn, query, table, nil, opts)
}

func Post(ctx context.Context, dsn string, query string, table string, postBody io.Reader, opts Options) ([]byte, error) {
	return do(ctx, dsn, query, table, postBody, false, opts)
}

func PostGzip(ctx context.Context, dsn string, query string, table string, postBody io.Reader, opts Options) ([]byte, error) {
	return do(ctx, dsn, query, table, postBody, true, opts)
}

func Reader(ctx context.Context, dsn string, query string, table string, opts Options) (io.ReadCloser, error) {
	return reader(ctx, dsn, query, table, nil, false, opts)
}

func reader(ctx context.Context, dsn string, query string, table string, postBody io.Reader, gzip bool, opts Options) (bodyReader io.ReadCloser, err error) {
	start := time.Now()

	var requestID string
	if value, ok := ctx.Value("requestID").(string); ok {
		requestID = value
	}

	queryForLogger := query
	if len(queryForLogger) > 500 {
		queryForLogger = queryForLogger[:395] + "<...>" + queryForLogger[len(queryForLogger)-100:]
	}
	logger := zapwriter.Logger("query").With(zap.String("query", formatSQL(queryForLogger)), zap.String("request_id", requestID))

	defer func() {
		d := time.Since(start)
		log := logger.With(
			zap.Duration("time", d),
		)
		// fmt.Println(time.Since(start), formatSQL(queryForLogger))
		if err != nil {
			log.Error("query", zap.Error(err))
		} else {
			log.Info("query")
		}
	}()

	p, err := url.Parse(dsn)
	if err != nil {
		return
	}

	var b [8]byte
	binary.LittleEndian.PutUint64(b[:], rand.Uint64())
	queryID := fmt.Sprintf("%x", b)

	q := p.Query()
	q.Set("query_id", fmt.Sprintf("%s::%s", requestID, queryID))
	p.RawQuery = q.Encode()

	if postBody != nil {
		q := p.Query()
		q.Set("query", query)
		p.RawQuery = q.Encode()
	} else {
		postBody = strings.NewReader(query)
	}

	url := p.String()

	req, err := http.NewRequest("POST", url, postBody)
	if err != nil {
		return
	}

	req.Header.Add("User-Agent", fmt.Sprintf("graphite-clickhouse/%s (table:%s)", version.Version, table))

	if gzip {
		req.Header.Add("Content-Encoding", "gzip")
	}

	client := &http.Client{
		Timeout: opts.Timeout,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: opts.ConnectTimeout,
			}).Dial,
			DisableKeepAlives: true,
		},
	}
	resp, err := client.Do(req)
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		body, _ := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		err = fmt.Errorf("clickhouse response status %d: %s", resp.StatusCode, string(body))
		return
	}

	bodyReader = resp.Body
	return
}

func do(ctx context.Context, dsn string, query string, table string, postBody io.Reader, gzip bool, opts Options) ([]byte, error) {
	bodyReader, err := reader(ctx, dsn, query, table, postBody, gzip, opts)<commit_msg>fix query time logging<commit_after>package clickhouse

import (
	"context"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/url"
	"strings"
	"time"

	"github.com/lomik/graphite-clickhouse/helper/version"
	"github.com/lomik/zapwriter"

	"go.uber.org/zap"
)

var ErrUvarintRead = errors.New("ReadUvarint: Malformed array")
var ErrUvarintOverflow = errors.New("ReadUvarint: varint overflows a 64-bit integer")
var ErrClickHouseResponse = errors.New("Malformed response from clickhouse")

type Options struct {
	Timeout        time.Duration
	ConnectTimeout time.Duration
}

type loggedReader struct {
	reader   io.ReadCloser
	logger   *zap.Logger
	start    time.Time
	finished bool
}

func (r *loggedReader) Read(p []byte) (int, error) {
	n, err := r.reader.Read(p)
	if err != nil && !r.finished {
		r.finished = true
		r.logger.Info("query", zap.Duration("time", time.Since(r.start)))
	}
	return n, err
}

func (r *loggedReader) Close() error {
	err := r.reader.Close()
	if !r.finished {
		r.finished = true
		r.logger.Info("query", zap.Duration("time", time.Since(r.start)))
	}
	return err
}

func formatSQL(q string) string {
	s := strings.Split(q, "\n")
	for i := 0; i < len(s); i++ {
		s[i] = strings.TrimSpace(s[i])
	}

	return strings.Join(s, " ")
}

func Escape(s string) string {
	s = strings.Replace(s, `\`, `\\`, -1)
	s = strings.Replace(s, `'`, `\'`, -1)
	return s
}

func Query(ctx context.Context, dsn string, query string, table string, opts Options) ([]byte, error) {
	return Post(ctx, dsn, query, table, nil, opts)
}

func Post(ctx context.Context, dsn string, query string, table string, postBody io.Reader, opts Options) ([]byte, error) {
	return do(ctx, dsn, query, table, postBody, false, opts)
}

func PostGzip(ctx context.Context, dsn string, query string, table string, postBody io.Reader, opts Options) ([]byte, error) {
	return do(ctx, dsn, query, table, postBody, true, opts)
}

func Reader(ctx context.Context, dsn string, query string, table string, opts Options) (io.ReadCloser, error) {
	return reader(ctx, dsn, query, table, nil, false, opts)
}

func reader(ctx context.Context, dsn string, query string, table string, postBody io.Reader, gzip bool, opts Options) (bodyReader io.ReadCloser, err error) {
	start := time.Now()

	var requestID string
	if value, ok := ctx.Value("requestID").(string); ok {
		requestID = value
	}

	queryForLogger := query
	if len(queryForLogger) > 500 {
		queryForLogger = queryForLogger[:395] + "<...>" + queryForLogger[len(queryForLogger)-100:]
	}
	logger := zapwriter.Logger("query").With(zap.String("query", formatSQL(queryForLogger)), zap.String("request_id", requestID))

	defer func() {
		// fmt.Println(time.Since(start), formatSQL(queryForLogger))
		if err != nil {
			logger.Error("query", zap.Error(err), zap.Duration("time", time.Since(start)))
		}
	}()

	p, err := url.Parse(dsn)
	if err != nil {
		return
	}

	var b [8]byte
	binary.LittleEndian.PutUint64(b[:], rand.Uint64())
	queryID := fmt.Sprintf("%x", b)

	q := p.Query()
	q.Set("query_id", fmt.Sprintf("%s::%s", requestID, queryID))
	p.RawQuery = q.Encode()

	if postBody != nil {
		q := p.Query()
		q.Set("query", query)
		p.RawQuery = q.Encode()
	} else {
		postBody = strings.NewReader(query)
	}

	url := p.String()

	req, err := http.NewRequest("POST", url, postBody)
	if err != nil {
		return
	}

	req.Header.Add("User-Agent", fmt.Sprintf("graphite-clickhouse/%s (table:%s)", version.Version, table))

	if gzip {
		req.Header.Add("Content-Encoding", "gzip")
	}

	client := &http.Client{
		Timeout: opts.Timeout,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: opts.ConnectTimeout,
			}).Dial,
			DisableKeepAlives: true,
		},
	}
	resp, err := client.Do(req)
	if err != nil {
		return
	}

	if resp.StatusCode != 200 {
		body, _ := ioutil.ReadAll(resp.Body)
		resp.Body.Close()
		err = fmt.Errorf("clickhouse response status %d: %s", resp.StatusCode, string(body))
		return
	}

	bodyReader = &loggedReader{
		reader: resp.Body,
		logger: logger,
		start:  start,
	}

	return
}

func do(ctx context.Context, dsn string, query string, table string, postBody io.Reader, gzip bool, opts Options) ([]byte, error) {
	bodyReader, err := reader(ctx, dsn, query, table, postBody, gzip, opts)
{'fraction_non_alphanumeric': 0.16429818000498628, 'fraction_numerical': 0.005484916479680878, 'mean_word_length': 4.172791747259832, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, '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': '2112449', 'n_tokens_mistral': 3071, 'n_tokens_neox': 2926, 'n_words': 731}