github.com/ydb-platform/ydb-go-sdk/v3@v3.57.0/internal/discovery/discovery.go (about)

     1  package discovery
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  	"net"
     7  	"strconv"
     8  
     9  	"github.com/ydb-platform/ydb-go-genproto/Ydb_Discovery_V1"
    10  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb"
    11  	"github.com/ydb-platform/ydb-go-genproto/protos/Ydb_Discovery"
    12  	"google.golang.org/grpc"
    13  
    14  	"github.com/ydb-platform/ydb-go-sdk/v3/discovery"
    15  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/discovery/config"
    16  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
    17  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stack"
    18  	"github.com/ydb-platform/ydb-go-sdk/v3/internal/xerrors"
    19  	"github.com/ydb-platform/ydb-go-sdk/v3/trace"
    20  )
    21  
    22  func New(ctx context.Context, cc grpc.ClientConnInterface, config *config.Config) (*Client, error) {
    23  	return &Client{
    24  		config: config,
    25  		cc:     cc,
    26  		client: Ydb_Discovery_V1.NewDiscoveryServiceClient(cc),
    27  	}, nil
    28  }
    29  
    30  var _ discovery.Client = &Client{}
    31  
    32  type Client struct {
    33  	config *config.Config
    34  	cc     grpc.ClientConnInterface
    35  	client Ydb_Discovery_V1.DiscoveryServiceClient
    36  }
    37  
    38  // Discover cluster endpoints
    39  func (c *Client) Discover(ctx context.Context) (endpoints []endpoint.Endpoint, err error) {
    40  	var (
    41  		onDone = trace.DiscoveryOnDiscover(
    42  			c.config.Trace(), &ctx,
    43  			stack.FunctionID(""),
    44  			c.config.Endpoint(), c.config.Database(),
    45  		)
    46  		request = Ydb_Discovery.ListEndpointsRequest{
    47  			Database: c.config.Database(),
    48  		}
    49  		response *Ydb_Discovery.ListEndpointsResponse
    50  		result   Ydb_Discovery.ListEndpointsResult
    51  		location string
    52  	)
    53  	defer func() {
    54  		nodes := make([]trace.EndpointInfo, 0, len(endpoints))
    55  		for _, e := range endpoints {
    56  			nodes = append(nodes, e.Copy())
    57  		}
    58  		onDone(location, nodes, err)
    59  	}()
    60  
    61  	ctx, err = c.config.Meta().Context(ctx)
    62  	if err != nil {
    63  		return nil, xerrors.WithStackTrace(err)
    64  	}
    65  
    66  	response, err = c.client.ListEndpoints(ctx, &request)
    67  	if err != nil {
    68  		return nil, xerrors.WithStackTrace(err)
    69  	}
    70  
    71  	if response.GetOperation().GetStatus() != Ydb.StatusIds_SUCCESS {
    72  		return nil, xerrors.WithStackTrace(
    73  			xerrors.FromOperation(response.GetOperation()),
    74  		)
    75  	}
    76  
    77  	err = response.GetOperation().GetResult().UnmarshalTo(&result)
    78  	if err != nil {
    79  		return nil, xerrors.WithStackTrace(err)
    80  	}
    81  
    82  	location = result.GetSelfLocation()
    83  	endpoints = make([]endpoint.Endpoint, 0, len(result.GetEndpoints()))
    84  	for _, e := range result.GetEndpoints() {
    85  		if e.GetSsl() == c.config.Secure() {
    86  			endpoints = append(endpoints, endpoint.New(
    87  				net.JoinHostPort(e.GetAddress(), strconv.Itoa(int(e.GetPort()))),
    88  				endpoint.WithLocation(e.GetLocation()),
    89  				endpoint.WithID(e.GetNodeId()),
    90  				endpoint.WithLoadFactor(e.GetLoadFactor()),
    91  				endpoint.WithLocalDC(e.GetLocation() == location),
    92  				endpoint.WithServices(e.GetService()),
    93  			))
    94  		}
    95  	}
    96  
    97  	return endpoints, nil
    98  }
    99  
   100  func (c *Client) WhoAmI(ctx context.Context) (whoAmI *discovery.WhoAmI, err error) {
   101  	var (
   102  		onDone             = trace.DiscoveryOnWhoAmI(c.config.Trace(), &ctx, stack.FunctionID(""))
   103  		request            = Ydb_Discovery.WhoAmIRequest{}
   104  		response           *Ydb_Discovery.WhoAmIResponse
   105  		whoAmIResultResult Ydb_Discovery.WhoAmIResult
   106  	)
   107  	defer func() {
   108  		if err != nil {
   109  			onDone("", nil, err)
   110  		} else {
   111  			onDone(whoAmI.User, whoAmI.Groups, err)
   112  		}
   113  	}()
   114  
   115  	ctx, err = c.config.Meta().Context(ctx)
   116  	if err != nil {
   117  		return nil, xerrors.WithStackTrace(err)
   118  	}
   119  
   120  	response, err = c.client.WhoAmI(ctx, &request)
   121  	if err != nil {
   122  		return nil, xerrors.WithStackTrace(err)
   123  	}
   124  
   125  	if response.GetOperation().GetStatus() != Ydb.StatusIds_SUCCESS {
   126  		return nil, xerrors.WithStackTrace(
   127  			xerrors.FromOperation(
   128  				response.GetOperation(),
   129  			),
   130  		)
   131  	}
   132  
   133  	result := response.GetOperation().GetResult()
   134  	if result == nil {
   135  		return &discovery.WhoAmI{}, nil
   136  	}
   137  
   138  	err = response.GetOperation().GetResult().UnmarshalTo(&whoAmIResultResult)
   139  	if err != nil {
   140  		return nil, xerrors.WithStackTrace(err)
   141  	}
   142  
   143  	return &discovery.WhoAmI{
   144  		User:   whoAmIResultResult.GetUser(),
   145  		Groups: whoAmIResultResult.GetGroups(),
   146  	}, nil
   147  }
   148  
   149  func (c *Client) Close(context.Context) error {
   150  	if cc, has := c.cc.(io.Closer); has {
   151  		return cc.Close()
   152  	}
   153  
   154  	return nil
   155  }