Add configurable DB option

darenliang · Jul 3, 2023 · 201bac4 · 201bac4
1 parent 2558daf
commit 201bac4
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/db.go b/db.go
@@ -125,6 +125,17 @@ func (iter *MapDBIterator) Next() (string, *Tx, bool) {
 	return entry.Key, entry.Tx, entry.Ok
 }
 
+func GetNewDB(dbType string) DB {
+	switch dbType {
+	case "radix":
+		return NewRadixDB()
+	case "map":
+		return NewMapDB()
+	default:
+		panic("unknown db type")
+	}
+}
+
 // prepareChannels prepares the tx and data channels
 // creating channels if necessary
 func prepareChannels(dg *discordgo.Session, guildID string) (*discordgo.Channel, *discordgo.Channel, error) {
@@ -161,8 +172,8 @@ func prepareChannels(dg *discordgo.Session, guildID string) (*discordgo.Channel,
 // setupDB setups the in-mem database
 // This function needs to be refactored; it looks really gross in its current
 // state.
-func setupDB(dg *discordgo.Session, txChannel *discordgo.Channel, compact bool) (DB, error) {
-	db := NewRadixDB()
+func setupDB(dg *discordgo.Session, txChannel *discordgo.Channel, compact bool, dbType string) (DB, error) {
+	db := GetNewDB(dbType)
 
 	txBuffer := &bytes.Buffer{}
 	var pinnedMsg *discordgo.Message

diff --git a/main.go b/main.go
@@ -21,7 +21,8 @@ var (
 	guildID   string
 	mount     string
 	compact   bool
-	cache     string
+	cacheType string
+	dbType    string
 	debug     bool
 	port      int
 	options   []string
@@ -36,7 +37,8 @@ func main() {
 	kingpin.Flag("user", "Token is a user token").Short('u').BoolVar(&userToken)
 	kingpin.Flag("mount", "Mount point").Short('m').StringVar(&mount)
 	kingpin.Flag("compact", "Compact transactions").Short('x').BoolVar(&compact)
-	kingpin.Flag("cache", "Cache type").Short('c').Default("disk").EnumVar(&cache, "disk", "memory")
+	kingpin.Flag("cache", "Cache type").Short('c').Default("disk").EnumVar(&cacheType, "disk", "memory")
+	kingpin.Flag("db", "Database type").Short('d').Default("radix").EnumVar(&dbType, "radix", "map")
 	kingpin.Flag("verbose", "Enable pprof and print debug logs").Short('v').BoolVar(&debug)
 	kingpin.Flag("port", "Port to run pprof on").Short('p').Default("8000").IntVar(&port)
 	kingpin.Flag("options", "FUSE options").Short('o').StringsVar(&options)
@@ -100,15 +102,15 @@ func main() {
 		return
 	}
 
-	db, err := setupDB(dg, txChannel, compact)
+	db, err := setupDB(dg, txChannel, compact, dbType)
 	if err != nil {
 		zap.S().Error(err)
 		return
 	}
 
 	writer := setupWriter(dg, txChannel.ID, dataChannel.ID)
 
-	dsfs = NewDsfs(dg, db, writer, txChannel, dataChannel, cache)
+	dsfs = NewDsfs(dg, db, writer, txChannel, dataChannel, cacheType)
 	dsfsReady.Store(true)
 
 	host := fuse.NewFileSystemHost(dsfs)