milvus/internal/indexnode/client/client.go

package indexnodeclient

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"time"

	"github.com/zilliztech/milvus-distributed/internal/errors"
	"google.golang.org/grpc"

	"github.com/zilliztech/milvus-distributed/internal/proto/commonpb"
	"github.com/zilliztech/milvus-distributed/internal/proto/indexpb"
	"github.com/zilliztech/milvus-distributed/internal/util/typeutil"
)

type UniqueID = typeutil.UniqueID

type Client struct {
	client  indexpb.IndexServiceClient
	address string
	ctx     context.Context
}

func NewBuildIndexClient(ctx context.Context, address string) (*Client, error) {
	return &Client{
		address: address,
		ctx:     ctx,
	}, nil
}

func parseTS(t int64) time.Time {
	return time.Unix(0, t)
}

func (c *Client) tryConnect() error {
	if c.client != nil {
		return nil
	}
	conn, err := grpc.DialContext(c.ctx, c.address, grpc.WithInsecure(), grpc.WithBlock())
	if err != nil {
		return err
	}
	c.client = indexpb.NewIndexServiceClient(conn)
	return nil
}

func (c *Client) BuildIndex(columnDataPaths []string, typeParams map[string]string, indexParams map[string]string) (UniqueID, error) {
	if c.tryConnect() != nil {
		panic("BuildIndexWithoutID: failed to connect index builder")
	}
	parseMap := func(mStr string) (map[string]string, error) {
		buffer := make(map[string]interface{})
		err := json.Unmarshal([]byte(mStr), &buffer)
		if err != nil {
			return nil, errors.New("Unmarshal params failed")
		}
		ret := make(map[string]string)
		for key, value := range buffer {
			valueStr := fmt.Sprintf("%v", value)
			ret[key] = valueStr
		}
		return ret, nil
	}
	var typeParamsKV []*commonpb.KeyValuePair
	for key := range typeParams {
		if key == "params" {
			mapParams, err := parseMap(typeParams[key])
			if err != nil {
				log.Println("parse params error: ", err)
			}
			for pk, pv := range mapParams {
				typeParamsKV = append(typeParamsKV, &commonpb.KeyValuePair{
					Key:   pk,
					Value: pv,
				})
			}
		} else {
			typeParamsKV = append(typeParamsKV, &commonpb.KeyValuePair{
				Key:   key,
				Value: typeParams[key],
			})
		}
	}

	var indexParamsKV []*commonpb.KeyValuePair
	for key := range indexParams {
		if key == "params" {
			mapParams, err := parseMap(indexParams[key])
			if err != nil {
				log.Println("parse params error: ", err)
			}
			for pk, pv := range mapParams {
				indexParamsKV = append(indexParamsKV, &commonpb.KeyValuePair{
					Key:   pk,
					Value: pv,
				})
			}
		} else {
			indexParamsKV = append(indexParamsKV, &commonpb.KeyValuePair{
				Key:   key,
				Value: indexParams[key],
			})
		}
	}

	ctx := context.TODO()
	requset := &indexpb.BuildIndexRequest{
		DataPaths:   columnDataPaths,
		TypeParams:  typeParamsKV,
		IndexParams: indexParamsKV,
	}
	response, err := c.client.BuildIndex(ctx, requset)
	if err != nil {
		return 0, err
	}

	indexID := response.IndexID
	return indexID, err
}

func (c *Client) GetIndexStates(indexIDs []UniqueID) (*indexpb.IndexStatesResponse, error) {
	if c.tryConnect() != nil {
		panic("DescribeIndex: failed to connect index builder")
	}
	ctx := context.TODO()
	request := &indexpb.IndexStatesRequest{
		IndexID: indexIDs,
	}

	response, err := c.client.GetIndexStates(ctx, request)
	return response, err
}

func (c *Client) GetIndexFilePaths(indexID UniqueID) ([]string, error) {
	if c.tryConnect() != nil {
		panic("GetIndexFilePaths: failed to connect index builder")
	}
	ctx := context.TODO()
	request := &indexpb.IndexFilePathRequest{
		IndexID: indexID,
	}

	response, err := c.client.GetIndexFilePaths(ctx, request)
	if err != nil {
		return nil, err
	}

	return response.IndexFilePaths, nil
}