Merge pull request #44 from kqzh/feat_get_space_used

internal/clients/spaceusage.go

@@ -0,0 +1,51 @@
+package clients
+
+import (
+	"os"
+	"path/filepath"
+	"strconv"
+
+	pb "github.com/vesoft-inc/nebula-agent/v3/pkg/proto"
+)
+
+type SpaceUsage struct {
+	DataPath string
+}
+
+func NewSpaceUsage(dataPath string) *SpaceUsage {
+	return &SpaceUsage{
+		DataPath: dataPath,
+	}
+}
+
+func (s *SpaceUsage) GetSpaceUsages() (*pb.GetSpaceUsagesResponse, error) {
+	spaceUsages := make([]*pb.GetSpaceUsagesResponse_SpaceUsageItem, 0)
+	dirs, err := filepath.Glob(filepath.Join(s.DataPath, "*/"))
+	if err != nil {
+		return nil, err
+	}
+	for _, dir := range dirs {
+		id, err := strconv.Atoi(filepath.Base(dir))
+		if err != nil {
+			continue
+		}
+
+		var size int64
+		err = filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
+			if !info.IsDir() {
+				size += info.Size()
+			}
+			return err
+		})
+		if err != nil {
+			return nil, err
+		}
+
+		spaceUsages = append(spaceUsages, &pb.GetSpaceUsagesResponse_SpaceUsageItem{
+			Id:    int64(id),
+			Usage: size,
+		})
+	}
+
+	return &pb.GetSpaceUsagesResponse{SpaceUsages: spaceUsages}, nil
+}

internal/server/agent.go

@@ -94,3 +94,7 @@ func (a *AgentServer) HealthCheck(ctx context.Context, req *pb.HealthCheckReques
 		Status: "healthy",
 	}, nil
 }
+
+func (a *AgentServer) GetSpaceUsages(ctx context.Context, req *pb.GetSpaceUsagesRequest) (*pb.GetSpaceUsagesResponse, error) {
+	return clients.NewSpaceUsage(req.DataPath).GetSpaceUsages()
+}

pkg/client/client.go

@@ -33,6 +33,7 @@ type Client interface {
 	ExistDir(req *pb.ExistDirRequest) (*pb.ExistDirResponse, error)
 	StopAgent(req *pb.StopAgentRequest) (*pb.StopAgentResponse, error)
 	HealthCheck(req *pb.HealthCheckRequest) (*pb.HealthCheckResponse, error)
+	GetSpaceUsages(req *pb.GetSpaceUsagesRequest) (*pb.GetSpaceUsagesResponse, error)
 }
 
 func genSessionId() string {
@@ -199,3 +200,13 @@ func (c *client) HealthCheck(req *pb.HealthCheckRequest) (resp *pb.HealthCheckRe
 
 	return c.agent.HealthCheck(c.ctx, req)
 }
+
+func (c *client) GetSpaceUsages(req *pb.GetSpaceUsagesRequest) (resp *pb.GetSpaceUsagesResponse, err error) {
+	defer func() {
+		if err != nil {
+			err = fmt.Errorf("agent, get space usages failed: %w", err)
+		}
+	}()
+
+	return c.agent.GetSpaceUsages(c.ctx, req)
+}