fix: add more requests to the database interceptor (#30452)

issue: #30368

Signed-off-by: SimFG <bang.fu@zilliz.com>
pull/30479/head
SimFG 2024-02-02 16:49:06 +08:00 committed by GitHub
parent 0a78b38bb8
commit 02de84239d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 1 deletions

View File

@ -128,6 +128,11 @@ func fillDatabase(ctx context.Context, req interface{}) (context.Context, interf
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
return ctx, r
case *milvuspb.AlterIndexRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
return ctx, r
case *milvuspb.GetIndexBuildProgressRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
@ -153,6 +158,11 @@ func fillDatabase(ctx context.Context, req interface{}) (context.Context, interf
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
return ctx, r
case *milvuspb.HybridSearchRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
return ctx, r
case *milvuspb.FlushRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
@ -178,6 +188,15 @@ func fillDatabase(ctx context.Context, req interface{}) (context.Context, interf
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
return ctx, r
case *milvuspb.DescribeAliasRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
return ctx, r
case *milvuspb.ListAliasesRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
}
case *milvuspb.ImportRequest:
if r.DbName == "" {
r.DbName = GetCurDBNameFromContextOrDefault(ctx)
@ -247,6 +266,6 @@ func fillDatabase(ctx context.Context, req interface{}) (context.Context, interf
}
return ctx, r
default:
return ctx, req
}
return ctx, req
}

View File

@ -68,17 +68,21 @@ func TestDatabaseInterceptor(t *testing.T) {
&milvuspb.CreateIndexRequest{},
&milvuspb.DescribeIndexRequest{},
&milvuspb.DropIndexRequest{},
&milvuspb.AlterIndexRequest{},
&milvuspb.GetIndexBuildProgressRequest{},
&milvuspb.GetIndexStateRequest{},
&milvuspb.InsertRequest{},
&milvuspb.DeleteRequest{},
&milvuspb.SearchRequest{},
&milvuspb.HybridSearchRequest{},
&milvuspb.FlushRequest{},
&milvuspb.GetFlushStateRequest{},
&milvuspb.QueryRequest{},
&milvuspb.CreateAliasRequest{},
&milvuspb.DropAliasRequest{},
&milvuspb.AlterAliasRequest{},
&milvuspb.ListAliasesRequest{},
&milvuspb.DescribeAliasRequest{},
&milvuspb.GetPersistentSegmentInfoRequest{},
&milvuspb.GetQuerySegmentInfoRequest{},
&milvuspb.LoadBalanceRequest{},