feat(db) no progress when --log-json option (#147)

vulsio · Jul 4, 2024 · 0c8a289 · 0c8a289
1 parent 04189c0
commit 0c8a289
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/db/rdb.go b/db/rdb.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"time"
@@ -183,7 +184,12 @@ func (r *RDBDriver) InsertMetasploit(records []models.Metasploit) (err error) {
 }
 
 func (r *RDBDriver) deleteAndInsertMetasploit(records []models.Metasploit) (err error) {
-	bar := pb.StartNew(len(records))
+	bar := pb.StartNew(len(records)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	tx := r.conn.Begin()
 	defer func() {
 		if err != nil {

diff --git a/db/redis.go b/db/redis.go
@@ -6,6 +6,8 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
+	"os"
 	"strconv"
 	"strings"
 	"time"
@@ -202,7 +204,12 @@ func (r *RedisDriver) InsertMetasploit(records []models.Metasploit) (err error)
 	}
 
 	log15.Info("Inserting Modules having CVEs...")
-	bar := pb.StartNew(len(records))
+	bar := pb.StartNew(len(records)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	for idx := range chunkSlice(len(records), batchSize) {
 		pipe := r.conn.Pipeline()
 		for _, record := range records[idx.From:idx.To] {