mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
380 lines
8.3 KiB
Go
380 lines
8.3 KiB
Go
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
|
|
//
|
|
// 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 paramtable
|
|
|
|
import (
|
|
"os"
|
|
"path"
|
|
"runtime"
|
|
"strconv"
|
|
"strings"
|
|
|
|
"go.uber.org/zap"
|
|
|
|
memkv "github.com/milvus-io/milvus/internal/kv/mem"
|
|
"github.com/milvus-io/milvus/internal/log"
|
|
"github.com/milvus-io/milvus/internal/proto/commonpb"
|
|
"github.com/milvus-io/milvus/internal/util/typeutil"
|
|
"github.com/spf13/cast"
|
|
"github.com/spf13/viper"
|
|
)
|
|
|
|
type UniqueID = typeutil.UniqueID
|
|
|
|
type Base interface {
|
|
Load(key string) (string, error)
|
|
LoadRange(key, endKey string, limit int) ([]string, []string, error)
|
|
LoadYaml(fileName string) error
|
|
Remove(key string) error
|
|
Save(key, value string) error
|
|
Init()
|
|
}
|
|
|
|
type BaseTable struct {
|
|
params *memkv.MemoryKV
|
|
configDir string
|
|
}
|
|
|
|
func (gp *BaseTable) Init() {
|
|
gp.params = memkv.NewMemoryKV()
|
|
|
|
_, fpath, _, _ := runtime.Caller(0)
|
|
configDir := path.Dir(fpath) + "/../../../configs/"
|
|
if _, err := os.Stat(configDir); err != nil {
|
|
log.Warn("cannot access config directory", zap.String("configDir", configDir), zap.Error(err))
|
|
if runPath, err1 := os.Getwd(); err1 != nil {
|
|
panic(err1.Error())
|
|
} else {
|
|
configDir = runPath + "/configs/"
|
|
}
|
|
}
|
|
gp.configDir = configDir
|
|
log.Debug("config directory", zap.String("configDir", gp.configDir))
|
|
|
|
if err := gp.LoadYaml("milvus.yaml"); err != nil {
|
|
panic(err)
|
|
}
|
|
if err := gp.LoadYaml("advanced/common.yaml"); err != nil {
|
|
panic(err)
|
|
}
|
|
if err := gp.LoadYaml("advanced/channel.yaml"); err != nil {
|
|
panic(err)
|
|
}
|
|
gp.tryloadFromEnv()
|
|
}
|
|
|
|
func (gp *BaseTable) GetConfigDir() string {
|
|
return gp.configDir
|
|
}
|
|
|
|
func (gp *BaseTable) LoadFromKVPair(kvPairs []*commonpb.KeyValuePair) error {
|
|
for _, pair := range kvPairs {
|
|
err := gp.Save(pair.Key, pair.Value)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (gp *BaseTable) tryloadFromEnv() {
|
|
minioAddress := os.Getenv("MINIO_ADDRESS")
|
|
if minioAddress == "" {
|
|
minioHost, err := gp.Load("minio.address")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
port, err := gp.Load("minio.port")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
minioAddress = minioHost + ":" + port
|
|
}
|
|
err := gp.Save("_MinioAddress", minioAddress)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
etcdEndpoints := os.Getenv("ETCD_ENDPOINTS")
|
|
if etcdEndpoints == "" {
|
|
etcdEndpoints, err = gp.Load("etcd.endpoints")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
err = gp.Save("_EtcdEndpoints", etcdEndpoints)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
pulsarAddress := os.Getenv("PULSAR_ADDRESS")
|
|
if pulsarAddress == "" {
|
|
pulsarHost, err := gp.Load("pulsar.address")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
port, err := gp.Load("pulsar.port")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
pulsarAddress = "pulsar://" + pulsarHost + ":" + port
|
|
}
|
|
err = gp.Save("_PulsarAddress", pulsarAddress)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
rocksmqPath := os.Getenv("ROCKSMQ_PATH")
|
|
if rocksmqPath == "" {
|
|
path, err := gp.Load("rocksmq.path")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
rocksmqPath = path
|
|
}
|
|
err = gp.Save("_RocksmqPath", rocksmqPath)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
rootCoordAddress := os.Getenv("ROOT_COORD_ADDRESS")
|
|
if rootCoordAddress == "" {
|
|
rootCoordHost, err := gp.Load("rootCoord.address")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
port, err := gp.Load("rootCoord.port")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
rootCoordAddress = rootCoordHost + ":" + port
|
|
}
|
|
err = gp.Save("_RootCoordAddress", rootCoordAddress)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
indexCoordAddress := os.Getenv("INDEX_COORD_ADDRESS")
|
|
if indexCoordAddress == "" {
|
|
indexCoordHost, err := gp.Load("indexCoord.address")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
port, err := gp.Load("indexCoord.port")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
indexCoordAddress = indexCoordHost + ":" + port
|
|
}
|
|
err = gp.Save("_IndexCoordAddress", indexCoordAddress)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
queryCoordAddress := os.Getenv("QUERY_COORD_ADDRESS")
|
|
if queryCoordAddress == "" {
|
|
serviceHost, err := gp.Load("queryCoord.address")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
port, err := gp.Load("queryCoord.port")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
queryCoordAddress = serviceHost + ":" + port
|
|
}
|
|
err = gp.Save("_QueryCoordAddress", queryCoordAddress)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
dataCoordAddress := os.Getenv("DATA_COORD_ADDRESS")
|
|
if dataCoordAddress == "" {
|
|
serviceHost, err := gp.Load("dataCoord.address")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
port, err := gp.Load("dataCoord.port")
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
dataCoordAddress = serviceHost + ":" + port
|
|
}
|
|
err = gp.Save("_DataCoordAddress", dataCoordAddress)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func (gp *BaseTable) Load(key string) (string, error) {
|
|
return gp.params.Load(strings.ToLower(key))
|
|
}
|
|
|
|
func (gp *BaseTable) LoadWithDefault(key string, defaultValue string) (string, error) {
|
|
return gp.params.LoadWithDefault(strings.ToLower(key), defaultValue)
|
|
}
|
|
|
|
func (gp *BaseTable) LoadRange(key, endKey string, limit int) ([]string, []string, error) {
|
|
return gp.params.LoadRange(strings.ToLower(key), strings.ToLower(endKey), limit)
|
|
}
|
|
|
|
func (gp *BaseTable) LoadYaml(fileName string) error {
|
|
config := viper.New()
|
|
configFile := gp.configDir + fileName
|
|
if _, err := os.Stat(configFile); err != nil {
|
|
panic("cannot access config file: " + configFile)
|
|
}
|
|
|
|
config.SetConfigFile(configFile)
|
|
if err := config.ReadInConfig(); err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
for _, key := range config.AllKeys() {
|
|
val := config.Get(key)
|
|
str, err := cast.ToStringE(val)
|
|
if err != nil {
|
|
switch val := val.(type) {
|
|
case []interface{}:
|
|
str = str[:0]
|
|
for _, v := range val {
|
|
ss, err := cast.ToStringE(v)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
if len(str) == 0 {
|
|
str = ss
|
|
} else {
|
|
str = str + "," + ss
|
|
}
|
|
}
|
|
|
|
default:
|
|
panic("undefined config type, key=" + key)
|
|
}
|
|
}
|
|
err = gp.params.Save(strings.ToLower(key), str)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (gp *BaseTable) Remove(key string) error {
|
|
return gp.params.Remove(strings.ToLower(key))
|
|
}
|
|
|
|
func (gp *BaseTable) Save(key, value string) error {
|
|
return gp.params.Save(strings.ToLower(key), value)
|
|
}
|
|
|
|
func (gp *BaseTable) ParseBool(key string, defaultValue bool) bool {
|
|
valueStr, err := gp.LoadWithDefault(key, strconv.FormatBool(defaultValue))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
value, err := strconv.ParseBool(valueStr)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return value
|
|
}
|
|
|
|
func (gp *BaseTable) ParseFloat(key string) float64 {
|
|
valueStr, err := gp.Load(key)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
value, err := strconv.ParseFloat(valueStr, 64)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return value
|
|
}
|
|
|
|
func (gp *BaseTable) ParseInt64(key string) int64 {
|
|
valueStr, err := gp.Load(key)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
value, err := strconv.ParseInt(valueStr, 10, 64)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return value
|
|
}
|
|
|
|
func (gp *BaseTable) ParseInt32(key string) int32 {
|
|
valueStr, err := gp.Load(key)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
value, err := strconv.ParseInt(valueStr, 10, 32)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return int32(value)
|
|
}
|
|
|
|
func (gp *BaseTable) ParseInt(key string) int {
|
|
valueStr, err := gp.Load(key)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
value, err := strconv.Atoi(valueStr)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return value
|
|
}
|
|
|
|
// package methods
|
|
|
|
func ConvertRangeToIntRange(rangeStr, sep string) []int {
|
|
items := strings.Split(rangeStr, sep)
|
|
if len(items) != 2 {
|
|
panic("Illegal range ")
|
|
}
|
|
|
|
startStr := items[0]
|
|
endStr := items[1]
|
|
start, err := strconv.Atoi(startStr)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
end, err := strconv.Atoi(endStr)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
if start < 0 || end < 0 {
|
|
panic("Illegal range value")
|
|
}
|
|
if start > end {
|
|
panic("Illegal range value, start > end")
|
|
}
|
|
return []int{start, end}
|
|
}
|
|
|
|
func ConvertRangeToIntSlice(rangeStr, sep string) []int {
|
|
rangeSlice := ConvertRangeToIntRange(rangeStr, sep)
|
|
start, end := rangeSlice[0], rangeSlice[1]
|
|
var ret []int
|
|
for i := start; i < end; i++ {
|
|
ret = append(ret, i)
|
|
}
|
|
return ret
|
|
}
|