From 2cdb7fda2683df10087e0cd995d17b24e509aa3c Mon Sep 17 00:00:00 2001
From: Sadayuki Matsuno <sadayuki.matsuno@gmail.com>
Date: Wed, 3 Jul 2024 13:55:42 +0900
Subject: [PATCH] feat(db) no progress when --log-json option

---
 db/rdb.go   | 29 +++++++++++++++++++++++++----
 db/redis.go | 29 +++++++++++++++++++++++++----
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/db/rdb.go b/db/rdb.go
index f7142a0..aa4cb38 100644
--- a/db/rdb.go
+++ b/db/rdb.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"os"
 	"time"
@@ -628,7 +629,12 @@ func deleteJvn(tx *gorm.DB) error {
 }
 
 func insertJvn(tx *gorm.DB, cves []models.Jvn, batchSize int) error {
-	bar := pb.StartNew(len(cves))
+	bar := pb.StartNew(len(cves)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	for _, cve := range cves {
 		if err := tx.Omit("Cpes").Create(&cve).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
@@ -727,7 +733,12 @@ func deleteNvd(tx *gorm.DB) error {
 }
 
 func insertNvd(tx *gorm.DB, cves []models.Nvd, batchSize int) error {
-	bar := pb.StartNew(len(cves))
+	bar := pb.StartNew(len(cves)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	for _, cve := range cves {
 		if err := tx.Omit("Cpes").Create(&cve).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
@@ -782,7 +793,12 @@ func (r *RDBDriver) InsertFortinet(advs []models.Fortinet) (err error) {
 	}
 
 	logger.Infof("Inserting fetched %d CVEs...", len(advs))
-	bar := pb.StartNew(len(advs))
+	bar := pb.StartNew(len(advs)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	for _, adv := range advs {
 		if err := tx.Omit("Cpes").Create(&adv).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
@@ -930,7 +946,12 @@ func deleteMitre(tx *gorm.DB) error {
 }
 
 func insertMitre(tx *gorm.DB, cves []models.Mitre, _ int) error {
-	bar := pb.StartNew(len(cves))
+	bar := pb.StartNew(len(cves)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	for _, cve := range cves {
 		if err := tx.Create(&cve).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
diff --git a/db/redis.go b/db/redis.go
index 98aae66..0b5ec7e 100644
--- a/db/redis.go
+++ b/db/redis.go
@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 	"strconv"
 	"strings"
@@ -443,7 +444,12 @@ func (r *RedisDriver) InsertJvn(years []string) error {
 		delete(uniqCves, year)
 
 		log.Infof("Inserting fetched CVEs(%s)...", year)
-		bar := pb.StartNew(len(cves))
+		bar := pb.StartNew(len(cves)).SetWriter(func() io.Writer {
+			if viper.GetBool("log-json") {
+				return io.Discard
+			}
+			return os.Stderr
+		}())
 		for idx := range chunkSlice(len(cves), batchSize) {
 			pipe := r.conn.Pipeline()
 			for _, cve := range cves[idx.From:idx.To] {
@@ -575,7 +581,12 @@ func (r *RedisDriver) InsertNvd(years []string) error {
 		}
 
 		log.Infof("Inserting fetched CVEs(%s)...", year)
-		bar := pb.StartNew(len(cves))
+		bar := pb.StartNew(len(cves)).SetWriter(func() io.Writer {
+			if viper.GetBool("log-json") {
+				return io.Discard
+			}
+			return os.Stderr
+		}())
 		for idx := range chunkSlice(len(cves), batchSize) {
 			pipe := r.conn.Pipeline()
 			for _, cve := range cves[idx.From:idx.To] {
@@ -681,7 +692,12 @@ func (r *RedisDriver) InsertFortinet(advs []models.Fortinet) error {
 	}
 
 	log.Infof("Inserting fetched %d CVEs...", len(advs))
-	bar := pb.StartNew(len(advs))
+	bar := pb.StartNew(len(advs)).SetWriter(func() io.Writer {
+		if viper.GetBool("log-json") {
+			return io.Discard
+		}
+		return os.Stderr
+	}())
 	for idx := range chunkSlice(len(advs), batchSize) {
 		pipe := r.conn.Pipeline()
 		for _, adv := range advs[idx.From:idx.To] {
@@ -810,7 +826,12 @@ func (r *RedisDriver) InsertMitre(years []string) error {
 		}
 
 		log.Infof("Inserting fetched CVEs(%s)...", year)
-		bar := pb.StartNew(len(cves))
+		bar := pb.StartNew(len(cves)).SetWriter(func() io.Writer {
+			if viper.GetBool("log-json") {
+				return io.Discard
+			}
+			return os.Stderr
+		}())
 		for idx := range chunkSlice(len(cves), batchSize) {
 			pipe := r.conn.Pipeline()
 			for _, cve := range cves[idx.From:idx.To] {