refactor: use std slices, maps package (#278)

vulsio · Sep 27, 2024 · bbe0b9f · bbe0b9f
1 parent 11de4e4
commit bbe0b9f
Show file tree

Hide file tree

Showing 12 changed files with 74 additions and 92 deletions.
diff --git a/db/arch.go b/db/arch.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"slices"
 
 	"github.com/cheggaaa/pb/v3"
 	"github.com/spf13/viper"
@@ -82,11 +83,11 @@ func (r *RDBDriver) deleteAndInsertArch(advs []models.ArchADV) (err error) {
 		return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
 	}
 
-	for idx := range chunkSlice(len(advs), batchSize) {
-		if err = tx.Create(advs[idx.From:idx.To]).Error; err != nil {
+	for chunk := range slices.Chunk(advs, batchSize) {
+		if err = tx.Create(chunk).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
 		}
-		bar.Add(idx.To - idx.From)
+		bar.Add(len(chunk))
 	}
 	bar.Finish()
 

diff --git a/db/db.go b/db/db.go
@@ -91,26 +91,3 @@ func newDB(dbType string) (DB, error) {
 	}
 	return nil, fmt.Errorf("Invalid database dialect. dbType: %s", dbType)
 }
-
-// IndexChunk has a starting point and an ending point for Chunk
-type IndexChunk struct {
-	From, To int
-}
-
-func chunkSlice(length int, chunkSize int) <-chan IndexChunk {
-	ch := make(chan IndexChunk)
-
-	go func() {
-		defer close(ch)
-
-		for i := 0; i < length; i += chunkSize {
-			idx := IndexChunk{i, i + chunkSize}
-			if length < idx.To {
-				idx.To = length
-			}
-			ch <- idx
-		}
-	}()
-
-	return ch
-}
diff --git a/db/debian.go b/db/debian.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"slices"
 
 	"github.com/cheggaaa/pb/v3"
 	"github.com/spf13/viper"
@@ -90,11 +91,11 @@ func (r *RDBDriver) deleteAndInsertDebian(cves []models.DebianCVE) (err error) {
 		return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
 	}
 
-	for idx := range chunkSlice(len(cves), batchSize) {
-		if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil {
+	for chunk := range slices.Chunk(cves, batchSize) {
+		if err = tx.Create(chunk).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
 		}
-		bar.Add(idx.To - idx.From)
+		bar.Add(len(chunk))
 	}
 	bar.Finish()
 

diff --git a/db/microsoft.go b/db/microsoft.go
@@ -4,15 +4,15 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"maps"
 	"os"
+	"slices"
 	"strconv"
 	"strings"
 
 	"github.com/cheggaaa/pb/v3"
 	"github.com/inconshreveable/log15"
 	"github.com/spf13/viper"
-	"golang.org/x/exp/maps"
-	"golang.org/x/exp/slices"
 	"golang.org/x/xerrors"
 	"gorm.io/gorm"
 
@@ -69,7 +69,7 @@ func (r *RDBDriver) GetExpandKB(applied []string, unapplied []string) ([]string,
 		uniqUnappliedKBIDs[kbID] = struct{}{}
 		delete(uniqAppliedKBIDs, kbID)
 	}
-	applied = maps.Keys(uniqAppliedKBIDs)
+	applied = slices.Collect(maps.Keys(uniqAppliedKBIDs))
 
 	if len(applied) > 0 {
 		relations := []models.MicrosoftKBRelation{}
@@ -102,7 +102,7 @@ func (r *RDBDriver) GetExpandKB(applied []string, unapplied []string) ([]string,
 
 		if err := r.conn.
 			Preload("SupersededBy").
-			Where("kb_id IN ?", maps.Keys(uniqUnappliedKBIDs)).
+			Where("kb_id IN ?", slices.Collect(maps.Keys(uniqUnappliedKBIDs))).
 			Find(&relations).Error; err != nil {
 			return nil, nil, xerrors.Errorf("Failed to get KB Relation by unapplied KBID: %q. err: %w", unapplied, err)
 		}
@@ -114,7 +114,7 @@ func (r *RDBDriver) GetExpandKB(applied []string, unapplied []string) ([]string,
 		}
 	}
 
-	return applied, maps.Keys(uniqUnappliedKBIDs), nil
+	return applied, slices.Collect(maps.Keys(uniqUnappliedKBIDs)), nil
 }
 
 // GetRelatedProducts :
@@ -284,11 +284,11 @@ func (r *RDBDriver) deleteAndInsertMicrosoft(cves []models.MicrosoftCVE) (err er
 		return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
 	}
 
-	for idx := range chunkSlice(len(cves), batchSize) {
-		if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil {
+	for chunk := range slices.Chunk(cves, batchSize) {
+		if err = tx.Create(chunk).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
 		}
-		bar.Add(idx.To - idx.From)
+		bar.Add(len(chunk))
 	}
 	bar.Finish()
 
@@ -324,11 +324,11 @@ func (r *RDBDriver) deleteAndInsertMicrosoftKBRelation(kbs []models.MicrosoftKBR
 		return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
 	}
 
-	for idx := range chunkSlice(len(kbs), batchSize) {
-		if err = tx.Create(kbs[idx.From:idx.To]).Error; err != nil {
+	for chunk := range slices.Chunk(kbs, batchSize) {
+		if err = tx.Create(chunk).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
 		}
-		bar.Add(idx.To - idx.From)
+		bar.Add(len(chunk))
 	}
 	bar.Finish()
 	return nil

diff --git a/db/rdb.go b/db/rdb.go
@@ -186,9 +186,7 @@ func (r *RDBDriver) MigrateDB() error {
 				}
 			}
 		case dialectMysql, dialectPostgreSQL:
-			if err != nil {
-				return xerrors.Errorf("Failed to migrate. err: %w", err)
-			}
+			return xerrors.Errorf("Failed to migrate. err: %w", err)
 		default:
 			return xerrors.Errorf("Not Supported DB dialects. r.name: %s", r.name)
 		}

diff --git a/db/redhat.go b/db/redhat.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"slices"
 	"strings"
 	"time"
 
@@ -222,11 +223,11 @@ func (r *RDBDriver) deleteAndInsertRedhat(cves []models.RedhatCVE) (err error) {
 		return fmt.Errorf("Failed to set batch-size. err: batch-size option is not set properly")
 	}
 
-	for idx := range chunkSlice(len(cves), batchSize) {
-		if err = tx.Create(cves[idx.From:idx.To]).Error; err != nil {
+	for chunk := range slices.Chunk(cves, batchSize) {
+		if err = tx.Create(chunk).Error; err != nil {
 			return xerrors.Errorf("Failed to insert. err: %w", err)
 		}
-		bar.Add(idx.To - idx.From)
+		bar.Add(len(chunk))
 	}
 	bar.Finish()