diff --git a/tools/db/main.go b/tools/db/main.go index 0b3504a..3ce77e6 100644 --- a/tools/db/main.go +++ b/tools/db/main.go @@ -1,13 +1,22 @@ package main import ( + "flag" "log" + "strings" "github.com/dgraph-io/badger/v2" ) +var ( + deleteRepoName = flag.String("deleterepo", "", "repo to delete from badger cache") + removeDuplicatesFlag = flag.Bool("removeduplicates", false, "remove non-lowercase duplicates from badger cache") + dryRun = flag.Bool("dryrun", false, "dry run mode") +) + func main() { - // delete a repo from badger cache + flag.Parse() + db, err := badger.Open(badger.DefaultOptions("/usr/local/badger").WithTruncate(true)) if err != nil { log.Fatal("ERROR: could not open badger db: ", err) @@ -15,13 +24,85 @@ func main() { defer db.Close() - repo := "[repo]" + if *deleteRepoName != "" { + err := deleteRepo(db, *deleteRepoName, *dryRun) + if err != nil { + log.Fatal(err) + } + } - err = db.Update(func(txn *badger.Txn) error { - return txn.Delete([]byte("repos-" + repo)) - }) - - if err != nil { - log.Fatal("Badger writing error:", err) + if *removeDuplicatesFlag { + err := removeDuplicates(db, *dryRun) + if err != nil { + log.Fatal(err) + } } } + +// delete a repo from badger cache +func deleteRepo(db *badger.DB, repo string, dryRun bool) error { + if !dryRun { + return db.Update(func(txn *badger.Txn) error { + return txn.Delete([]byte("repos-" + repo)) + }) + } + + log.Println("would delete", "repos-"+repo) + + return nil +} + +// remove duplicates from badger cache +func removeDuplicates(db *badger.DB, dryRun bool) error { + var toRemove []string + + err := db.View(func(txn *badger.Txn) error { + opts := badger.DefaultIteratorOptions + opts.PrefetchSize = 10 + it := txn.NewIterator(opts) + defer it.Close() + for it.Rewind(); it.Valid(); it.Next() { + item := it.Item() + k := string(item.Key()) + + if strings.HasPrefix(k, "repos-") { + if strings.ToLower(k) != k { + toRemove = append(toRemove, k) + } + } + } + + return nil + }) + + for _, k := range toRemove { + err := db.View(func(txn *badger.Txn) error { + _, err := txn.Get([]byte(strings.ToLower(k))) + + return err + }) + + if err == badger.ErrKeyNotFound { + continue + } + + if err != nil { + return err + } + + if !dryRun { + err = db.Update(func(txn *badger.Txn) error { + log.Println("Deleting", k) + return txn.Delete([]byte(k)) + }) + + if err != nil { + return err + } + } else { + log.Println("would delete", k) + } + } + + return err +}