2021-11-10 23:55:32 +08:00
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
2021-04-19 10:09:43 +08:00
// with the License. You may obtain a copy of the License at
//
2021-11-10 23:55:32 +08:00
// http://www.apache.org/licenses/LICENSE-2.0
2021-04-19 10:09:43 +08:00
//
2021-11-10 23:55:32 +08:00
// 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.
2021-04-19 10:09:43 +08:00
2021-06-22 14:40:07 +08:00
package proxy
2020-11-26 16:01:31 +08:00
import (
2022-08-04 11:04:34 +08:00
"context"
2021-03-05 10:15:27 +08:00
"fmt"
2020-11-26 16:01:31 +08:00
"strconv"
"strings"
2022-05-24 12:05:59 +08:00
"time"
2020-11-26 16:01:31 +08:00
2023-02-26 11:31:49 +08:00
"github.com/cockroachdb/errors"
2022-08-04 11:04:34 +08:00
"go.uber.org/zap"
2022-08-25 15:48:54 +08:00
"golang.org/x/crypto/bcrypt"
2022-08-04 11:04:34 +08:00
"google.golang.org/grpc/metadata"
2022-10-16 20:49:27 +08:00
"github.com/milvus-io/milvus-proto/go-api/commonpb"
2023-01-04 17:21:36 +08:00
"github.com/milvus-io/milvus-proto/go-api/milvuspb"
2022-10-16 20:49:27 +08:00
"github.com/milvus-io/milvus-proto/go-api/schemapb"
2023-04-06 19:14:32 +08:00
"github.com/milvus-io/milvus/internal/proto/querypb"
"github.com/milvus-io/milvus/internal/types"
2023-05-16 17:41:22 +08:00
"github.com/milvus-io/milvus/pkg/common"
2023-04-06 19:14:32 +08:00
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/mq/msgstream"
"github.com/milvus-io/milvus/pkg/util"
"github.com/milvus-io/milvus/pkg/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/util/crypto"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/tsoutil"
"github.com/milvus-io/milvus/pkg/util/typeutil"
2020-11-26 16:01:31 +08:00
)
2022-08-25 15:48:54 +08:00
const (
strongTS = 0
boundedTS = 2
// enableMultipleVectorFields indicates whether to enable multiple vector fields.
enableMultipleVectorFields = false
2022-05-24 12:05:59 +08:00
2022-08-25 15:48:54 +08:00
defaultMaxVarCharLength = 65535
// DefaultIndexType name of default index type for scalar field
DefaultIndexType = "STL_SORT"
// DefaultStringIndexType name of default index type for varChar/string field
DefaultStringIndexType = "Trie"
)
2022-03-14 23:20:02 +08:00
2022-08-04 11:04:34 +08:00
var logger = log . L ( ) . WithOptions ( zap . Fields ( zap . String ( "role" , typeutil . ProxyRole ) ) )
2021-12-24 19:27:38 +08:00
// isAlpha check if c is alpha.
2020-11-26 16:01:31 +08:00
func isAlpha ( c uint8 ) bool {
if ( c < 'A' || c > 'Z' ) && ( c < 'a' || c > 'z' ) {
return false
}
return true
}
2021-12-17 21:52:46 +08:00
// isNumber check if c is a number.
2020-11-26 16:01:31 +08:00
func isNumber ( c uint8 ) bool {
if c < '0' || c > '9' {
return false
}
return true
}
2022-09-26 18:00:57 +08:00
func validateLimit ( limit int64 ) error {
2022-12-15 14:03:22 +08:00
if limit <= 0 || limit > Params . CommonCfg . TopKLimit . GetAsInt64 ( ) {
return fmt . Errorf ( "should be in range [1, %d], but got %d" , Params . CommonCfg . TopKLimit . GetAsInt64 ( ) , limit )
2022-05-24 21:27:59 +08:00
}
return nil
}
2021-09-18 11:13:51 +08:00
func validateCollectionNameOrAlias ( entity , entityType string ) error {
entity = strings . TrimSpace ( entity )
2020-11-26 16:01:31 +08:00
2021-09-18 11:13:51 +08:00
if entity == "" {
2021-11-29 12:25:17 +08:00
return fmt . Errorf ( "collection %s should not be empty" , entityType )
2020-11-26 16:01:31 +08:00
}
2021-09-18 11:13:51 +08:00
invalidMsg := fmt . Sprintf ( "Invalid collection %s: %s. " , entityType , entity )
2022-12-07 18:01:19 +08:00
if len ( entity ) > Params . ProxyCfg . MaxNameLength . GetAsInt ( ) {
2023-02-07 17:52:31 +08:00
return fmt . Errorf ( "%s the length of a collection %s must be less than %s characters" , invalidMsg , entityType ,
Params . ProxyCfg . MaxNameLength . GetValue ( ) )
2020-11-26 16:01:31 +08:00
}
2021-09-18 11:13:51 +08:00
firstChar := entity [ 0 ]
2020-11-26 16:01:31 +08:00
if firstChar != '_' && ! isAlpha ( firstChar ) {
2023-02-07 17:52:31 +08:00
return fmt . Errorf ( "%s the first character of a collection %s must be an underscore or letter" , invalidMsg , entityType )
2020-11-26 16:01:31 +08:00
}
2021-09-18 11:13:51 +08:00
for i := 1 ; i < len ( entity ) ; i ++ {
c := entity [ i ]
2022-09-08 10:16:34 +08:00
if c != '_' && ! isAlpha ( c ) && ! isNumber ( c ) {
2023-02-07 17:52:31 +08:00
return fmt . Errorf ( "%s collection %s can only contain numbers, letters and underscores" , invalidMsg , entityType )
}
}
return nil
}
func ValidateResourceGroupName ( entity string ) error {
if entity == "" {
2023-02-16 10:48:34 +08:00
return errors . New ( "resource group name couldn't be empty" )
2023-02-07 17:52:31 +08:00
}
invalidMsg := fmt . Sprintf ( "Invalid resource group name %s." , entity )
if len ( entity ) > Params . ProxyCfg . MaxNameLength . GetAsInt ( ) {
return fmt . Errorf ( "%s the length of a resource group name must be less than %s characters" ,
invalidMsg , Params . ProxyCfg . MaxNameLength . GetValue ( ) )
}
firstChar := entity [ 0 ]
if firstChar != '_' && ! isAlpha ( firstChar ) {
return fmt . Errorf ( "%s the first character of a resource group name must be an underscore or letter" , invalidMsg )
}
for i := 1 ; i < len ( entity ) ; i ++ {
c := entity [ i ]
if c != '_' && ! isAlpha ( c ) && ! isNumber ( c ) {
return fmt . Errorf ( "%s resource group name can only contain numbers, letters and underscores" , invalidMsg )
2020-11-26 16:01:31 +08:00
}
}
return nil
}
2021-10-02 12:04:12 +08:00
// ValidateCollectionAlias returns true if collAlias is a valid alias name for collection, otherwise returns false.
2021-09-18 11:13:51 +08:00
func ValidateCollectionAlias ( collAlias string ) error {
return validateCollectionNameOrAlias ( collAlias , "alias" )
}
2021-10-23 10:53:12 +08:00
func validateCollectionName ( collName string ) error {
2021-09-18 11:13:51 +08:00
return validateCollectionNameOrAlias ( collName , "name" )
}
2021-10-23 18:23:44 +08:00
func validatePartitionTag ( partitionTag string , strictCheck bool ) error {
2020-11-26 16:01:31 +08:00
partitionTag = strings . TrimSpace ( partitionTag )
2021-06-17 16:08:01 +08:00
invalidMsg := "Invalid partition name: " + partitionTag + ". "
2020-11-26 16:01:31 +08:00
if partitionTag == "" {
2021-06-17 16:08:01 +08:00
msg := invalidMsg + "Partition name should not be empty."
2020-11-26 16:01:31 +08:00
return errors . New ( msg )
}
2022-12-07 18:01:19 +08:00
if len ( partitionTag ) > Params . ProxyCfg . MaxNameLength . GetAsInt ( ) {
msg := invalidMsg + "The length of a partition name must be less than " + Params . ProxyCfg . MaxNameLength . GetValue ( ) + " characters."
2020-11-26 16:01:31 +08:00
return errors . New ( msg )
}
if strictCheck {
firstChar := partitionTag [ 0 ]
if firstChar != '_' && ! isAlpha ( firstChar ) && ! isNumber ( firstChar ) {
2022-09-08 10:16:34 +08:00
msg := invalidMsg + "The first character of a partition name must be an underscore or letter."
2020-11-26 16:01:31 +08:00
return errors . New ( msg )
}
tagSize := len ( partitionTag )
for i := 1 ; i < tagSize ; i ++ {
c := partitionTag [ i ]
2022-09-08 10:16:34 +08:00
if c != '_' && ! isAlpha ( c ) && ! isNumber ( c ) {
msg := invalidMsg + "Partition name can only contain numbers, letters and underscores."
2020-11-26 16:01:31 +08:00
return errors . New ( msg )
}
}
}
return nil
}
2021-10-23 18:25:37 +08:00
func validateFieldName ( fieldName string ) error {
2020-11-26 16:01:31 +08:00
fieldName = strings . TrimSpace ( fieldName )
if fieldName == "" {
2021-11-29 12:25:17 +08:00
return errors . New ( "field name should not be empty" )
2020-11-26 16:01:31 +08:00
}
invalidMsg := "Invalid field name: " + fieldName + ". "
2022-12-07 18:01:19 +08:00
if len ( fieldName ) > Params . ProxyCfg . MaxNameLength . GetAsInt ( ) {
msg := invalidMsg + "The length of a field name must be less than " + Params . ProxyCfg . MaxNameLength . GetValue ( ) + " characters."
2020-11-26 16:01:31 +08:00
return errors . New ( msg )
}
firstChar := fieldName [ 0 ]
if firstChar != '_' && ! isAlpha ( firstChar ) {
msg := invalidMsg + "The first character of a field name must be an underscore or letter."
return errors . New ( msg )
}
fieldNameSize := len ( fieldName )
for i := 1 ; i < fieldNameSize ; i ++ {
c := fieldName [ i ]
if c != '_' && ! isAlpha ( c ) && ! isNumber ( c ) {
msg := invalidMsg + "Field name cannot only contain numbers, letters, and underscores."
return errors . New ( msg )
}
}
return nil
}
2022-04-29 13:35:49 +08:00
func validateDimension ( field * schemapb . FieldSchema ) error {
exist := false
var dim int64
for _ , param := range field . TypeParams {
2023-05-16 17:41:22 +08:00
if param . Key == common . DimKey {
2022-04-29 13:35:49 +08:00
exist = true
tmp , err := strconv . ParseInt ( param . Value , 10 , 64 )
if err != nil {
return err
}
dim = tmp
break
}
}
if ! exist {
return errors . New ( "dimension is not defined in field type params, check type param `dim` for vector field" )
}
2022-12-07 18:01:19 +08:00
if dim <= 0 || dim > Params . ProxyCfg . MaxDimension . GetAsInt64 ( ) {
return fmt . Errorf ( "invalid dimension: %d. should be in range 1 ~ %d" , dim , Params . ProxyCfg . MaxDimension . GetAsInt ( ) )
2020-11-26 16:01:31 +08:00
}
2022-04-29 13:35:49 +08:00
if field . DataType == schemapb . DataType_BinaryVector && dim % 8 != 0 {
2021-03-05 10:15:27 +08:00
return fmt . Errorf ( "invalid dimension: %d. should be multiple of 8. " , dim )
2020-11-26 16:01:31 +08:00
}
return nil
}
2020-11-30 19:38:23 +08:00
2022-04-29 13:35:49 +08:00
func validateMaxLengthPerRow ( collectionName string , field * schemapb . FieldSchema ) error {
exist := false
for _ , param := range field . TypeParams {
2023-05-16 17:41:22 +08:00
if param . Key != common . MaxLengthKey {
2022-06-07 15:58:06 +08:00
return fmt . Errorf ( "type param key(max_length) should be specified for varChar field, not %s" , param . Key )
2022-04-29 13:35:49 +08:00
}
maxLengthPerRow , err := strconv . ParseInt ( param . Value , 10 , 64 )
if err != nil {
return err
}
if maxLengthPerRow > defaultMaxVarCharLength || maxLengthPerRow <= 0 {
return fmt . Errorf ( "the maximum length specified for a VarChar shoule be in (0, 65535]" )
}
exist = true
}
2022-06-07 15:58:06 +08:00
// if not exist type params max_length, return error
2022-04-29 13:35:49 +08:00
if ! exist {
2022-06-07 15:58:06 +08:00
return fmt . Errorf ( "type param(max_length) should be specified for varChar field of collection %s" , collectionName )
2022-04-29 13:35:49 +08:00
}
return nil
}
2021-10-25 23:32:20 +08:00
func validateVectorFieldMetricType ( field * schemapb . FieldSchema ) error {
2021-03-12 14:22:09 +08:00
if ( field . DataType != schemapb . DataType_FloatVector ) && ( field . DataType != schemapb . DataType_BinaryVector ) {
2020-11-30 19:38:23 +08:00
return nil
}
for _ , params := range field . IndexParams {
2023-05-16 17:41:22 +08:00
if params . Key == common . MetricTypeKey {
2020-11-30 19:38:23 +08:00
return nil
}
}
return errors . New ( "vector float without metric_type" )
}
2021-10-26 10:38:41 +08:00
func validateDuplicatedFieldName ( fields [ ] * schemapb . FieldSchema ) error {
2020-11-30 19:38:23 +08:00
names := make ( map [ string ] bool )
for _ , field := range fields {
_ , ok := names [ field . Name ]
if ok {
2021-03-08 12:41:46 +08:00
return errors . New ( "duplicated field name" )
2020-11-30 19:38:23 +08:00
}
names [ field . Name ] = true
}
return nil
}
2022-03-03 16:57:56 +08:00
func validateFieldType ( schema * schemapb . CollectionSchema ) error {
for _ , field := range schema . GetFields ( ) {
switch field . GetDataType ( ) {
case schemapb . DataType_String :
2022-03-14 23:20:02 +08:00
return errors . New ( "string data type not supported yet, please use VarChar type instead" )
2022-03-03 16:57:56 +08:00
case schemapb . DataType_None :
return errors . New ( "data type None is not valid" )
}
}
return nil
}
2022-10-12 18:37:23 +08:00
// ValidateFieldAutoID call after validatePrimaryKey
2021-06-21 11:42:18 +08:00
func ValidateFieldAutoID ( coll * schemapb . CollectionSchema ) error {
var idx = - 1
for i , field := range coll . Fields {
if field . AutoID {
if idx != - 1 {
return fmt . Errorf ( "only one field can speficy AutoID with true, field name = %s, %s" , coll . Fields [ idx ] . Name , field . Name )
}
idx = i
if ! field . IsPrimaryKey {
return fmt . Errorf ( "only primary field can speficy AutoID with true, field name = %s" , field . Name )
2020-11-30 19:38:23 +08:00
}
}
}
2021-06-21 11:42:18 +08:00
return nil
}
2021-10-25 23:42:29 +08:00
func validatePrimaryKey ( coll * schemapb . CollectionSchema ) error {
2020-11-30 19:38:23 +08:00
idx := - 1
for i , field := range coll . Fields {
if field . IsPrimaryKey {
if idx != - 1 {
2021-03-08 12:41:46 +08:00
return fmt . Errorf ( "there are more than one primary key, field name = %s, %s" , coll . Fields [ idx ] . Name , field . Name )
2020-11-30 19:38:23 +08:00
}
2022-03-14 23:20:02 +08:00
// The type of the primary key field can only be int64 and varchar
if field . DataType != schemapb . DataType_Int64 && field . DataType != schemapb . DataType_VarChar {
return errors . New ( "the data type of primary key should be Int64 or VarChar" )
2020-11-30 19:38:23 +08:00
}
2022-03-14 23:20:02 +08:00
// varchar field do not support autoID
// If autoID is required, it is recommended to use int64 field as the primary key
if field . DataType == schemapb . DataType_VarChar {
if field . AutoID {
return fmt . Errorf ( "autoID is not supported when the VarChar field is the primary key" )
}
}
2020-11-30 19:38:23 +08:00
idx = i
}
}
if idx == - 1 {
2021-06-21 11:42:18 +08:00
return errors . New ( "primary key is not specified" )
2020-11-30 19:38:23 +08:00
}
return nil
}
2021-03-08 12:41:46 +08:00
2023-05-18 09:33:24 +08:00
func validateDynamicField ( coll * schemapb . CollectionSchema ) error {
for _ , field := range coll . Fields {
if field . IsDynamic {
return fmt . Errorf ( "cannot explicitly set a field as a dynamic field" )
}
}
return nil
}
2021-11-02 21:43:08 +08:00
// RepeatedKeyValToMap transfer the kv pairs to map.
2021-03-08 12:41:46 +08:00
func RepeatedKeyValToMap ( kvPairs [ ] * commonpb . KeyValuePair ) ( map [ string ] string , error ) {
resMap := make ( map [ string ] string )
for _ , kv := range kvPairs {
_ , ok := resMap [ kv . Key ]
if ok {
return nil , fmt . Errorf ( "duplicated param key: %s" , kv . Key )
}
resMap [ kv . Key ] = kv . Value
}
return resMap , nil
}
2021-12-17 21:54:42 +08:00
// isVector check if dataType belongs to vector type.
2021-03-08 12:41:46 +08:00
func isVector ( dataType schemapb . DataType ) ( bool , error ) {
switch dataType {
2021-03-12 14:22:09 +08:00
case schemapb . DataType_Bool , schemapb . DataType_Int8 ,
schemapb . DataType_Int16 , schemapb . DataType_Int32 ,
schemapb . DataType_Int64 ,
schemapb . DataType_Float , schemapb . DataType_Double :
2021-03-08 12:41:46 +08:00
return false , nil
2021-03-12 14:22:09 +08:00
case schemapb . DataType_FloatVector , schemapb . DataType_BinaryVector :
2021-03-08 12:41:46 +08:00
return true , nil
}
return false , fmt . Errorf ( "invalid data type: %d" , dataType )
}
2021-10-25 23:44:21 +08:00
func validateMetricType ( dataType schemapb . DataType , metricTypeStrRaw string ) error {
2021-03-08 12:41:46 +08:00
metricTypeStr := strings . ToUpper ( metricTypeStrRaw )
switch metricTypeStr {
case "L2" , "IP" :
2021-03-12 14:22:09 +08:00
if dataType == schemapb . DataType_FloatVector {
2021-03-08 12:41:46 +08:00
return nil
}
case "JACCARD" , "HAMMING" , "TANIMOTO" , "SUBSTRUCTURE" , "SUBPERSTURCTURE" :
2021-03-12 14:22:09 +08:00
if dataType == schemapb . DataType_BinaryVector {
2021-03-08 12:41:46 +08:00
return nil
}
}
return fmt . Errorf ( "data_type %s mismatch with metric_type %s" , dataType . String ( ) , metricTypeStrRaw )
}
2021-10-25 23:56:19 +08:00
func validateSchema ( coll * schemapb . CollectionSchema ) error {
2021-03-08 12:41:46 +08:00
autoID := coll . AutoID
primaryIdx := - 1
idMap := make ( map [ int64 ] int ) // fieldId -> idx
nameMap := make ( map [ string ] int ) // name -> idx
for idx , field := range coll . Fields {
// check system field
if field . FieldID < 100 {
// System Fields, not injected yet
2021-11-29 12:25:17 +08:00
return fmt . Errorf ( "fieldID(%d) that is less than 100 is reserved for system fields: %s" , field . FieldID , field . Name )
2021-03-08 12:41:46 +08:00
}
// primary key detector
if field . IsPrimaryKey {
if autoID {
return fmt . Errorf ( "autoId forbids primary key" )
} else if primaryIdx != - 1 {
return fmt . Errorf ( "there are more than one primary key, field name = %s, %s" , coll . Fields [ primaryIdx ] . Name , field . Name )
}
2021-03-12 14:22:09 +08:00
if field . DataType != schemapb . DataType_Int64 {
2021-03-08 12:41:46 +08:00
return fmt . Errorf ( "type of primary key shoule be int64" )
}
primaryIdx = idx
}
// check unique
elemIdx , ok := idMap [ field . FieldID ]
if ok {
return fmt . Errorf ( "duplicate field ids: %d" , coll . Fields [ elemIdx ] . FieldID )
}
idMap [ field . FieldID ] = idx
elemIdx , ok = nameMap [ field . Name ]
if ok {
return fmt . Errorf ( "duplicate field names: %s" , coll . Fields [ elemIdx ] . Name )
}
nameMap [ field . Name ] = idx
isVec , err3 := isVector ( field . DataType )
if err3 != nil {
return err3
}
if isVec {
indexKv , err1 := RepeatedKeyValToMap ( field . IndexParams )
if err1 != nil {
return err1
}
typeKv , err2 := RepeatedKeyValToMap ( field . TypeParams )
if err2 != nil {
return err2
}
2023-05-16 17:41:22 +08:00
dimStr , ok := typeKv [ common . DimKey ]
2021-03-08 12:41:46 +08:00
if ! ok {
return fmt . Errorf ( "dim not found in type_params for vector field %s(%d)" , field . Name , field . FieldID )
}
dim , err := strconv . Atoi ( dimStr )
if err != nil || dim < 0 {
return fmt . Errorf ( "invalid dim; %s" , dimStr )
}
2023-05-16 17:41:22 +08:00
metricTypeStr , ok := indexKv [ common . MetricTypeKey ]
2021-03-08 12:41:46 +08:00
if ok {
2021-10-25 23:44:21 +08:00
err4 := validateMetricType ( field . DataType , metricTypeStr )
2021-03-08 12:41:46 +08:00
if err4 != nil {
return err4
}
} else {
// in C++, default type will be specified
// do nothing
}
} else {
if len ( field . IndexParams ) != 0 {
return fmt . Errorf ( "index params is not empty for scalar field: %s(%d)" , field . Name , field . FieldID )
}
if len ( field . TypeParams ) != 0 {
return fmt . Errorf ( "type params is not empty for scalar field: %s(%d)" , field . Name , field . FieldID )
}
}
}
if ! autoID && primaryIdx == - 1 {
return fmt . Errorf ( "primary key is required for non autoid mode" )
}
return nil
}
2021-12-06 10:03:34 +08:00
2021-12-17 22:19:03 +08:00
// validateMultipleVectorFields check if schema has multiple vector fields.
2021-12-06 10:03:34 +08:00
func validateMultipleVectorFields ( schema * schemapb . CollectionSchema ) error {
vecExist := false
var vecName string
for i := range schema . Fields {
name := schema . Fields [ i ] . Name
dType := schema . Fields [ i ] . DataType
2022-03-25 14:27:25 +08:00
isVec := dType == schemapb . DataType_BinaryVector || dType == schemapb . DataType_FloatVector
2021-12-06 10:03:34 +08:00
if isVec && vecExist && ! enableMultipleVectorFields {
return fmt . Errorf (
"multiple vector fields is not supported, fields name: %s, %s" ,
vecName ,
name ,
)
} else if isVec {
vecExist = true
vecName = name
}
}
return nil
}
2022-03-25 14:27:25 +08:00
// parsePrimaryFieldData2IDs get IDs to fill grpc result, for example insert request, delete request etc.
func parsePrimaryFieldData2IDs ( fieldData * schemapb . FieldData ) ( * schemapb . IDs , error ) {
primaryData := & schemapb . IDs { }
switch fieldData . Field . ( type ) {
case * schemapb . FieldData_Scalars :
scalarField := fieldData . GetScalars ( )
switch scalarField . Data . ( type ) {
case * schemapb . ScalarField_LongData :
primaryData . IdField = & schemapb . IDs_IntId {
IntId : scalarField . GetLongData ( ) ,
}
case * schemapb . ScalarField_StringData :
primaryData . IdField = & schemapb . IDs_StrId {
StrId : scalarField . GetStringData ( ) ,
}
default :
return nil , errors . New ( "currently only support DataType Int64 or VarChar as PrimaryField" )
}
default :
2022-06-20 22:06:12 +08:00
return nil , errors . New ( "currently not support vector field as PrimaryField" )
2022-03-25 14:27:25 +08:00
}
return primaryData , nil
}
// autoGenPrimaryFieldData generate primary data when autoID == true
func autoGenPrimaryFieldData ( fieldSchema * schemapb . FieldSchema , data interface { } ) ( * schemapb . FieldData , error ) {
var fieldData schemapb . FieldData
fieldData . FieldName = fieldSchema . Name
fieldData . Type = fieldSchema . DataType
switch data := data . ( type ) {
case [ ] int64 :
if fieldSchema . DataType != schemapb . DataType_Int64 {
return nil , errors . New ( "the data type of the data and the schema do not match" )
}
fieldData . Field = & schemapb . FieldData_Scalars {
Scalars : & schemapb . ScalarField {
Data : & schemapb . ScalarField_LongData {
LongData : & schemapb . LongArray {
Data : data ,
} ,
} ,
} ,
}
default :
return nil , errors . New ( "currently only support autoID for int64 PrimaryField" )
}
return & fieldData , nil
}
2023-05-18 09:33:24 +08:00
func autoGenDynamicFieldData ( data [ ] [ ] byte ) ( * schemapb . FieldData , error ) {
fieldData := & schemapb . FieldData {
FieldName : common . MetaFieldName ,
Type : schemapb . DataType_JSON ,
Field : & schemapb . FieldData_Scalars {
Scalars : & schemapb . ScalarField {
Data : & schemapb . ScalarField_JsonData {
JsonData : & schemapb . JSONArray {
Data : data ,
} ,
} ,
} ,
} ,
}
return fieldData , nil
}
2022-03-25 14:27:25 +08:00
// fillFieldIDBySchema set fieldID to fieldData according FieldSchemas
func fillFieldIDBySchema ( columns [ ] * schemapb . FieldData , schema * schemapb . CollectionSchema ) error {
if len ( columns ) != len ( schema . GetFields ( ) ) {
return fmt . Errorf ( "len(columns) mismatch the len(fields), len(columns): %d, len(fields): %d" ,
len ( columns ) , len ( schema . GetFields ( ) ) )
}
fieldName2Schema := make ( map [ string ] * schemapb . FieldSchema )
for _ , field := range schema . GetFields ( ) {
fieldName2Schema [ field . Name ] = field
}
for _ , fieldData := range columns {
if fieldSchema , ok := fieldName2Schema [ fieldData . FieldName ] ; ok {
fieldData . FieldId = fieldSchema . FieldID
fieldData . Type = fieldSchema . DataType
} else {
return fmt . Errorf ( "fieldName %v not exist in collection schema" , fieldData . FieldName )
}
}
return nil
}
2022-04-11 19:49:34 +08:00
func ValidateUsername ( username string ) error {
username = strings . TrimSpace ( username )
if username == "" {
return errors . New ( "username should not be empty" )
}
invalidMsg := "Invalid username: " + username + ". "
2022-12-07 18:01:19 +08:00
if len ( username ) > Params . ProxyCfg . MaxUsernameLength . GetAsInt ( ) {
msg := invalidMsg + "The length of username must be less than " + Params . ProxyCfg . MaxUsernameLength . GetValue ( ) + " characters."
2022-04-11 19:49:34 +08:00
return errors . New ( msg )
}
firstChar := username [ 0 ]
if ! isAlpha ( firstChar ) {
msg := invalidMsg + "The first character of username must be a letter."
return errors . New ( msg )
}
usernameSize := len ( username )
for i := 1 ; i < usernameSize ; i ++ {
c := username [ i ]
if c != '_' && ! isAlpha ( c ) && ! isNumber ( c ) {
msg := invalidMsg + "Username should only contain numbers, letters, and underscores."
return errors . New ( msg )
}
}
return nil
}
func ValidatePassword ( password string ) error {
2022-12-07 18:01:19 +08:00
if len ( password ) < Params . ProxyCfg . MinPasswordLength . GetAsInt ( ) || len ( password ) > Params . ProxyCfg . MaxPasswordLength . GetAsInt ( ) {
msg := "The length of password must be great than " + Params . ProxyCfg . MinPasswordLength . GetValue ( ) +
" and less than " + Params . ProxyCfg . MaxPasswordLength . GetValue ( ) + " characters."
2022-04-11 19:49:34 +08:00
return errors . New ( msg )
}
return nil
}
2022-05-11 09:47:53 +08:00
2022-05-24 12:05:59 +08:00
func validateTravelTimestamp ( travelTs , tMax typeutil . Timestamp ) error {
durationSeconds := tsoutil . CalculateDuration ( tMax , travelTs ) / 1000
2022-12-07 18:01:19 +08:00
if durationSeconds > Params . CommonCfg . RetentionDuration . GetAsInt64 ( ) {
2022-10-12 15:17:23 +08:00
durationIn := time . Second * time . Duration ( durationSeconds )
2022-12-07 18:01:19 +08:00
durationSupport := time . Second * time . Duration ( Params . CommonCfg . RetentionDuration . GetAsInt64 ( ) )
2022-10-12 15:17:23 +08:00
return fmt . Errorf ( "only support to travel back to %v so far, but got %v" , durationSupport , durationIn )
2022-05-24 12:05:59 +08:00
}
return nil
}
2022-05-11 09:47:53 +08:00
func ReplaceID2Name ( oldStr string , id int64 , name string ) string {
return strings . ReplaceAll ( oldStr , strconv . FormatInt ( id , 10 ) , name )
}
2022-05-24 12:05:59 +08:00
func parseGuaranteeTs ( ts , tMax typeutil . Timestamp ) typeutil . Timestamp {
switch ts {
case strongTS :
ts = tMax
case boundedTS :
2023-02-15 17:22:34 +08:00
ratio := Params . CommonCfg . GracefulTime . GetAsDuration ( time . Millisecond )
ts = tsoutil . AddPhysicalDurationOnTs ( tMax , - ratio )
2022-05-24 12:05:59 +08:00
}
return ts
}
2022-08-04 11:04:34 +08:00
func validateName ( entity string , nameType string ) error {
entity = strings . TrimSpace ( entity )
if entity == "" {
return fmt . Errorf ( "%s should not be empty" , nameType )
}
invalidMsg := fmt . Sprintf ( "invalid %s: %s. " , nameType , entity )
2022-12-07 18:01:19 +08:00
if len ( entity ) > Params . ProxyCfg . MaxNameLength . GetAsInt ( ) {
msg := invalidMsg + fmt . Sprintf ( "the length of %s must be less than " , nameType ) + Params . ProxyCfg . MaxNameLength . GetValue ( ) + " characters."
2022-08-04 11:04:34 +08:00
return errors . New ( msg )
}
firstChar := entity [ 0 ]
if firstChar != '_' && ! isAlpha ( firstChar ) {
msg := invalidMsg + fmt . Sprintf ( "the first character of %s must be an underscore or letter." , nameType )
return errors . New ( msg )
}
for i := 1 ; i < len ( entity ) ; i ++ {
c := entity [ i ]
if c != '_' && c != '$' && ! isAlpha ( c ) && ! isNumber ( c ) {
msg := invalidMsg + fmt . Sprintf ( "%s can only contain numbers, letters, dollars and underscores." , nameType )
return errors . New ( msg )
}
}
return nil
}
func ValidateRoleName ( entity string ) error {
return validateName ( entity , "role name" )
}
2022-08-15 16:40:48 +08:00
func IsDefaultRole ( roleName string ) bool {
for _ , defaultRole := range util . DefaultRoles {
if defaultRole == roleName {
return true
}
2022-08-04 11:04:34 +08:00
}
2022-08-15 16:40:48 +08:00
return false
}
2022-08-04 11:04:34 +08:00
2022-08-15 16:40:48 +08:00
func ValidateObjectName ( entity string ) error {
if util . IsAnyWord ( entity ) {
return nil
2022-08-04 11:04:34 +08:00
}
2022-08-15 16:40:48 +08:00
return validateName ( entity , "role name" )
2022-08-04 11:04:34 +08:00
}
func ValidateObjectType ( entity string ) error {
return validateName ( entity , "ObjectType" )
}
func ValidatePrincipalName ( entity string ) error {
return validateName ( entity , "PrincipalName" )
}
func ValidatePrincipalType ( entity string ) error {
return validateName ( entity , "PrincipalType" )
}
func ValidatePrivilege ( entity string ) error {
2022-08-15 16:40:48 +08:00
if util . IsAnyWord ( entity ) {
return nil
}
2022-08-04 11:04:34 +08:00
return validateName ( entity , "Privilege" )
}
func GetCurUserFromContext ( ctx context . Context ) ( string , error ) {
md , ok := metadata . FromIncomingContext ( ctx )
if ! ok {
return "" , fmt . Errorf ( "fail to get md from the context" )
}
authorization := md [ strings . ToLower ( util . HeaderAuthorize ) ]
if len ( authorization ) < 1 {
return "" , fmt . Errorf ( "fail to get authorization from the md, authorize:[%s]" , util . HeaderAuthorize )
}
token := authorization [ 0 ]
rawToken , err := crypto . Base64Decode ( token )
if err != nil {
return "" , fmt . Errorf ( "fail to decode the token, token: %s" , token )
}
secrets := strings . SplitN ( rawToken , util . CredentialSeperator , 2 )
if len ( secrets ) < 2 {
return "" , fmt . Errorf ( "fail to get user info from the raw token, raw token: %s" , rawToken )
}
username := secrets [ 0 ]
return username , nil
}
func GetRole ( username string ) ( [ ] string , error ) {
if globalMetaCache == nil {
2023-03-24 15:27:58 +08:00
return [ ] string { } , merr . WrapErrServiceUnavailable ( "internal: Milvus Proxy is not ready yet. please wait" )
2022-08-04 11:04:34 +08:00
}
return globalMetaCache . GetUserRole ( username ) , nil
}
2022-08-19 19:42:50 +08:00
// PasswordVerify verify password
func passwordVerify ( ctx context . Context , username , rawPwd string , globalMetaCache Cache ) bool {
// it represents the cache miss if Sha256Password is empty within credInfo, which shall be updated first connection.
// meanwhile, generating Sha256Password depends on raw password and encrypted password will not cache.
credInfo , err := globalMetaCache . GetCredentialInfo ( ctx , username )
if err != nil {
log . Error ( "found no credential" , zap . String ( "username" , username ) , zap . Error ( err ) )
return false
}
// hit cache
sha256Pwd := crypto . SHA256 ( rawPwd , credInfo . Username )
if credInfo . Sha256Password != "" {
return sha256Pwd == credInfo . Sha256Password
}
// miss cache, verify against encrypted password from etcd
if err := bcrypt . CompareHashAndPassword ( [ ] byte ( credInfo . EncryptedPassword ) , [ ] byte ( rawPwd ) ) ; err != nil {
log . Error ( "Verify password failed" , zap . Error ( err ) )
return false
}
// update cache after miss cache
credInfo . Sha256Password = sha256Pwd
log . Debug ( "get credential miss cache, update cache with" , zap . Any ( "credential" , credInfo ) )
globalMetaCache . UpdateCredential ( credInfo )
return true
}
2022-10-08 15:38:58 +08:00
// Support wildcard in output fields:
2022-10-12 18:37:23 +08:00
//
2023-05-17 12:41:22 +08:00
// "*" - all fields
2022-10-12 18:37:23 +08:00
//
2022-10-08 15:38:58 +08:00
// For example, A and B are scalar fields, C and D are vector fields, duplicated fields will automatically be removed.
2022-10-12 18:37:23 +08:00
//
2023-05-17 12:41:22 +08:00
// output_fields=["*"] ==> [A,B,C,D]
// output_fields=["*",A] ==> [A,B,C,D]
// output_fields=["*",C] ==> [A,B,C,D]
2022-10-08 15:38:58 +08:00
func translateOutputFields ( outputFields [ ] string , schema * schemapb . CollectionSchema , addPrimary bool ) ( [ ] string , error ) {
var primaryFieldName string
2023-05-19 09:41:25 +08:00
allFieldNameMap := make ( map [ string ] bool )
2022-10-08 15:38:58 +08:00
resultFieldNameMap := make ( map [ string ] bool )
resultFieldNames := make ( [ ] string , 0 )
for _ , field := range schema . Fields {
if field . IsPrimaryKey {
primaryFieldName = field . Name
}
2023-05-19 09:41:25 +08:00
allFieldNameMap [ field . Name ] = true
2022-10-08 15:38:58 +08:00
}
for _ , outputFieldName := range outputFields {
outputFieldName = strings . TrimSpace ( outputFieldName )
if outputFieldName == "*" {
2023-05-19 09:41:25 +08:00
for fieldName := range allFieldNameMap {
2022-10-08 15:38:58 +08:00
resultFieldNameMap [ fieldName ] = true
}
} else {
2023-05-19 09:41:25 +08:00
if _ , ok := allFieldNameMap [ outputFieldName ] ; ok {
resultFieldNameMap [ outputFieldName ] = true
} else {
if schema . EnableDynamicField {
resultFieldNameMap [ common . MetaFieldName ] = true
} else {
return nil , fmt . Errorf ( "field %s not exist" , outputFieldName )
}
}
2022-10-08 15:38:58 +08:00
}
}
if addPrimary {
resultFieldNameMap [ primaryFieldName ] = true
}
for fieldName := range resultFieldNameMap {
resultFieldNames = append ( resultFieldNames , fieldName )
}
return resultFieldNames , nil
}
2022-10-16 21:05:25 +08:00
func validateIndexName ( indexName string ) error {
indexName = strings . TrimSpace ( indexName )
if indexName == "" {
return nil
}
invalidMsg := "Invalid index name: " + indexName + ". "
2022-12-07 18:01:19 +08:00
if len ( indexName ) > Params . ProxyCfg . MaxNameLength . GetAsInt ( ) {
msg := invalidMsg + "The length of a index name must be less than " + Params . ProxyCfg . MaxNameLength . GetValue ( ) + " characters."
2022-10-16 21:05:25 +08:00
return errors . New ( msg )
}
firstChar := indexName [ 0 ]
if firstChar != '_' && ! isAlpha ( firstChar ) {
msg := invalidMsg + "The first character of a index name must be an underscore or letter."
return errors . New ( msg )
}
indexNameSize := len ( indexName )
for i := 1 ; i < indexNameSize ; i ++ {
c := indexName [ i ]
if c != '_' && ! isAlpha ( c ) && ! isNumber ( c ) {
msg := invalidMsg + "Index name cannot only contain numbers, letters, and underscores."
return errors . New ( msg )
}
}
return nil
}
2022-10-21 14:41:28 +08:00
2022-10-27 13:05:31 +08:00
func isCollectionLoaded ( ctx context . Context , qc types . QueryCoord , collID int64 ) ( bool , error ) {
2022-10-21 14:41:28 +08:00
// get all loading collections
resp , err := qc . ShowCollections ( ctx , & querypb . ShowCollectionsRequest {
CollectionIDs : nil ,
} )
if err != nil {
return false , err
}
if resp . Status . ErrorCode != commonpb . ErrorCode_Success {
return false , errors . New ( resp . Status . Reason )
}
for _ , loadedCollID := range resp . GetCollectionIDs ( ) {
2022-10-27 13:05:31 +08:00
if collID == loadedCollID {
return true , nil
2022-10-21 14:41:28 +08:00
}
}
2022-10-27 13:05:31 +08:00
return false , nil
2022-10-21 14:41:28 +08:00
}
2022-10-27 13:05:31 +08:00
func isPartitionLoaded ( ctx context . Context , qc types . QueryCoord , collID int64 , partIDs [ ] int64 ) ( bool , error ) {
2022-10-21 14:41:28 +08:00
// get all loading collections
resp , err := qc . ShowPartitions ( ctx , & querypb . ShowPartitionsRequest {
2022-10-27 13:05:31 +08:00
CollectionID : collID ,
2022-10-21 14:41:28 +08:00
PartitionIDs : nil ,
} )
if err != nil {
return false , err
}
if resp . Status . ErrorCode != commonpb . ErrorCode_Success {
return false , errors . New ( resp . Status . Reason )
}
for _ , loadedPartID := range resp . GetPartitionIDs ( ) {
for _ , partID := range partIDs {
if partID == loadedPartID {
2022-10-27 13:05:31 +08:00
return true , nil
2022-10-21 14:41:28 +08:00
}
}
}
2022-10-27 13:05:31 +08:00
return false , nil
2022-10-21 14:41:28 +08:00
}
2022-12-08 18:37:19 +08:00
2023-05-15 16:15:21 +08:00
func fillFieldsDataBySchema ( schema * schemapb . CollectionSchema , insertMsg * msgstream . InsertMsg ) error {
2022-12-08 18:37:19 +08:00
neededFieldsNum := 0
2023-05-15 16:15:21 +08:00
isPrimaryKeyNum := 0
var dataNameSet = typeutil . NewSet [ string ] ( )
for _ , data := range insertMsg . FieldsData {
dataNameSet . Insert ( data . GetFieldName ( ) )
}
for _ , fieldSchema := range schema . Fields {
if fieldSchema . AutoID && ! fieldSchema . IsPrimaryKey {
log . Error ( "not primary key field, but set autoID true" , zap . String ( "fieldSchemaName" , fieldSchema . GetName ( ) ) )
return merr . WrapErrParameterInvalid ( "only primary key field can set autoID true" , "" )
}
if fieldSchema . GetDefaultValue ( ) != nil && fieldSchema . IsPrimaryKey {
return merr . WrapErrParameterInvalid ( "no default data" , "" , "pk field schema can not set default value" )
}
if ! fieldSchema . AutoID {
2022-12-08 18:37:19 +08:00
neededFieldsNum ++
}
2023-05-15 16:15:21 +08:00
// if has no field pass in, consider use default value
// so complete it with field schema
if _ , ok := dataNameSet [ fieldSchema . GetName ( ) ] ; ! ok {
// primary key can not use default value
if fieldSchema . IsPrimaryKey {
isPrimaryKeyNum ++
continue
}
dataToAppend := & schemapb . FieldData {
Type : fieldSchema . GetDataType ( ) ,
FieldName : fieldSchema . GetName ( ) ,
}
insertMsg . FieldsData = append ( insertMsg . FieldsData , dataToAppend )
}
}
if isPrimaryKeyNum > 1 {
log . Error ( "the number of passed primary key fields is more than 1" ,
zap . Int64 ( "primaryKeyNum" , int64 ( isPrimaryKeyNum ) ) ,
zap . String ( "CollectionSchemaName" , schema . GetName ( ) ) )
return merr . WrapErrParameterInvalid ( "0 or 1" , fmt . Sprint ( isPrimaryKeyNum ) , "the number of passed primary key fields is more than 1" )
2022-12-08 18:37:19 +08:00
}
2023-05-15 16:15:21 +08:00
if len ( insertMsg . FieldsData ) != neededFieldsNum {
log . Error ( "the length of passed fields is not equal to needed" ,
zap . Int ( "expectFieldNumber" , neededFieldsNum ) ,
zap . Int ( "passFieldNumber" , len ( insertMsg . FieldsData ) ) ,
zap . String ( "CollectionSchemaName" , schema . GetName ( ) ) )
return merr . WrapErrParameterInvalid ( neededFieldsNum , len ( insertMsg . FieldsData ) , "the length of passed fields is equal to needed" )
2022-12-08 18:37:19 +08:00
}
return nil
}
2023-01-17 17:53:42 +08:00
func checkPrimaryFieldData ( schema * schemapb . CollectionSchema , result * milvuspb . MutationResult , insertMsg * msgstream . InsertMsg , inInsert bool ) ( * schemapb . IDs , error ) {
2022-12-08 18:37:19 +08:00
rowNums := uint32 ( insertMsg . NRows ( ) )
// TODO(dragondriver): in fact, NumRows is not trustable, we should check all input fields
if insertMsg . NRows ( ) <= 0 {
2023-03-24 15:27:58 +08:00
return nil , merr . WrapErrParameterInvalid ( "invalid num_rows" , fmt . Sprint ( rowNums ) , "num_rows should be greater than 0" )
2022-12-08 18:37:19 +08:00
}
2023-05-15 16:15:21 +08:00
if err := fillFieldsDataBySchema ( schema , insertMsg ) ; err != nil {
2022-12-08 18:37:19 +08:00
return nil , err
}
primaryFieldSchema , err := typeutil . GetPrimaryFieldSchema ( schema )
if err != nil {
log . Error ( "get primary field schema failed" , zap . String ( "collectionName" , insertMsg . CollectionName ) , zap . Any ( "schema" , schema ) , zap . Error ( err ) )
return nil , err
}
// get primaryFieldData whether autoID is true or not
var primaryFieldData * schemapb . FieldData
2023-01-17 17:53:42 +08:00
if inInsert {
// when checkPrimaryFieldData in insert
if ! primaryFieldSchema . AutoID {
primaryFieldData , err = typeutil . GetPrimaryFieldData ( insertMsg . GetFieldsData ( ) , primaryFieldSchema )
if err != nil {
log . Error ( "get primary field data failed" , zap . String ( "collectionName" , insertMsg . CollectionName ) , zap . Error ( err ) )
return nil , err
}
} else {
// check primary key data not exist
if typeutil . IsPrimaryFieldDataExist ( insertMsg . GetFieldsData ( ) , primaryFieldSchema ) {
return nil , fmt . Errorf ( "can not assign primary field data when auto id enabled %v" , primaryFieldSchema . Name )
}
// if autoID == true, currently only support autoID for int64 PrimaryField
primaryFieldData , err = autoGenPrimaryFieldData ( primaryFieldSchema , insertMsg . GetRowIDs ( ) )
if err != nil {
log . Error ( "generate primary field data failed when autoID == true" , zap . String ( "collectionName" , insertMsg . CollectionName ) , zap . Error ( err ) )
return nil , err
}
// if autoID == true, set the primary field data
// insertMsg.fieldsData need append primaryFieldData
insertMsg . FieldsData = append ( insertMsg . FieldsData , primaryFieldData )
2022-12-08 18:37:19 +08:00
}
} else {
2023-01-17 17:53:42 +08:00
// when checkPrimaryFieldData in upsert
if primaryFieldSchema . AutoID {
// upsert has not supported when autoID == true
log . Info ( "can not upsert when auto id enabled" ,
zap . String ( "primaryFieldSchemaName" , primaryFieldSchema . Name ) )
result . Status . ErrorCode = commonpb . ErrorCode_UpsertAutoIDTrue
return nil , fmt . Errorf ( "upsert can not assign primary field data when auto id enabled %v" , primaryFieldSchema . Name )
2022-12-08 18:37:19 +08:00
}
2023-01-17 17:53:42 +08:00
primaryFieldData , err = typeutil . GetPrimaryFieldData ( insertMsg . GetFieldsData ( ) , primaryFieldSchema )
2022-12-08 18:37:19 +08:00
if err != nil {
2023-01-17 17:53:42 +08:00
log . Error ( "get primary field data failed when upsert" , zap . String ( "collectionName" , insertMsg . CollectionName ) , zap . Error ( err ) )
2022-12-08 18:37:19 +08:00
return nil , err
}
2023-01-04 17:21:36 +08:00
}
// parse primaryFieldData to result.IDs, and as returned primary keys
ids , err := parsePrimaryFieldData2IDs ( primaryFieldData )
if err != nil {
log . Error ( "parse primary field data to IDs failed" , zap . String ( "collectionName" , insertMsg . CollectionName ) , zap . Error ( err ) )
return nil , err
}
return ids , nil
}
2023-04-12 15:06:28 +08:00
func getCollectionProgress (
ctx context . Context ,
queryCoord types . QueryCoord ,
msgBase * commonpb . MsgBase ,
collectionID int64 ,
) ( loadProgress int64 , refreshProgress int64 , err error ) {
2022-12-16 14:39:24 +08:00
resp , err := queryCoord . ShowCollections ( ctx , & querypb . ShowCollectionsRequest {
Base : commonpbutil . UpdateMsgBase (
msgBase ,
commonpbutil . WithMsgType ( commonpb . MsgType_DescribeCollection ) ,
) ,
CollectionIDs : [ ] int64 { collectionID } ,
} )
if err != nil {
log . Warn ( "fail to show collections" , zap . Int64 ( "collection_id" , collectionID ) , zap . Error ( err ) )
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
2023-01-10 20:35:39 +08:00
if resp . Status . ErrorCode == commonpb . ErrorCode_InsufficientMemoryToLoad {
2023-04-12 15:06:28 +08:00
err = ErrInsufficientMemory
2023-01-10 20:35:39 +08:00
log . Warn ( "detected insufficientMemoryError when getCollectionProgress" , zap . Int64 ( "collection_id" , collectionID ) , zap . String ( "reason" , resp . GetStatus ( ) . GetReason ( ) ) )
2023-04-12 15:06:28 +08:00
return
2023-01-10 20:35:39 +08:00
}
2022-12-16 14:39:24 +08:00
if resp . Status . ErrorCode != commonpb . ErrorCode_Success {
2023-04-12 15:06:28 +08:00
err = merr . Error ( resp . GetStatus ( ) )
2022-12-16 14:39:24 +08:00
log . Warn ( "fail to show collections" , zap . Int64 ( "collection_id" , collectionID ) ,
zap . String ( "reason" , resp . Status . Reason ) )
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
if len ( resp . InMemoryPercentages ) == 0 {
errMsg := "fail to show collections from the querycoord, no data"
2023-04-12 15:06:28 +08:00
err = errors . New ( errMsg )
2022-12-16 14:39:24 +08:00
log . Warn ( errMsg , zap . Int64 ( "collection_id" , collectionID ) )
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
2023-04-12 15:06:28 +08:00
loadProgress = resp . GetInMemoryPercentages ( ) [ 0 ]
if len ( resp . GetRefreshProgress ( ) ) > 0 { // Compatibility for new Proxy with old QueryCoord
refreshProgress = resp . GetRefreshProgress ( ) [ 0 ]
}
return
2022-12-16 14:39:24 +08:00
}
2023-04-12 15:06:28 +08:00
func getPartitionProgress (
ctx context . Context ,
queryCoord types . QueryCoord ,
msgBase * commonpb . MsgBase ,
partitionNames [ ] string ,
collectionName string ,
collectionID int64 ,
) ( loadProgress int64 , refreshProgress int64 , err error ) {
2022-12-16 14:39:24 +08:00
IDs2Names := make ( map [ int64 ] string )
partitionIDs := make ( [ ] int64 , 0 )
for _ , partitionName := range partitionNames {
2023-04-12 15:06:28 +08:00
var partitionID int64
partitionID , err = globalMetaCache . GetPartitionID ( ctx , collectionName , partitionName )
2022-12-16 14:39:24 +08:00
if err != nil {
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
IDs2Names [ partitionID ] = partitionName
partitionIDs = append ( partitionIDs , partitionID )
}
resp , err := queryCoord . ShowPartitions ( ctx , & querypb . ShowPartitionsRequest {
Base : commonpbutil . UpdateMsgBase (
msgBase ,
commonpbutil . WithMsgType ( commonpb . MsgType_ShowPartitions ) ,
) ,
CollectionID : collectionID ,
PartitionIDs : partitionIDs ,
} )
if err != nil {
log . Warn ( "fail to show partitions" , zap . Int64 ( "collection_id" , collectionID ) ,
zap . String ( "collection_name" , collectionName ) ,
zap . Strings ( "partition_names" , partitionNames ) ,
zap . Error ( err ) )
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
2023-01-10 20:35:39 +08:00
if resp . GetStatus ( ) . GetErrorCode ( ) == commonpb . ErrorCode_InsufficientMemoryToLoad {
2023-04-12 15:06:28 +08:00
err = ErrInsufficientMemory
log . Warn ( "detected insufficientMemoryError when getPartitionProgress" ,
zap . Int64 ( "collection_id" , collectionID ) ,
zap . String ( "collection_name" , collectionName ) ,
zap . Strings ( "partition_names" , partitionNames ) ,
zap . String ( "reason" , resp . GetStatus ( ) . GetReason ( ) ) ,
)
return
2023-01-10 20:35:39 +08:00
}
2022-12-16 14:39:24 +08:00
if len ( resp . InMemoryPercentages ) != len ( partitionIDs ) {
errMsg := "fail to show partitions from the querycoord, invalid data num"
2023-04-12 15:06:28 +08:00
err = errors . New ( errMsg )
2022-12-16 14:39:24 +08:00
log . Warn ( errMsg , zap . Int64 ( "collection_id" , collectionID ) ,
zap . String ( "collection_name" , collectionName ) ,
zap . Strings ( "partition_names" , partitionNames ) )
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
for _ , p := range resp . InMemoryPercentages {
2023-04-12 15:06:28 +08:00
loadProgress += p
}
loadProgress /= int64 ( len ( partitionIDs ) )
if len ( resp . GetRefreshProgress ( ) ) > 0 { // Compatibility for new Proxy with old QueryCoord
refreshProgress = resp . GetRefreshProgress ( ) [ 0 ]
2022-12-16 14:39:24 +08:00
}
2023-04-12 15:06:28 +08:00
return
2022-12-16 14:39:24 +08:00
}
2023-05-15 16:15:21 +08:00
func memsetLoop [ T any ] ( v T , numRows int ) [ ] T {
ret := make ( [ ] T , 0 , numRows )
for i := 0 ; i < numRows ; i ++ {
ret = append ( ret , v )
}
return ret
}