@@ -11,6 +11,7 @@ import (
1111 "github.com/cheggaaa/pb/v3"
1212 "github.com/inconshreveable/log15"
1313 sqlite3 "github.com/mattn/go-sqlite3"
14+ "github.com/spf13/viper"
1415 "github.com/vulsio/go-cpe-dictionary/config"
1516 "github.com/vulsio/go-cpe-dictionary/models"
1617 "golang.org/x/xerrors"
@@ -43,12 +44,17 @@ func (r *RDBDriver) Name() string {
4344func (r * RDBDriver ) OpenDB (dbType , dbPath string , debugSQL bool ) (locked bool , err error ) {
4445 gormConfig := gorm.Config {
4546 DisableForeignKeyConstraintWhenMigrating : true ,
46- Logger : logger .Default .LogMode (logger .Silent ),
47+ Logger : logger .New (
48+ log .New (os .Stderr , "\r \n " , log .LstdFlags ),
49+ logger.Config {
50+ LogLevel : logger .Silent ,
51+ },
52+ ),
4753 }
4854
4955 if debugSQL {
5056 gormConfig .Logger = logger .New (
51- log .New (os .Stdout , "\r \n " , log .LstdFlags ),
57+ log .New (os .Stderr , "\r \n " , log .LstdFlags ),
5258 logger.Config {
5359 SlowThreshold : time .Second ,
5460 LogLevel : logger .Info ,
@@ -171,7 +177,7 @@ func (r *RDBDriver) GetVendorProducts() (vendorProducts []string, err error) {
171177 }
172178
173179 for _ , vp := range results {
174- vendorProducts = append (vendorProducts , fmt .Sprintf ("%s:: %s" , vp .Vendor , vp .Product ))
180+ vendorProducts = append (vendorProducts , fmt .Sprintf ("%s# %s" , vp .Vendor , vp .Product ))
175181 }
176182 return
177183}
@@ -209,6 +215,11 @@ func (r *RDBDriver) deleteAndInsertCpes(conn *gorm.DB, fetchType models.FetchTyp
209215 tx .Commit ()
210216 }()
211217
218+ batchSize := viper .GetInt ("batch-size" )
219+ if batchSize < 1 {
220+ return xerrors .New ("Failed to set batch-size. err: batch-size option is not set properly" )
221+ }
222+
212223 // Delete all old records
213224 oldIDs := []int64 {}
214225 result := tx .Model (models.CategorizedCpe {}).Select ("id" ).Where ("fetch_type = ?" , fetchType ).Find (& oldIDs )
@@ -218,15 +229,15 @@ func (r *RDBDriver) deleteAndInsertCpes(conn *gorm.DB, fetchType models.FetchTyp
218229
219230 if result .RowsAffected > 0 {
220231 log15 .Info (fmt .Sprintf ("Deleting records that match fetch_type = %s from your DB. This will take some time." , fetchType ))
221- for idx := range chunkSlice (len (oldIDs ), 10000 ) {
232+ for idx := range chunkSlice (len (oldIDs ), batchSize ) {
222233 if err := tx .Where ("id IN ?" , oldIDs [idx .From :idx .To ]).Delete (& models.CategorizedCpe {}).Error ; err != nil {
223234 return xerrors .Errorf ("Failed to delete: %w" , err )
224235 }
225236 }
226237 }
227238
228239 bar := pb .StartNew (len (cpes ))
229- for idx := range chunkSlice (len (cpes ), 2000 ) {
240+ for idx := range chunkSlice (len (cpes ), batchSize ) {
230241 if err := tx .Create (cpes [idx .From :idx .To ]).Error ; err != nil {
231242 return xerrors .Errorf ("Failed to insert. err: %w" , err )
232243 }
0 commit comments