PackageManager/Client/Filesystem/PickFiles.go

package Filesystem

import (
	"os"

	"github.com/vbauerster/mpb"
	bolt "go.etcd.io/bbolt"

	"PackageManager/Client/Database"
	"PackageManager/Client/ProgressBar"
	"PackageManager/Variables"
)

func pickFilesSingle(rootPath string) error {
	var (
		indexBucket *bolt.Bucket
		picksBucket *bolt.Bucket
		e           error
	)
	e = Database.FsDB.Batch(func(tx *bolt.Tx) error {
		indexBucket = tx.Bucket(Variables.FsHashIndexBucket)
		picksBucket = tx.Bucket(Variables.FsHashPicksBucket)

		e = AddFileToBucket(picksBucket, rootPath)
		if e != nil {
			return e
		}
		return RemoveFileFromBucket(indexBucket, rootPath)
	})
	return e
}

func pickFilesRecursive(rootPath string) error {
	var (
		fsStatus    FilesystemStatus
		indexBucket *bolt.Bucket
		picksBucket *bolt.Bucket
		bar         *mpb.Bar
		totalLen    int
		f           string
		e           error
	)

	fsStatus, e = GetFilesystemDiff(rootPath)
	if e != nil {
		return e
	}

	totalLen = len(fsStatus.NewFiles) + len(fsStatus.ModifiedFiles) + len(fsStatus.MissingFiles)

	if totalLen == 0 {
		return nil
	}

	bar = ProgressBar.InitBar("Adding...", totalLen)

	e = Database.FsDB.Batch(func(tx *bolt.Tx) error {
		indexBucket = tx.Bucket(Variables.FsHashIndexBucket)
		picksBucket = tx.Bucket(Variables.FsHashPicksBucket)

		if len(fsStatus.NewFiles) > 0 {
			for _, f = range fsStatus.NewFiles {
				bar.Increment()
				e = AddFileToBucket(picksBucket, f)
				if e != nil {
					return e
				}
			}
		}

		if len(fsStatus.ModifiedFiles) > 0 {
			for _, f = range fsStatus.ModifiedFiles {
				bar.Increment()
				e = AddFileToBucket(picksBucket, f)
				if e != nil {
					return e
				}
			}
		}

		if len(fsStatus.MissingFiles) > 0 {
			for _, f = range fsStatus.MissingFiles {
				bar.Increment()
				e = RemoveFileFromBucket(indexBucket, f)
				if e != nil {
					return e
				}
				e = RemoveFileFromBucket(picksBucket, f)
				if e != nil {
					return e
				}
			}
		}

		return nil
	})

	return e
}

func PickFiles(rootPath string) error {
	var (
		rootStat os.FileInfo
		e        error
	)

	rootStat, e = os.Stat(rootPath)
	if e != nil {
		return e
	}

	if !rootStat.IsDir() {
		return pickFilesSingle(rootPath)
	}

	return pickFilesRecursive(rootPath)
}

func ResetAllPickedFiles() error {
	var (
		e error
	)

	e = Database.FsDB.Batch(func(tx *bolt.Tx) error {
		return tx.DeleteBucket(Variables.FsHashPicksBucket)
	})

	return e
}