feat(redis): add expire option (#35)

* feat(redis): add expire option * feat(redis): make the Key persistent when expire is 0 * chore: do not repeat the same string concat
vulsio · Aug 12, 2021 · a707cd1 · a707cd1
1 parent a56e1b6
commit a707cd1
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 11 deletions.
diff --git a/commands/fetchjvn.go b/commands/fetchjvn.go
@@ -19,9 +19,6 @@ var fetchJvnCmd = &cobra.Command{
 
 func init() {
 	RootCmd.AddCommand(fetchJvnCmd)
-
-	fetchJvnCmd.PersistentFlags().Bool("stdout", false, "display all CPEs to stdout")
-	_ = viper.BindPFlag("stdout", fetchJvnCmd.PersistentFlags().Lookup("stdout"))
 }
 
 func fetchJvn(cmd *cobra.Command, args []string) (err error) {

diff --git a/commands/fetchnvd.go b/commands/fetchnvd.go
@@ -19,9 +19,6 @@ var fetchNvdCmd = &cobra.Command{
 
 func init() {
 	RootCmd.AddCommand(fetchNvdCmd)
-
-	fetchNvdCmd.PersistentFlags().Bool("stdout", false, "display all CPEs to stdout")
-	_ = viper.BindPFlag("stdout", fetchNvdCmd.PersistentFlags().Lookup("stdout"))
 }
 
 func fetchNvd(cmd *cobra.Command, args []string) (err error) {

diff --git a/commands/root.go b/commands/root.go
@@ -49,6 +49,12 @@ func init() {
 
 	RootCmd.PersistentFlags().String("http-proxy", "", "http://proxy-url:port (default: empty)")
 	_ = viper.BindPFlag("http-proxy", RootCmd.PersistentFlags().Lookup("http-proxy"))
+
+	RootCmd.PersistentFlags().Bool("stdout", false, "display all CPEs to stdout")
+	_ = viper.BindPFlag("stdout", RootCmd.PersistentFlags().Lookup("stdout"))
+
+	RootCmd.PersistentFlags().Uint("expire", 0, "timeout to set for Redis keys in seconds. If set to 0, the key is persistent.")
+	_ = viper.BindPFlag("expire", RootCmd.PersistentFlags().Lookup("expire"))
 }
 
 // initConfig reads in config file and ENV variables if set.

diff --git a/db/redis.go b/db/redis.go
@@ -9,6 +9,7 @@ import (
 	"github.com/go-redis/redis/v8"
 	"github.com/inconshreveable/log15"
 	"github.com/kotakanbe/go-cpe-dictionary/models"
+	"github.com/spf13/viper"
 )
 
 const (
@@ -101,22 +102,40 @@ func (r *RedisDriver) GetCpesByVendorProduct(vendor, product string) ([]string,
 
 // InsertCpes Select Cve information from DB.
 func (r *RedisDriver) InsertCpes(cpes []models.CategorizedCpe) (err error) {
+	expire := viper.GetUint("expire")
+
 	ctx := context.Background()
 	bar := pb.New(len(cpes))
 	bar.Start()
+	vendorProductRootKey := hKeyPrefix + "VendorProduct"
 	for chunked := range chunkSlice(cpes, 10) {
-		var pipe redis.Pipeliner
-		pipe = r.conn.Pipeline()
+		pipe := r.conn.Pipeline()
 		for _, c := range chunked {
 			bar.Increment()
-			if result := pipe.ZAdd(ctx, hKeyPrefix+"VendorProduct", &redis.Z{Score: 0, Member: c.Vendor + sep + c.Product}); result.Err() != nil {
+			key := hKeyPrefix + c.Vendor + sep + c.Product
+			if result := pipe.ZAdd(ctx, vendorProductRootKey, &redis.Z{Score: 0, Member: c.Vendor + sep + c.Product}); result.Err() != nil {
 				return fmt.Errorf("Failed to ZAdd vendorProduct. err: %s", result.Err())
 			}
-			if result := pipe.ZAdd(ctx, hKeyPrefix+c.Vendor+sep+c.Product, &redis.Z{Score: 0, Member: c.CpeURI}); result.Err() != nil {
+			if result := pipe.ZAdd(ctx, key, &redis.Z{Score: 0, Member: c.CpeURI}); result.Err() != nil {
 				return fmt.Errorf("Failed to ZAdd CpeURI. err: %s", result.Err())
 			}
+			if expire > 0 {
+				if err := pipe.Expire(ctx, vendorProductRootKey, time.Duration(expire*uint(time.Second))).Err(); err != nil {
+					return fmt.Errorf("Failed to set Expire to Key. err: %s", err)
+				}
+				if err := pipe.Expire(ctx, key, time.Duration(expire*uint(time.Second))).Err(); err != nil {
+					return fmt.Errorf("Failed to set Expire to Key. err: %s", err)
+				}
+			} else {
+				if err := pipe.Persist(ctx, vendorProductRootKey).Err(); err != nil {
+					return fmt.Errorf("Failed to remove the existing timeout on Key. err: %s", err)
+				}
+				if err := pipe.Persist(ctx, key).Err(); err != nil {
+					return fmt.Errorf("Failed to remove the existing timeout on Key. err: %s", err)
+				}
+			}
 			if c.Deprecated {
-				if result := pipe.Set(ctx, fmt.Sprintf("%s%s", deprecatedPrefix, c.CpeURI), "true", time.Duration(0)); result.Err() != nil {
+				if result := pipe.Set(ctx, fmt.Sprintf("%s%s", deprecatedPrefix, c.CpeURI), "true", time.Duration(expire*uint(time.Second))); result.Err() != nil {
 					return fmt.Errorf("Failed to set to deprecated CPE. err: %s", result.Err())
 				}
 			}