migrating db in vulndb commands

Author: timbastin

README.md

@@ -105,7 +105,7 @@ DevGuard uses [golang-migrate](https://github.com/golang-migrate/migrate) for da
 ### How Database Migrations Work
 
 1. **Automatic Migration**: By default, migrations run automatically when the application starts
-2. **Environment Control**: Set `AUTO_MIGRATE=false` to disable automatic migrations
+2. **Environment Control**: Set `DISABLE_AUTOMIGRATE=true` to disable automatic migrations
 3. **Embedded Migrations**: Migration files are embedded in the binary for easy deployment
 4. **Idempotent**: Migrations can be run multiple times safely
 

cmd/devguard-cli/commands/vulndb.go

@@ -1,13 +1,17 @@
 package commands
 
 import (
+	"errors"
 	"log/slog"
+	"os"
 	"regexp"
 	"strings"
 	"time"
 
 	"github.com/l3montree-dev/devguard/internal/core"
 	"github.com/l3montree-dev/devguard/internal/core/vulndb"
+	"github.com/l3montree-dev/devguard/internal/database"
+	"github.com/l3montree-dev/devguard/internal/database/models"
 	"github.com/l3montree-dev/devguard/internal/database/repositories"
 	"github.com/spf13/cobra"
 )
@@ -51,19 +55,41 @@ func isValidCVE(cveID string) bool {
 	return r.MatchString(cveID)
 }
 
+func migrateDB(db core.DB) {
+	// Run database migrations using the existing database connection
+	disableAutoMigrate := os.Getenv("DISABLE_AUTOMIGRATE")
+	if disableAutoMigrate != "true" {
+		slog.Info("running database migrations...")
+		if err := database.RunMigrationsWithDB(db); err != nil {
+			slog.Error("failed to run database migrations", "error", err)
+			panic(errors.New("Failed to run database migrations"))
+		}
+
+		// Run hash migrations if needed (when algorithm version changes)
+		if err := models.RunHashMigrationsIfNeeded(db); err != nil {
+			slog.Error("failed to run hash migrations", "error", err)
+			panic(errors.New("Failed to run hash migrations"))
+		}
+	} else {
+		slog.Info("automatic migrations disabled via DISABLE_AUTOMIGRATE=true")
+	}
+}
+
 func newImportCVECommand() *cobra.Command {
 	importCmd := &cobra.Command{
 		Use:   "import-cve",
 		Short: "Will import the vulnerability database",
 		Args:  cobra.ExactArgs(1),
 		Run: func(cmd *cobra.Command, args []string) {
 			core.LoadConfig() // nolint
-			database, err := core.DatabaseFactory()
+			db, err := core.DatabaseFactory()
 			if err != nil {
 				slog.Error("could not connect to database", "err", err)
 				return
 			}
 
+			migrateDB(db)
+
 			cveID := args[0]
 			cveID = strings.TrimSpace(strings.ToUpper(cveID))
 			// check if first argument is valid cve
@@ -72,9 +98,9 @@ func newImportCVECommand() *cobra.Command {
 				return
 			}
 
-			cveRepository := repositories.NewCVERepository(database)
+			cveRepository := repositories.NewCVERepository(db)
 			nvdService := vulndb.NewNVDService(cveRepository)
-			osvService := vulndb.NewOSVService(repositories.NewAffectedComponentRepository(database))
+			osvService := vulndb.NewOSVService(repositories.NewAffectedComponentRepository(db))
 
 			cve, err := nvdService.ImportCVE(cveID)
 
@@ -155,27 +181,29 @@ func newSyncCommand() *cobra.Command {
 
 			core.LoadConfig() // nolint
 
-			database, err := core.DatabaseFactory()
+			db, err := core.DatabaseFactory()
 			if err != nil {
 				slog.Error("could not connect to database", "err", err)
 				return
 			}
 
+			migrateDB(db)
+
 			databasesToSync, _ := cmd.Flags().GetStringArray("databases")
 
-			cveRepository := repositories.NewCVERepository(database)
-			cweRepository := repositories.NewCWERepository(database)
-			affectedCmpRepository := repositories.NewAffectedComponentRepository(database)
+			cveRepository := repositories.NewCVERepository(db)
+			cweRepository := repositories.NewCWERepository(db)
+			affectedCmpRepository := repositories.NewAffectedComponentRepository(db)
 			nvdService := vulndb.NewNVDService(cveRepository)
 			mitreService := vulndb.NewMitreService(cweRepository)
 			epssService := vulndb.NewEPSSService(nvdService, cveRepository)
 			osvService := vulndb.NewOSVService(affectedCmpRepository)
 			// cvelistService := vulndb.NewCVEListService(cveRepository)
 			debianSecurityTracker := vulndb.NewDebianSecurityTracker(affectedCmpRepository)
 
-			expoitDBService := vulndb.NewExploitDBService(nvdService, repositories.NewExploitRepository(database))
+			expoitDBService := vulndb.NewExploitDBService(nvdService, repositories.NewExploitRepository(db))
 
-			githubExploitDBService := vulndb.NewGithubExploitDBService(repositories.NewExploitRepository(database))
+			githubExploitDBService := vulndb.NewGithubExploitDBService(repositories.NewExploitRepository(db))
 
 			if emptyOrContains(databasesToSync, "cwe") {
 				now := time.Now()

cmd/devguard/main.go

@@ -72,8 +72,8 @@ func main() {
 	}
 
 	// Run database migrations using the existing database connection
-	autoMigrate := os.Getenv("AUTO_MIGRATE")
-	if autoMigrate == "" || autoMigrate == "true" {
+	disableAutoMigrate := os.Getenv("DISABLE_AUTOMIGRATE")
+	if disableAutoMigrate != "true" {
 		slog.Info("running database migrations...")
 		if err := database.RunMigrationsWithDB(db); err != nil {
 			slog.Error("failed to run database migrations", "error", err)
@@ -86,7 +86,7 @@ func main() {
 			panic(errors.New("Failed to run hash migrations"))
 		}
 	} else {
-		slog.Info("automatic migrations disabled via AUTO_MIGRATE=false")
+		slog.Info("automatic migrations disabled via DISABLE_AUTOMIGRATE=true")
 	}
 
 	daemon.Start(db)