Add ALTER REFERENCE TABLE SET STORAGE method

coordinator/pkg/clustered_coord.go

@@ -2644,6 +2644,52 @@ func (qc *ClusteredCoordinator) SyncReferenceRelations(ctx context.Context, relN
 	})
 }
 
+// AlterReferenceRelationStorage implements meta.EntityMgr.
+func (qc *ClusteredCoordinator) AlterReferenceRelationStorageAdvanced(ctx context.Context, relName *rfqn.RelationFQN, shs []string) error {
+	rel, err := qc.GetReferenceRelation(ctx, relName)
+	if err != nil {
+		return err
+	}
+	shardsExSet := make(map[string]struct{})
+	shardsToAdd := make([]string, 0)
+	shardsIntersect := make([]string, 0)
+	for _, sh := range rel.ShardIds {
+		shardsExSet[sh] = struct{}{}
+	}
+	for _, sh := range shs {
+		if _, ok := shardsExSet[sh]; !ok {
+			shardsToAdd = append(shardsToAdd, sh)
+		} else if ok {
+			shardsIntersect = append(shardsIntersect, sh)
+		}
+	}
+
+	if len(shardsIntersect) < len(rel.ShardIds) {
+		// We need to drop shards
+		if err := qc.db.AlterReferenceRelationStorage(ctx, relName, shardsIntersect); err != nil {
+			return fmt.Errorf("failed to alter reference relation storage: failed to remove excess shards in coordinator: %s", err)
+		}
+		if err := qc.traverseRouters(ctx, func(cc *grpc.ClientConn) error {
+			c := proto.NewReferenceRelationsServiceClient(cc)
+			_, err := c.AlterReferenceRelationStorage(ctx, &proto.AlterReferenceRelationStorageRequest{
+				Relation: rfqn.RelationFQNToProto(relName),
+				ShardIds: shardsIntersect,
+			})
+			return err
+		}); err != nil {
+			return fmt.Errorf("failed to alter reference relation storage: failed to remove excess shards in routers: %s", err)
+		}
+	}
+
+	rels := []*rfqn.RelationFQN{relName}
+	for _, sh := range shardsToAdd {
+		if err := qc.SyncReferenceRelations(ctx, rels, sh); err != nil {
+			return fmt.Errorf("failed to alter reference relation storage: failed to sync relation on shard \"%s\": %s", sh, err)
+		}
+	}
+	return nil
+}
+
 // TODO: unit tests
 func (qc *ClusteredCoordinator) DropReferenceRelation(ctx context.Context,
 	relName *rfqn.RelationFQN) error {

coordinator/provider/rrelations.go

@@ -43,3 +43,12 @@ func (rr *ReferenceRelationServer) DropReferenceRelations(ctx context.Context, r
 	}
 	return nil, nil
 }
+
+func (rr *ReferenceRelationServer) AlterReferenceRelationStorageAdvanced(ctx context.Context, req *protos.AlterReferenceRelationStorageRequest) (*emptypb.Empty, error) {
+	err := rr.impl.AlterReferenceRelationStorage(ctx, &rfqn.RelationFQN{
+		RelationName: req.Relation.RelationName,
+		SchemaName:   req.Relation.SchemaName,
+	}, req.ShardIds)
+
+	return nil, err
+}

pkg/coord/adapter.go

@@ -108,6 +108,16 @@ func (a *Adapter) AlterReferenceRelationStorage(ctx context.Context, relName *rf
 	return fmt.Errorf("AlterReferenceRelationStorage should not be used in proxy adapter")
 }
 
+// AlterReferenceRelationStorage implements meta.EntityMgr.
+func (a *Adapter) AlterReferenceRelationStorageAdvanced(ctx context.Context, relName *rfqn.RelationFQN, shs []string) error {
+	c := proto.NewReferenceRelationsServiceClient(a.conn)
+	_, err := c.AlterReferenceRelationStorageAdvanced(ctx, &proto.AlterReferenceRelationStorageRequest{
+		Relation: rfqn.RelationFQNToProto(relName),
+		ShardIds: shs,
+	})
+	return err
+}
+
 // CreateReferenceRelation implements meta.EntityMgr.
 func (a *Adapter) CreateReferenceRelation(ctx context.Context, r *rrelation.ReferenceRelation, entry []*rrelation.AutoIncrementEntry) error {
 	c := proto.NewReferenceRelationsServiceClient(a.conn)

pkg/coord/clocal.go

@@ -440,11 +440,6 @@ func (lc *LocalInstanceMetadataMgr) StopMoveTaskGroup(_ context.Context, _ strin
 	return ErrNotCoordinator
 }
 
-// AlterReferenceRelationStorage implements meta.EntityMgr.
-func (lc *LocalInstanceMetadataMgr) AlterReferenceRelationStorage(ctx context.Context, relName *rfqn.RelationFQN, shs []string) error {
-	return lc.qdb.AlterReferenceRelationStorage(ctx, relName, shs)
-}
-
 // SyncReferenceRelations implements meta.EntityMgr.
 func (lc *LocalInstanceMetadataMgr) SyncReferenceRelations(ctx context.Context, ids []*rfqn.RelationFQN, destShard string) error {
 	return ErrNotCoordinator

pkg/coord/coord.go

@@ -37,11 +37,16 @@ func NewCoordinator(q qdb.XQDB, d qdb.DCStateKeeper) Coordinator {
 	}
 }
 
-// AlterReferenceRelationStorage implements meta.EntityMgr.
+// AlterReferenceRelationStorage alters shards, on which reference relation is contained.
 func (lc *Coordinator) AlterReferenceRelationStorage(ctx context.Context, relName *rfqn.RelationFQN, shs []string) error {
 	return lc.qdb.AlterReferenceRelationStorage(ctx, relName, shs)
 }
 
+// AlterReferenceRelationStorage implements meta.EntityMgr.
+func (lc *Coordinator) AlterReferenceRelationStorageAdvanced(ctx context.Context, relName *rfqn.RelationFQN, shs []string) error {
+	return ErrNotCoordinator
+}
+
 // SyncReferenceRelations implements meta.EntityMgr.
 func (lc *Coordinator) SyncReferenceRelations(ctx context.Context, relNames []*rfqn.RelationFQN, destShard string) error {
 	for _, qualName := range relNames {

pkg/meta/meta.go

@@ -1181,6 +1181,17 @@ func ProcMetadataCommand(ctx context.Context,
 			}},
 		}
 
+		return tts, nil
+	case *spqrparser.AlterReferenceTableStorage:
+		if err := mgr.AlterReferenceRelationStorageAdvanced(ctx, stmt.RelationSelector, stmt.Shards); err != nil {
+			return nil, err
+		}
+		tts := &tupleslot.TupleTableSlot{
+			Desc: engine.GetVPHeader("relation", "shard"),
+		}
+		for _, sh := range stmt.Shards {
+			tts.WriteDataRow(stmt.RelationSelector.String(), sh)
+		}
 		return tts, nil
 	default:
 		return nil, ErrUnknownCoordinatorCommand

pkg/models/rrelation/rmgr.go

@@ -17,4 +17,5 @@ type ReferenceRelationMgr interface {
 
 	/* Method for cluster-wide reference relation data replication */
 	SyncReferenceRelations(ctx context.Context, ids []*rfqn.RelationFQN, destShard string) error
+	AlterReferenceRelationStorageAdvanced(ctx context.Context, relName *rfqn.RelationFQN, shs []string) error
 }

protos/reference_relation.proto

@@ -51,4 +51,5 @@ service ReferenceRelationsService {
   rpc AlterReferenceRelationStorage(AlterReferenceRelationStorageRequest) returns (google.protobuf.Empty) {}
 
   rpc SyncReferenceRelations(SyncReferenceRelationsRequest) returns (google.protobuf.Empty) {} 
+  rpc AlterReferenceRelationStorageAdvanced(AlterReferenceRelationStorageRequest) returns (google.protobuf.Empty) {}
 }