collection.go 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  1. // Licensed to the LF AI & Data foundation under one
  2. // or more contributor license agreements. See the NOTICE file
  3. // distributed with this work for additional information
  4. // regarding copyright ownership. The ASF licenses this file
  5. // to you under the Apache License, Version 2.0 (the
  6. // "License"); you may not use this file except in compliance
  7. // with the License. You may obtain a copy of the License at
  8. //
  9. // http://www.apache.org/licenses/LICENSE-2.0
  10. //
  11. // Unless required by applicable law or agreed to in writing, software
  12. // distributed under the License is distributed on an "AS IS" BASIS,
  13. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. // See the License for the specific language governing permissions and
  15. // limitations under the License.
  16. package client
  17. import (
  18. "context"
  19. "github.com/cockroachdb/errors"
  20. "google.golang.org/grpc"
  21. "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
  22. "github.com/milvus-io/milvus/client/v2/entity"
  23. "github.com/milvus-io/milvus/pkg/util/merr"
  24. )
  25. // CreateCollection is the API for create a collection in Milvus.
  26. func (c *Client) CreateCollection(ctx context.Context, option CreateCollectionOption, callOptions ...grpc.CallOption) error {
  27. req := option.Request()
  28. err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
  29. resp, err := milvusService.CreateCollection(ctx, req, callOptions...)
  30. return merr.CheckRPCCall(resp, err)
  31. })
  32. if err != nil {
  33. return err
  34. }
  35. indexes := option.Indexes()
  36. for _, indexOption := range indexes {
  37. task, err := c.CreateIndex(ctx, indexOption, callOptions...)
  38. if err != nil {
  39. return err
  40. }
  41. err = task.Await(ctx)
  42. if err != nil {
  43. return nil
  44. }
  45. }
  46. if option.IsFast() {
  47. task, err := c.LoadCollection(ctx, NewLoadCollectionOption(req.GetCollectionName()))
  48. if err != nil {
  49. return err
  50. }
  51. return task.Await(ctx)
  52. }
  53. return nil
  54. }
  55. func (c *Client) ListCollections(ctx context.Context, option ListCollectionOption, callOptions ...grpc.CallOption) (collectionNames []string, err error) {
  56. req := option.Request()
  57. err = c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
  58. resp, err := milvusService.ShowCollections(ctx, req, callOptions...)
  59. err = merr.CheckRPCCall(resp, err)
  60. if err != nil {
  61. return err
  62. }
  63. collectionNames = resp.GetCollectionNames()
  64. return nil
  65. })
  66. return collectionNames, err
  67. }
  68. func (c *Client) DescribeCollection(ctx context.Context, option DescribeCollectionOption, callOptions ...grpc.CallOption) (collection *entity.Collection, err error) {
  69. req := option.Request()
  70. err = c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
  71. resp, err := milvusService.DescribeCollection(ctx, req, callOptions...)
  72. err = merr.CheckRPCCall(resp, err)
  73. if err != nil {
  74. return err
  75. }
  76. collection = &entity.Collection{
  77. ID: resp.GetCollectionID(),
  78. Schema: entity.NewSchema().ReadProto(resp.GetSchema()),
  79. PhysicalChannels: resp.GetPhysicalChannelNames(),
  80. VirtualChannels: resp.GetVirtualChannelNames(),
  81. ConsistencyLevel: entity.ConsistencyLevel(resp.ConsistencyLevel),
  82. ShardNum: resp.GetShardsNum(),
  83. Properties: entity.KvPairsMap(resp.GetProperties()),
  84. }
  85. collection.Name = collection.Schema.CollectionName
  86. return nil
  87. })
  88. return collection, err
  89. }
  90. func (c *Client) HasCollection(ctx context.Context, option HasCollectionOption, callOptions ...grpc.CallOption) (has bool, err error) {
  91. req := option.Request()
  92. err = c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
  93. resp, err := milvusService.DescribeCollection(ctx, req, callOptions...)
  94. err = merr.CheckRPCCall(resp, err)
  95. if err != nil {
  96. // ErrCollectionNotFound for collection not exist
  97. if errors.Is(err, merr.ErrCollectionNotFound) {
  98. return nil
  99. }
  100. return err
  101. }
  102. has = true
  103. return nil
  104. })
  105. return has, err
  106. }
  107. func (c *Client) DropCollection(ctx context.Context, option DropCollectionOption, callOptions ...grpc.CallOption) error {
  108. req := option.Request()
  109. err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
  110. resp, err := milvusService.DropCollection(ctx, req, callOptions...)
  111. return merr.CheckRPCCall(resp, err)
  112. })
  113. return err
  114. }
  115. func (c *Client) RenameCollection(ctx context.Context, option RenameCollectionOption, callOptions ...grpc.CallOption) error {
  116. req := option.Request()
  117. return c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
  118. resp, err := milvusService.RenameCollection(ctx, req, callOptions...)
  119. return merr.CheckRPCCall(resp, err)
  120. })
  121. }