seaweedfs/weed/shell/command_fs_tree.go
2019-04-05 11:40:54 -07:00

148 lines
3.2 KiB
Go

package shell
import (
"context"
"fmt"
"github.com/chrislusf/seaweedfs/weed/filer2"
"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
"io"
"strings"
)
func init() {
commands = append(commands, &commandFsTree{})
}
type commandFsTree struct {
}
func (c *commandFsTree) Name() string {
return "fs.tree"
}
func (c *commandFsTree) Help() string {
return `recursively list all files under a directory
fs.tree http://<filer_server>:<port>/dir/
`
}
func (c *commandFsTree) Do(args []string, commandEnv *commandEnv, writer io.Writer) (err error) {
filerServer, filerPort, path, err := commandEnv.parseUrl(findInputDirectory(args))
if err != nil {
return err
}
dir, name := filer2.FullPath(path).DirAndName()
ctx := context.Background()
return commandEnv.withFilerClient(ctx, filerServer, filerPort, func(client filer_pb.SeaweedFilerClient) error {
dirCount, fCount, terr := treeTraverseDirectory(ctx, writer, client, dir, name, newPrefix(), -1)
if terr == nil {
fmt.Fprintf(writer, "%d directories, %d files\n", dirCount, fCount)
}
return terr
})
}
func treeTraverseDirectory(ctx context.Context, writer io.Writer, client filer_pb.SeaweedFilerClient, dir, name string, prefix *Prefix, level int) (directoryCount, fileCount int64, err error) {
paginatedCount := -1
startFromFileName := ""
paginateSize := 1000
for paginatedCount == -1 || paginatedCount == paginateSize {
resp, listErr := client.ListEntries(ctx, &filer_pb.ListEntriesRequest{
Directory: dir,
Prefix: name,
StartFromFileName: startFromFileName,
InclusiveStartFrom: false,
Limit: uint32(paginateSize),
})
if listErr != nil {
err = listErr
return
}
paginatedCount = len(resp.Entries)
if paginatedCount > 0 {
prefix.addMarker(level)
}
for i, entry := range resp.Entries {
if level < 0 {
if entry.Name != name {
break
}
}
// 0.1% wrong prefix here, but fixing it would need to paginate to the next batch first
isLast := paginatedCount < paginateSize && i == paginatedCount-1
fmt.Fprintf(writer, "%s%s\n", prefix.getPrefix(level, isLast), entry.Name)
if entry.IsDirectory {
directoryCount++
subDir := fmt.Sprintf("%s/%s", dir, entry.Name)
if dir == "/" {
subDir = "/" + entry.Name
}
dirCount, fCount, terr := treeTraverseDirectory(ctx, writer, client, subDir, "", prefix, level+1)
directoryCount += dirCount
fileCount += fCount
err = terr
} else {
fileCount++
}
startFromFileName = entry.Name
}
}
return
}
type Prefix struct {
markers map[int]bool
}
func newPrefix() *Prefix {
return &Prefix{
markers: make(map[int]bool),
}
}
func (p *Prefix) addMarker(marker int) {
p.markers[marker] = true
}
func (p *Prefix) removeMarker(marker int) {
delete(p.markers, marker)
}
func (p *Prefix) getPrefix(level int, isLastChild bool) string {
var sb strings.Builder
if level < 0 {
return ""
}
for i := 0; i < level; i++ {
if _, ok := p.markers[i]; ok {
sb.WriteString("│")
} else {
sb.WriteString(" ")
}
sb.WriteString(" ")
}
if isLastChild {
sb.WriteString("└──")
p.removeMarker(level)
} else {
sb.WriteString("├──")
}
return sb.String()
}