github.com/sequix/cortex@v1.1.6/pkg/chunk/ts/client.go

package ts

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net/url"
	"strings"
	"time"

	"github.com/baidubce/bce-sdk-go/auth"
	"github.com/baidubce/bce-sdk-go/http"
	"github.com/go-kit/kit/log/level"

	"github.com/sequix/cortex/pkg/chunk"
	"github.com/sequix/cortex/pkg/util"
)

var (
	maxChunks = 5
	signer    = &auth.BceV1Signer{}
)

type Config struct {
	Endpoint  string `yaml:"endpoint"`
	AccessKey string `yaml:"ak"`
	SecretKey string `yaml:"sk"`
	Instance  string `yaml:"instance"`
	Table     string `yaml:"table"`
}

type client struct {
	credential    *auth.BceCredentials
	endpoint      string
	uri           string
	headersToSign map[string]struct{}
}

func New(c Config) (*client, error) {
	cred, err := auth.NewBceCredentials(c.AccessKey, c.SecretKey)
	if err != nil {
		return nil, err
	}

	client := &client{
		credential: cred,
		endpoint:   c.Endpoint,
		uri:        "/v1/instance/" + c.Instance + "/table/" + c.Table + "/rows",
		headersToSign: map[string]struct{}{
			"host":       struct{}{},
			"x-bce-date": struct{}{},
		},
	}
	return client, nil
}

func (c *client) Stop() {
	// 为了实现 type chunk.ObjectClient interface
}

func splitReq(chunks []chunk.Chunk, max int, f func([]chunk.Chunk) error) error {
	n := len(chunks)
	for i := 0; i < n; i += maxChunks {
		j := i + max
		if j > n {
			j = n
		}
		if err := f(chunks[i:j]); err != nil {
			return err
		}
	}
	return nil
}

func (c *client) PutChunks(ctx context.Context, chunks []chunk.Chunk) error {
	return splitReq(chunks, maxChunks, c.put)
}

func (c *client) GetChunks(ctx context.Context, chunks []chunk.Chunk) ([]chunk.Chunk, error) {
	if err := splitReq(chunks, maxChunks, c.get); err != nil {
		return nil, err
	}
	return chunks, nil
}

func (c *client) newRequest(now time.Time) *http.Request {
	req := &http.Request{}
	req.SetProtocol("https")
	req.SetEndpoint(c.endpoint)
	req.SetHeader("host", c.endpoint)
	req.SetHeader("x-bce-date", now.Format(time.RFC3339))
	req.SetHeader("Content-Type", "application/json")
	//req.SetHeader("Accept", "application/json")
	req.SetUri(c.uri)
	return req
}

func (c *client) put(chunks []chunk.Chunk) error {
	reqBody := &batchPutRowReq{
		Rows: make([]*row, 0, len(chunks)),
	}

	for i := range chunks {
		r, err := chunkToRow(&chunks[i])
		if err != nil {
			return err
		}
		reqBody.Rows = append(reqBody.Rows, r)
	}

	reqBodyBuf, err := json.Marshal(reqBody)
	if err != nil {
		return err
	}

	now := time.Now().UTC()
	req := c.newRequest(now)
	req.SetMethod(http.PUT)
	req.SetLength(int64(len(reqBodyBuf)))
	req.SetBody(ioutil.NopCloser(bytes.NewReader(reqBodyBuf)))

	base := time.Now()
	_, err = c.do(now, req)
	level.Debug(util.Logger).Log("module", "TableStorageClient", "msg", "Put Rows", "size", len(reqBodyBuf), "duration", time.Since(base), "err", err)
	if err != nil {
		return err
	}
	return nil
}

func (c *client) get(chunks []chunk.Chunk) error {
	reqBody := &batchGetRowReq{
		Rows: make([]*row, 0, len(chunks)),
	}

	for i := range chunks {
		reqBody.Rows = append(reqBody.Rows, chunkToRowWithoutContent(&chunks[i]))
	}

	reqBodyBuf, err := json.Marshal(reqBody)
	if err != nil {
		return err
	}

	now := time.Now().UTC()
	req := c.newRequest(now)
	req.SetMethod(http.GET)
	req.SetLength(int64(len(reqBodyBuf)))
	req.SetBody(ioutil.NopCloser(bytes.NewReader(reqBodyBuf)))

	base := time.Now()
	rsp, err := c.do(now, req)
	level.Debug(util.Logger).Log("module", "TableStorageClient", "msg", "Get Rows", "size", len(reqBodyBuf), "duration", time.Since(base), "err", err)
	if err != nil {
		return err
	}
	defer rsp.Body().Close()

	rspBodyBytes, err := ioutil.ReadAll(rsp.Body())
	if err != nil {
		return err
	}

	rspBody := &batchGetRowRsp{}
	err = json.Unmarshal(rspBodyBytes, rspBody)
	if err != nil {
		return err
	}
	rows := rspBody.Rows

	keyToRow := make(map[string]*row)
	for i := range rows {
		keyToRow[rawUrlDecode(rows[i].RowKey)] = rows[i]
	}

	for i := range chunks {
		r, ok := keyToRow[chunks[i].ExternalKey()]
		if !ok {
			continue
		}
		err := rowToChunk(r, &chunks[i])
		if err != nil {
			return err
		}
	}
	return nil
}

func (c *client) do(now time.Time, req *http.Request) (*http.Response, error) {
	signer.Sign(req, c.credential, &auth.SignOptions{
		HeadersToSign: c.headersToSign,
		Timestamp:     now.Unix(),
		ExpireSeconds: 180,
	})
	rsp, origErr := http.Execute(req)
	if origErr != nil {
		return rsp, nil
	}

	if rsp.StatusCode() > 399 {
		content, err := ioutil.ReadAll(io.LimitReader(rsp.Body(), 256))
		if err != nil {
			return nil, fmt.Errorf("ts client reading response body failed")
		}
		return nil, fmt.Errorf("ts client %s rows: %s", req.Method(), content)
	}
	return rsp, nil
}

func rawUrlEncode(s string) string {
	es := url.QueryEscape(s)
	return strings.ReplaceAll(es, "+", "%20")
}

func rawUrlDecode(s string) string {
	rs := strings.ReplaceAll(s, "%20", "+")
	ds, _ := url.QueryUnescape(rs)
	return ds
}

func chunkToRow(chunk *chunk.Chunk) (*row, error) {
	r := &row{
		RowKey: rawUrlEncode(chunk.ExternalKey()),
		Cells: []*cell{
			{
				Column: columnName,
			},
		},
	}
	buf, err := chunk.Encoded()
	if err != nil {
		return nil, err
	}
	r.Cells[0].Value = rawUrlEncode(string(buf))
	return r, nil
}

func chunkToRowWithoutContent(chunk *chunk.Chunk) *row {
	return &row{
		RowKey: rawUrlEncode(chunk.ExternalKey()),
		Cells: []*cell{
			{
				Column: columnName,
			},
		},
	}
}

func rowToChunk(r *row, c *chunk.Chunk) error {
	decodeCtx := chunk.NewDecodeContext()
	return c.Decode(decodeCtx, []byte(rawUrlDecode(r.Cells[0].Value)))
}