Add GetKeyRange method to proto & EntityMgr (#567)

* Add GetKeyRange to proto * Add GetKeyRange to KeyRangeMgr * Added GetKeyRange method to all of KeyRangeMgr's * Changed id -> ids in GetKeyRange * Implement GetKeyRange in KeyRangeServiceServer's
pg-sharding · Mar 27, 2024 · 35bca1b · 35bca1b
1 parent c6d975e
commit 35bca1b
Show file tree

Hide file tree

Showing 9 changed files with 258 additions and 79 deletions.
diff --git a/coordinator/provider/coordinator.go b/coordinator/provider/coordinator.go
@@ -468,6 +468,16 @@ func (qc *qdbCoordinator) AddKeyRange(ctx context.Context, keyRange *kr.KeyRange
 	})
 }
 
+// GetKeyRange gets key range by id
+// TODO unit tests
+func (qc *qdbCoordinator) GetKeyRange(ctx context.Context, krId string) (*kr.KeyRange, error) {
+	krDb, err := qc.db.GetKeyRange(ctx, krId)
+	if err != nil {
+		return nil, err
+	}
+	return kr.KeyRangeFromDB(krDb), nil
+}
+
 // TODO : unit tests
 func (qc *qdbCoordinator) ListKeyRanges(ctx context.Context, distribution string) ([]*kr.KeyRange, error) {
 	keyRanges, err := qc.db.ListKeyRanges(ctx, distribution)

diff --git a/coordinator/provider/keyranges.go b/coordinator/provider/keyranges.go
@@ -63,6 +63,22 @@ func (c *CoordinatorService) SplitKeyRange(ctx context.Context, request *protos.
 	return &protos.ModifyReply{}, nil
 }
 
+// GetKeyRange gets key ranges with given ids
+// TODO unit tests
+func (c *CoordinatorService) GetKeyRange(ctx context.Context, request *protos.GetKeyRangeRequest) (*protos.KeyRangeReply, error) {
+	res := make([]*protos.KeyRangeInfo, 0)
+	for _, id := range request.Ids {
+		krg, err := c.impl.GetKeyRange(ctx, id)
+		if err != nil {
+			return nil, err
+		}
+		if krg != nil {
+			res = append(res, krg.ToProto())
+		}
+	}
+	return &protos.KeyRangeReply{KeyRangesInfo: res}, nil
+}
+
 // TODO : unit tests
 func (c *CoordinatorService) ListKeyRange(ctx context.Context, request *protos.ListKeyRangeRequest) (*protos.KeyRangeReply, error) {
 

diff --git a/pkg/coord/adapter.go b/pkg/coord/adapter.go
@@ -39,6 +39,22 @@ func (a *Adapter) ShareKeyRange(id string) error {
 	return spqrerror.New(spqrerror.SPQR_NOT_IMPLEMENTED, "ShareKeyRange not implemented")
 }
 
+// GetKeyRange gets key range by id
+// TODO unit tests
+func (a *Adapter) GetKeyRange(ctx context.Context, krId string) (*kr.KeyRange, error) {
+	c := proto.NewKeyRangeServiceClient(a.conn)
+	reply, err := c.GetKeyRange(ctx, &proto.GetKeyRangeRequest{
+		Ids: []string{krId},
+	})
+	if err != nil {
+		return nil, err
+	}
+	if len(reply.KeyRangesInfo) == 0 {
+		return nil, nil
+	}
+	return kr.KeyRangeFromProto(reply.KeyRangesInfo[0]), nil
+}
+
 // TODO : unit tests
 func (a *Adapter) ListKeyRanges(ctx context.Context, distribution string) ([]*kr.KeyRange, error) {
 	c := proto.NewKeyRangeServiceClient(a.conn)

diff --git a/pkg/coord/local/clocal.go b/pkg/coord/local/clocal.go
@@ -400,6 +400,16 @@ func (qr *LocalCoordinator) Shards() []string {
 	return ret
 }
 
+// GetKeyRange gets key range by id
+// TODO unit tests
+func (lc *LocalCoordinator) GetKeyRange(ctx context.Context, krId string) (*kr.KeyRange, error) {
+	krDb, err := lc.qdb.GetKeyRange(ctx, krId)
+	if err != nil {
+		return nil, err
+	}
+	return kr.KeyRangeFromDB(krDb), nil
+}
+
 // TODO : unit tests
 func (qr *LocalCoordinator) ListKeyRanges(ctx context.Context, distribution string) ([]*kr.KeyRange, error) {
 	var ret []*kr.KeyRange

diff --git a/pkg/models/kr/keyrangemgr.go b/pkg/models/kr/keyrangemgr.go
@@ -20,6 +20,7 @@ type UniteKeyRange struct {
 }
 
 type KeyRangeMgr interface {
+	GetKeyRange(ctx context.Context, krId string) (*KeyRange, error)
 	ListKeyRanges(ctx context.Context, distribution string) ([]*KeyRange, error)
 	ListAllKeyRanges(ctx context.Context) ([]*KeyRange, error)
 	AddKeyRange(ctx context.Context, kr *KeyRange) error

diff --git a/pkg/protos/key_range.pb.go b/pkg/protos/key_range.pb.go
diff --git a/pkg/protos/key_range_grpc.pb.go b/pkg/protos/key_range_grpc.pb.go
diff --git a/protos/key_range.proto b/protos/key_range.proto
@@ -5,6 +5,7 @@ package spqr;
 option go_package = "spqr/proto";
 
 service KeyRangeService {
+  rpc GetKeyRange (GetKeyRangeRequest) returns (KeyRangeReply) {}
   rpc ListKeyRange (ListKeyRangeRequest) returns (KeyRangeReply) {}
   rpc ListAllKeyRanges (ListAllKeyRangesRequest) returns (KeyRangeReply) {}
   rpc LockKeyRange (LockKeyRangeRequest) returns (ModifyReply) {}
@@ -96,3 +97,7 @@ message ResolveKeyRangeRequest {
 message ResolveKeyRangeReply {
   repeated string key_range_d = 1;
 }
+
+message GetKeyRangeRequest {
+  repeated string ids = 1;
+}
diff --git a/router/grpc/qrouter.go b/router/grpc/qrouter.go
@@ -238,6 +238,22 @@ func (l *LocalQrouterServer) AddKeyRange(ctx context.Context, request *protos.Ad
 	return &protos.ModifyReply{}, nil
 }
 
+// GetKeyRange gets key ranges with given ids
+// TODO unit tests
+func (l *LocalQrouterServer) GetKeyRange(ctx context.Context, request *protos.GetKeyRangeRequest) (*protos.KeyRangeReply, error) {
+	res := make([]*protos.KeyRangeInfo, 0)
+	for _, id := range request.Ids {
+		krg, err := l.mgr.GetKeyRange(ctx, id)
+		if err != nil {
+			return nil, err
+		}
+		if krg != nil {
+			res = append(res, krg.ToProto())
+		}
+	}
+	return &protos.KeyRangeReply{KeyRangesInfo: res}, nil
+}
+
 // TODO : unit tests
 func (l *LocalQrouterServer) ListKeyRange(ctx context.Context, request *protos.ListKeyRangeRequest) (*protos.KeyRangeReply, error) {
 	var krs []*protos.KeyRangeInfo