github.com/milvus-io/milvus-sdk-go/v2@v2.4.1/client/grpc_client.go (about)

     1  package client
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"strconv"
     8  	"time"
     9  
    10  	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
    11  	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
    12  	"github.com/milvus-io/milvus-sdk-go/v2/common"
    13  	"github.com/milvus-io/milvus-sdk-go/v2/merr"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/codes"
    16  	"google.golang.org/grpc/status"
    17  )
    18  
    19  // Check if GrpcClient implement Client.
    20  var _ Client = &GrpcClient{}
    21  
    22  // GrpcClient  uses default grpc Service definition to connect with Milvus2.0
    23  type GrpcClient struct {
    24  	Conn    *grpc.ClientConn             // grpc connection instance
    25  	Service milvuspb.MilvusServiceClient // Service client stub
    26  
    27  	config *Config // No thread safety
    28  }
    29  
    30  // connect connect to Service
    31  func (c *GrpcClient) connect(ctx context.Context, addr string, opts ...grpc.DialOption) error {
    32  	if addr == "" {
    33  		return fmt.Errorf("address is empty")
    34  	}
    35  	conn, err := grpc.DialContext(ctx, addr, opts...)
    36  	if err != nil {
    37  		return err
    38  	}
    39  
    40  	c.Conn = conn
    41  	c.Service = milvuspb.NewMilvusServiceClient(c.Conn)
    42  
    43  	if !c.config.DisableConn {
    44  		err = c.connectInternal(ctx)
    45  		if err != nil {
    46  			return err
    47  		}
    48  	}
    49  
    50  	return nil
    51  }
    52  
    53  func (c *GrpcClient) connectInternal(ctx context.Context) error {
    54  	hostName, err := os.Hostname()
    55  	if err != nil {
    56  		return err
    57  	}
    58  
    59  	req := &milvuspb.ConnectRequest{
    60  		ClientInfo: &commonpb.ClientInfo{
    61  			SdkType:    "Golang",
    62  			SdkVersion: common.SDKVersion,
    63  			LocalTime:  time.Now().String(),
    64  			User:       c.config.Username,
    65  			Host:       hostName,
    66  		},
    67  	}
    68  
    69  	resp, err := c.Service.Connect(ctx, req)
    70  	if err != nil {
    71  		status, ok := status.FromError(err)
    72  		if ok {
    73  			if status.Code() == codes.Unimplemented {
    74  				// disable unsupported feature
    75  				c.config.addFlags(
    76  					disableDatabase |
    77  						disableJSON |
    78  						disableParitionKey |
    79  						disableDynamicSchema)
    80  				return nil
    81  			}
    82  		}
    83  		return err
    84  	}
    85  
    86  	if !merr.Ok(resp.GetStatus()) {
    87  		return fmt.Errorf("connect fail, %s", resp.GetStatus().GetReason())
    88  	}
    89  
    90  	c.config.Identifier = strconv.FormatInt(resp.GetIdentifier(), 10)
    91  	c.config.ServerVersion = resp.GetServerInfo().GetBuildTags()
    92  	return nil
    93  }
    94  
    95  // Close close the connection
    96  func (c *GrpcClient) Close() error {
    97  	if c.Conn != nil {
    98  		err := c.Conn.Close()
    99  		c.Conn = nil
   100  		return err
   101  	}
   102  	return nil
   103  }