implements partial imports, adds logs to health endpoint

Author: timbastin

cmd/devguard-cli/commands/vulndb_import.go

@@ -17,6 +17,7 @@ func newImportCommand() *cobra.Command {
 	var full bool
 	var batchSize int
 	var bulk bool
+	var limitedToTables []string
 
 	importCmd := &cobra.Command{
 		Use:   "import",
@@ -26,9 +27,10 @@ func newImportCommand() *cobra.Command {
 			shared.LoadConfig() // nolint
 			migrateDB()
 			opts := shared.ImportOptions{
-				Full:      full,
-				BatchSize: batchSize,
-				Bulk:      bulk,
+				Full:            full,
+				BatchSize:       batchSize,
+				Bulk:            bulk,
+				LimitedToTables: limitedToTables,
 			}
 			app := fx.New(
 				fx.NopLogger,
@@ -58,6 +60,7 @@ func newImportCommand() *cobra.Command {
 	importCmd.Flags().BoolVar(&full, "full", false, "Force a full import, ignoring the last-import watermark")
 	importCmd.Flags().IntVar(&batchSize, "batchSize", 5000, "Number of OSV entries per batch (default 5000)")
 	importCmd.Flags().BoolVar(&bulk, "bulk", false, "Load all gob data into RAM before writing (faster but uses ~2-3 GB memory)")
+	importCmd.Flags().StringSliceVar(&limitedToTables, "limitedToTables", []string{}, "Comma-separated list of tables to limit the import to (e.g. --limitedToTables=cves,exploits,malicious_packages)")
 
 	return importCmd
 }

router/apiv1_router.go

@@ -1,7 +1,9 @@
 package router
 
 import (
+	"context"
 	"encoding/json"
+	"log/slog"
 	"os"
 	"runtime"
 	"time"
@@ -179,13 +181,18 @@ func NewAPIV1Router(srv api.Server,
 		// Check database connectivity
 		sqlDB, err := db.DB()
 		if err != nil {
+			slog.Info("failed to get database instance", "error", err)
 			return ctx.JSON(503, map[string]string{
 				"status": "unhealthy",
 				"error":  "failed to get database instance",
 			})
 		}
 
-		if err := sqlDB.Ping(); err != nil {
+		ctxWithTimeout, cancel := context.WithTimeout(ctx.Request().Context(), 5*time.Second)
+		defer cancel()
+
+		if err := sqlDB.PingContext(ctxWithTimeout); err != nil {
+			slog.Info("database ping failed", "error", err)
 			return ctx.JSON(503, map[string]string{
 				"status": "unhealthy",
 				"error":  "database ping failed",

shared/common_interfaces.go

@@ -633,7 +633,8 @@ type ImportOptions struct {
 	BatchSize int
 	// Bulk loads all gob data into RAM before writing — no channels, single DB flush.
 	// Faster than streaming but uses significantly more memory (~2-3 GB).
-	Bulk bool
+	Bulk            bool
+	LimitedToTables []string
 }
 
 type VulnDBService interface {

utils/slice.go

@@ -119,12 +119,7 @@ func CompareSlices[T any, K comparable](a, b []T, serializer func(T) K) CompareR
 }
 
 func Any[T any](s []T, f func(T) bool) bool {
-	for _, v := range s {
-		if f(v) {
-			return true
-		}
-	}
-	return false
+	return slices.ContainsFunc(s, f)
 }
 
 func All[T any](s []T, f func(T) bool) bool {
@@ -157,3 +152,9 @@ func ContainsAll[T comparable](s []T, needed []T) bool {
 		return Contains(s, n)
 	})
 }
+
+func ContainsAny[T comparable](s []T, needed []T) bool {
+	return Any(needed, func(n T) bool {
+		return Contains(s, n)
+	})
+}

vulndb/integrity.go

@@ -103,31 +103,33 @@ type integrityInformation struct {
 	ImportTimestamp time.Time                   `json:"import_timestamp"`
 }
 
-func validateIntegrityInformation(workingDir string, groundTruth integrityInformation, localIntegrityInformation []tableIntegrityInformation) error {
-	didErr := false
+// returns a string slice with failing tables
+// if nil, then all tables are valid
+func validateIntegrityInformation(workingDir string, groundTruth integrityInformation, localIntegrityInformation []tableIntegrityInformation) ([]string, bool) {
+	failingTables := make([]string, 0)
 	for _, tableIntegrity := range localIntegrityInformation {
 		found := false
 		for _, tableGroundTruth := range groundTruth.TableIntegrity {
 			if tableGroundTruth.TableName == tableIntegrity.TableName {
 				if !tableIntegrity.isEqual(tableGroundTruth) {
 					slog.Error("invalid checksum when importing", "table", tableIntegrity.TableName, "expectedCount", tableGroundTruth.TotalCount, "actualCount", tableIntegrity.TotalCount, "expectedChecksum", fmt.Sprintf("%x", tableGroundTruth.Checksum), "actualChecksum", fmt.Sprintf("%x", tableIntegrity.Checksum))
-
-					didErr = true
+					failingTables = append(failingTables, tableIntegrity.TableName)
 				} else {
 					found = true
 					break
 				}
 			}
 		}
 		if !found {
-			return fmt.Errorf("could not find integrity information for table %s", tableIntegrity.TableName)
+			slog.Error("unexpected table found when importing", "table", tableIntegrity.TableName, "count", tableIntegrity.TotalCount, "checksum", fmt.Sprintf("%x", tableIntegrity.Checksum))
+			failingTables = append(failingTables, tableIntegrity.TableName)
 		}
 	}
-	if didErr {
-		return fmt.Errorf("integrity validation failed for one or more tables when importing from %s", workingDir)
+	if len(failingTables) > 0 {
+		return failingTables, false
 	}
 
-	return nil
+	return nil, true
 }
 
 func calculateTotalIntegrityInformation(ctx context.Context, tx pgx.Tx) ([]tableIntegrityInformation, error) {

vulndb/osv_service.go

@@ -172,7 +172,8 @@ func (s osvService) fetchAndImportOSV(ctx context.Context, tx pgx.Tx, importStar
 	if err := createStagingTables(ctx, tx); err != nil {
 		return nil, nil, fmt.Errorf("could not create staging tables: %w", err)
 	}
-	vulnRows, malRows := gobOSVToVulnAndMalFilterTransformer(ctx, time.Time{}, nil)(allOSVVulns)
+	malRows := gobOSVToMalFilterTransformer(time.Time{})(allOSVVulns)
+	vulnRows := gobOSVToVulnFilterTransformer(time.Time{}, nil)(allOSVVulns)
 	fakeRows, fakeComps := buildFakePackages()
 	malRows.pkgs = append(malRows.pkgs, fakeRows...)
 	malRows.comps = append(malRows.comps, fakeComps...)

vulndb/transformers.go

@@ -55,17 +55,39 @@ func osvEntryToMaliciousPackageTransformer(entry *dtos.OSV) (models.MaliciousPac
 	return pkg, components
 }
 
-func gobOSVToVulnAndMalFilterTransformer(ctx context.Context, lastImportTime time.Time, existing map[int64][]int64) func([]OSVEntry) (vulndbRows, malRows) {
+func gobOSVToMalFilterTransformer(lastImportTime time.Time) func([]OSVEntry) malRows {
+	return func(elements []OSVEntry) malRows {
+		malPkgs := make([]models.MaliciousPackage, 0)
+		malComps := make([]models.MaliciousAffectedComponent, 0)
+
+		for i := range elements {
+			if !lastImportTime.IsZero() && !elements[i].ModifiedTimestamp.After(lastImportTime) {
+				continue
+			}
+
+			// check if malicious package or vulnerability
+			if strings.HasPrefix(elements[i].OSV.ID, "MAL-") {
+				pkg, comps := osvEntryToMaliciousPackageTransformer(elements[i].OSV)
+				malPkgs = append(malPkgs, pkg)
+				malComps = append(malComps, comps...)
+				continue
+			}
+		}
+		return malRows{
+			pkgs:  malPkgs,
+			comps: malComps,
+		}
+	}
+}
+func gobOSVToVulnFilterTransformer(lastImportTime time.Time, existing map[int64][]int64) func([]OSVEntry) vulndbRows {
 	if existing == nil {
 		existing = make(map[int64][]int64)
 	}
-	return func(elements []OSVEntry) (vulndbRows, malRows) {
+	return func(elements []OSVEntry) vulndbRows {
 		cves := make([]models.CVE, 0, len(elements))
 		cveRelationships := make([]models.CVERelationship, 0, len(elements)*2)
 		affectedComponents := make([]models.AffectedComponent, 0, len(elements)*12)
 		cveAffectedComponents := make([]cveAffectedComponentRow, 0, len(elements)*55)
-		malPkgs := make([]models.MaliciousPackage, 0)
-		malComps := make([]models.MaliciousAffectedComponent, 0)
 
 		for i := range elements {
 			if !lastImportTime.IsZero() && !elements[i].ModifiedTimestamp.After(lastImportTime) {
@@ -74,9 +96,6 @@ func gobOSVToVulnAndMalFilterTransformer(ctx context.Context, lastImportTime tim
 
 			// check if malicious package or vulnerability
 			if strings.HasPrefix(elements[i].OSV.ID, "MAL-") {
-				pkg, comps := osvEntryToMaliciousPackageTransformer(elements[i].OSV)
-				malPkgs = append(malPkgs, pkg)
-				malComps = append(malComps, comps...)
 				continue
 			}
 			relationships := transformer.OSVToCVERelationships(elements[i].OSV)
@@ -107,28 +126,33 @@ func gobOSVToVulnAndMalFilterTransformer(ctx context.Context, lastImportTime tim
 			}
 		}
 		return vulndbRows{
-				CVEs:                  cves,
-				CVERelationships:      cveRelationships,
-				AffectedComponents:    affectedComponents,
-				CVEAffectedComponents: cveAffectedComponents,
-			}, malRows{
-				pkgs:  malPkgs,
-				comps: malComps,
-			}
+			CVEs:                  cves,
+			CVERelationships:      cveRelationships,
+			AffectedComponents:    affectedComponents,
+			CVEAffectedComponents: cveAffectedComponents,
+		}
 	}
 }
 
-func gobOSVEntryStreamer(ctx context.Context, lastImportTime time.Time, existing map[int64][]int64, vulndbChan chan<- vulndbRows, malPkgsChan chan<- malRows) func([]OSVEntry) error {
-	transform := gobOSVToVulnAndMalFilterTransformer(ctx, lastImportTime, existing)
+func gobOSVStreamer(ctx context.Context, lastImportTime time.Time, existing map[int64][]int64, vulndbChan chan<- vulndbRows) func([]OSVEntry) error {
+	transform := gobOSVToVulnFilterTransformer(lastImportTime, existing)
 	return func(elements []OSVEntry) error {
-		vulndbRows, malRows := transform(elements)
+		vulndbRows := transform(elements)
 		select {
-		case malPkgsChan <- malRows:
+		case vulndbChan <- vulndbRows:
 		case <-ctx.Done():
 			return ctx.Err()
 		}
+		return nil
+	}
+}
+
+func gobOSVMalPkgStreamer(ctx context.Context, lastImportTime time.Time, malPkgsChan chan<- malRows) func([]OSVEntry) error {
+	transform := gobOSVToMalFilterTransformer(lastImportTime)
+	return func(elements []OSVEntry) error {
+		malRows := transform(elements)
 		select {
-		case vulndbChan <- vulndbRows:
+		case malPkgsChan <- malRows:
 		case <-ctx.Done():
 			return ctx.Err()
 		}