diff --git a/drive/drive.go b/drive/drive.go index 62872cf..1072487 100644 --- a/drive/drive.go +++ b/drive/drive.go @@ -5,7 +5,6 @@ import ( "drivedlgo/db" "drivedlgo/utils" "fmt" - "github.com/vbauerster/mpb/v8" "io" "log" "net/http" @@ -15,11 +14,15 @@ import ( "sync" "time" + "github.com/fatih/color" + "github.com/vbauerster/mpb/v8" + "github.com/vbauerster/mpb/v8/decor" "golang.org/x/net/context" "golang.org/x/oauth2" "golang.org/x/oauth2/google" "google.golang.org/api/drive/v3" + "google.golang.org/api/option" ) var wg sync.WaitGroup @@ -34,6 +37,7 @@ type GoogleDriveClient struct { DriveSrv *drive.Service Progress *mpb.Progress abuse bool + numFilesDownloaded int channel chan int } @@ -55,39 +59,28 @@ func (G *GoogleDriveClient) SetConcurrency(count int) { G.channel = make(chan int, count) } +func (G *GoogleDriveClient) PrepareProgressBar(size int64, dec decor.Decorator) *mpb.Bar { + return G.Progress.AddBar(size, + mpb.PrependDecorators( + decor.Name("[ "), + dec, + decor.Name(" ] "), + decor.CountersKibiByte("% .2f / % .2f"), + ), + mpb.AppendDecorators( + decor.AverageETA(decor.ET_STYLE_GO), + decor.Name("]"), + decor.AverageSpeed(decor.SizeB1000(0), " % .2f"), + ), + ) +} + func (G *GoogleDriveClient) GetProgressBar(filename string, size int64) *mpb.Bar { - var bar *mpb.Bar if len(filename) > MAX_NAME_CHARACTERS { marquee := customdec.NewChangeNameDecor(filename, MAX_NAME_CHARACTERS) - bar = G.Progress.AddBar(size, - mpb.PrependDecorators( - decor.Name("[ "), - marquee.MarqueeText(), - decor.Name(" ] "), - decor.CountersKibiByte("% .2f / % .2f"), - ), - mpb.AppendDecorators( - decor.AverageETA(decor.ET_STYLE_GO), - decor.Name("]"), - decor.AverageSpeed(decor.UnitKiB, " % .2f"), - ), - ) - } else { - bar = G.Progress.AddBar(size, - mpb.PrependDecorators( - decor.Name("[ "), - decor.Name(filename, decor.WC{W: 5, C: decor.DidentRight}), - decor.Name(" ] "), - decor.CountersKibiByte("% .2f / % .2f"), - ), - mpb.AppendDecorators( - decor.AverageETA(decor.ET_STYLE_GO), - decor.Name("]"), - decor.AverageSpeed(decor.UnitKiB, " % .2f"), - ), - ) + return G.PrepareProgressBar(size, marquee.MarqueeText()) } - return bar + return G.PrepareProgressBar(size, decor.Name(filename, decor.WC{W: 5, C: decor.DSyncSpaceR})) } func (G *GoogleDriveClient) getClient(dbPath string, config *oauth2.Config, port int) *http.Client { @@ -170,7 +163,7 @@ func (G *GoogleDriveClient) Authorize(dbPath string, useSA bool, port int) { } client = G.getClient(dbPath, config, port) } - srv, err := drive.New(client) + srv, err := drive.NewService(context.Background(), option.WithHTTPClient(client)) if err != nil { log.Fatalf("Unable to retrieve Drive client: %v", err) } @@ -203,29 +196,36 @@ func (G *GoogleDriveClient) GetFilesByParentId(parentId string) []*drive.File { func (G *GoogleDriveClient) GetFileMetadata(fileId string) *drive.File { file, err := G.DriveSrv.Files.Get(fileId).Fields("name,mimeType,size,id,md5Checksum").SupportsAllDrives(true).Do() if err != nil { - log.Fatal(err) + fmt.Println(err) + os.Exit(0) } return file } func (G *GoogleDriveClient) Download(nodeId string, localPath string, outputPath string) { + startTime := time.Now() file := G.GetFileMetadata(nodeId) - fmt.Printf("Name: %s, MimeType: %s\n", file.Name, file.MimeType) if outputPath == "" { outputPath = utils.CleanupFilename(file.Name) } + fmt.Printf("%s(%s): %s -> %s/%s\n", color.HiBlueString("Download"), color.GreenString(file.MimeType), color.HiGreenString(file.Id), color.HiYellowString(localPath), color.HiYellowString(outputPath)) absPath := path.Join(localPath, outputPath) if file.MimeType == G.GDRIVE_DIR_MIMETYPE { err := os.MkdirAll(absPath, 0755) if err != nil { - log.Println("Error while creating directory: ", err.Error()) + fmt.Println("Error while creating directory: ", err.Error()) return } - G.TraverseNodes(file.Id, absPath) + files := G.GetFilesByParentId(file.Id) + if len(files) == 0 { + fmt.Println("google drive folder is empty.") + } else { + G.TraverseNodes(file.Id, absPath) + } } else { err := os.MkdirAll(localPath, 0755) if err != nil { - log.Println("Error while creating directory: ", err.Error()) + fmt.Println("Error while creating directory: ", err.Error()) return } G.channel <- 1 @@ -233,6 +233,8 @@ func (G *GoogleDriveClient) Download(nodeId string, localPath string, outputPath go G.HandleDownloadFile(file, absPath) } wg.Wait() + G.Progress.Wait() + fmt.Printf("%s", color.GreenString(fmt.Sprintf("Downloaded %d files in %s.\n", G.numFilesDownloaded, time.Now().Sub(startTime)))) } func (G *GoogleDriveClient) TraverseNodes(nodeId string, localPath string) { @@ -270,19 +272,17 @@ func (G *GoogleDriveClient) HandleDownloadFile(file *drive.File, absPath string) return } if bytesDled != 0 { - fmt.Printf("Resuming %s at offset %d\n", file.Name, bytesDled) + o := fmt.Sprintf("Resuming %s at offset %d\n", file.Name, bytesDled) + fmt.Printf("%s", color.GreenString(o)) } G.DownloadFile(file, absPath, bytesDled, 1) } func (G *GoogleDriveClient) DownloadFile(file *drive.File, localPath string, startByteIndex int64, retry int) bool { - cleanup := func() { - } writer, err := os.OpenFile(localPath, os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0644) defer writer.Close() if err != nil { log.Printf("[FileOpenError]: %v\n", err) - cleanup() return false } writer.Seek(startByteIndex, 0) @@ -296,7 +296,6 @@ func (G *GoogleDriveClient) DownloadFile(file *drive.File, localPath string, sta return G.DownloadFile(file, localPath, startByteIndex, retry+1) } log.Printf("[API-files:get]: (%s) %v\n", file.Id, err) - cleanup() return false } bar := G.GetProgressBar(file.Name, file.Size-startByteIndex) @@ -304,9 +303,10 @@ func (G *GoogleDriveClient) DownloadFile(file *drive.File, localPath string, sta defer proxyReader.Close() _, err = io.Copy(writer, proxyReader) if err != nil { - pos, posErr := writer.Seek(0, os.SEEK_CUR) + pos, posErr := writer.Seek(0, io.SeekCurrent) if posErr != nil { log.Printf("Error while getting current file offset, %v\n", err) + return false } else if retry <= MAX_RETRIES { log.Printf("err while copying stream: retrying download: %s: %v\n", file.Name, err) bar.Abort(true) @@ -315,8 +315,9 @@ func (G *GoogleDriveClient) DownloadFile(file *drive.File, localPath string, sta } else { log.Printf("Error while copying stream, %v\n", err) } + } else { + G.numFilesDownloaded += 1 } - cleanup() return true } diff --git a/go.mod b/go.mod index 8527af1..68a3aa0 100644 --- a/go.mod +++ b/go.mod @@ -6,7 +6,7 @@ require ( github.com/OpenPeeDeeP/xdg v1.0.0 github.com/prologic/bitcask v0.3.6 github.com/urfave/cli v1.22.10 - github.com/vbauerster/mpb/v8 v8.1.0 + github.com/vbauerster/mpb/v8 v8.7.1 golang.org/x/net v0.9.0 golang.org/x/oauth2 v0.7.0 ) @@ -17,6 +17,7 @@ require ( github.com/VividCortex/ewma v1.2.0 // indirect github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d // indirect github.com/cpuguy83/go-md2man/v2 v2.0.2 // indirect + github.com/fatih/color v1.16.0 // indirect github.com/gofrs/flock v0.8.1 // indirect github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect github.com/golang/protobuf v1.5.3 // indirect @@ -24,18 +25,19 @@ require ( github.com/google/uuid v1.3.0 // indirect github.com/googleapis/enterprise-certificate-proxy v0.2.3 // indirect github.com/googleapis/gax-go/v2 v2.8.0 // indirect - github.com/mattn/go-isatty v0.0.3 // indirect - github.com/mattn/go-runewidth v0.0.14 // indirect + github.com/mattn/go-colorable v0.1.13 // indirect + github.com/mattn/go-isatty v0.0.20 // indirect + github.com/mattn/go-runewidth v0.0.15 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/plar/go-adaptive-radix-tree v1.0.1 // indirect - github.com/rivo/uniseg v0.4.2 // indirect + github.com/rivo/uniseg v0.4.4 // indirect github.com/russross/blackfriday/v2 v2.1.0 // indirect github.com/stretchr/testify v1.8.1 // indirect github.com/vbauerster/mpb v3.4.0+incompatible // indirect go.opencensus.io v0.24.0 // indirect golang.org/x/crypto v0.1.0 // indirect golang.org/x/exp v0.0.0-20220921164117-439092de6870 // indirect - golang.org/x/sys v0.7.0 // indirect + golang.org/x/sys v0.15.0 // indirect golang.org/x/term v0.7.0 // indirect golang.org/x/text v0.9.0 // indirect google.golang.org/api v0.119.0 // indirect diff --git a/go.sum b/go.sum index 0276a9e..9fe3302 100644 --- a/go.sum +++ b/go.sum @@ -71,6 +71,8 @@ github.com/envoyproxy/go-control-plane v0.9.9-0.20201210154907-fd9021fe5dad/go.m github.com/envoyproxy/go-control-plane v0.9.10-0.20210907150352-cf90f659a021/go.mod h1:AFq3mo9L8Lqqiid3OhADV3RfLJnjiw63cSpi+fDTRC0= github.com/envoyproxy/protoc-gen-validate v0.1.0/go.mod h1:iSmxcyjqTsJpI2R4NaDN7+kN2VEUnK/pcBlmesArF7c= github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5KwzbycvMj4= +github.com/fatih/color v1.16.0 h1:zmkK9Ngbjj+K0yRhTVONQh1p/HknKYSlNT+vZCzyokM= +github.com/fatih/color v1.16.0/go.mod h1:fL2Sau1YI5c0pdGEVCbKQbLXB6edEj1ZgiY4NijnWvE= github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo= github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04= github.com/go-gl/glfw v0.0.0-20190409004039-e6da0acd62b1/go.mod h1:vR7hzQXu2zJy9AVAgeJqvqgH9Q5CA+iKCZ2gyEVpxRU= @@ -179,10 +181,17 @@ github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI= github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ= github.com/magiconair/properties v1.8.1/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ= github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU= +github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA= +github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg= github.com/mattn/go-isatty v0.0.3 h1:ns/ykhmWi7G9O+8a448SecJU3nSMBXJfqQkl0upE1jI= github.com/mattn/go-isatty v0.0.3/go.mod h1:M+lRXTBqGeGNdLjl/ufCoiOlB5xdOkqRJdNxMWT7Zi4= +github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM= +github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY= +github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y= github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWVwUuU= github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= +github.com/mattn/go-runewidth v0.0.15 h1:UNAjwbU9l54TA3KzvqLGxwWjHmMgBUVhBiTjelZgg3U= +github.com/mattn/go-runewidth v0.0.15/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w= github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0= github.com/miekg/dns v1.0.14/go.mod h1:W1PPwlIAgtquWBMBEV9nkV9Cazfe8ScdGz/Lj7v3Nrg= github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc= @@ -224,6 +233,8 @@ github.com/prometheus/tsdb v0.7.1/go.mod h1:qhTCs0VvXwvX/y3TZrWD7rabWM+ijKTux40T github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc= github.com/rivo/uniseg v0.4.2 h1:YwD0ulJSJytLpiaWua0sBDusfsCZohxjxzVTYjwxfV8= github.com/rivo/uniseg v0.4.2/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= +github.com/rivo/uniseg v0.4.4 h1:8TfxU8dW6PdqD27gjM8MVNuicgxIjxpm4K7x4jp8sis= +github.com/rivo/uniseg v0.4.4/go.mod h1:FN3SvrM+Zdj16jyLfmOkMNblXMcoc8DfTHruCPUcx88= github.com/rogpeppe/fastuuid v0.0.0-20150106093220-6724a57986af/go.mod h1:XWv6SoW27p1b0cqNHllgS5HIMJraePCO15w5zCzIWYg= github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ= github.com/rogpeppe/go-internal v1.3.0/go.mod h1:M8bDsm7K2OlrFYOpmOWEs/qY81heoFRclV5y23lUDJ4= @@ -275,6 +286,8 @@ github.com/vbauerster/mpb v3.4.0+incompatible h1:mfiiYw87ARaeRW6x5gWwYRUawxaW1tL github.com/vbauerster/mpb v3.4.0+incompatible/go.mod h1:zAHG26FUhVKETRu+MWqYXcI70POlC6N8up9p1dID7SU= github.com/vbauerster/mpb/v8 v8.1.0 h1:Yc3r28RoSLT3nxrCJVo5H07DL/JpF4yg3ToSeDtBK1c= github.com/vbauerster/mpb/v8 v8.1.0/go.mod h1:v0gv/von+tu9oTWk6TM9mJULwUZ2rcxI81u1YDIPJvg= +github.com/vbauerster/mpb/v8 v8.7.1 h1:bQoSMMTFAg/gjsLrBYmO8gbRcZt7aDq6WI2IMa9BTqM= +github.com/vbauerster/mpb/v8 v8.7.1/go.mod h1:fWgXcAu4W+0cBSUh4ZlaKJyC2KtgU27ZSTaiIk0QNsQ= github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU= github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q= go.etcd.io/bbolt v1.3.2/go.mod h1:IbVyRI1SCnLcuJnV2u8VeU0CEYM7e686BmAb1XKL+uU= @@ -378,10 +391,14 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w golang.org/x/sys v0.0.0-20210423082822-04245dca01da/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20211019181941-9d821ace8654/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8 h1:h+EGohizhe9XlX18rfpa8k8RAc5XyaeamM+0VHRd4lc= golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.7.0 h1:3jlCCIQZPdOYu1h8BkNvLz8Kgwtae2cagcG/VamtZRU= golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= +golang.org/x/sys v0.15.0 h1:h48lPFYpsTvQJZF4EKyI4aLHaev3CxivZmv7yZig9pc= +golang.org/x/sys v0.15.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.7.0 h1:BEvjmm5fURWqcfbSKTdpkDXYBrUS1c0m8agp14W48vQ= golang.org/x/term v0.7.0/go.mod h1:P32HKFT3hSsZrRxla30E9HqToFYAQPCMs/zFMBUFqPY=