Correctly removing from list

Signed-off-by: nathannaveen <[email protected]>

pkg/assembler/backends/keyvalue/certifyVuln.go

@@ -18,7 +18,6 @@ package keyvalue
 import (
 	"context"
 	"errors"
-	"fmt"
 	"reflect"
 	"sort"
 	"strings"
@@ -31,12 +30,6 @@ import (
 	"github.com/guacsec/guac/pkg/assembler/kv"
 )
 
-const (
-	certifyVulnLinkType = "certifyVuln"
-	hasSLSALinkType     = "hasSLSA"
-	hasSBOMLinkType     = "hasSBOM"
-)
-
 // Internal data: link between packages and vulnerabilities (certifyVulnerability)
 type certifyVulnerabilityLink struct {
 	ThisID          string
@@ -68,63 +61,15 @@ func (n *certifyVulnerabilityLink) Key() string {
 	}, ":"))
 }
 
-// Helper function to remove vulnerability links. This works by setting all the links expect the specified linkID.
-func (c *demoClient) removeLinks(ctx context.Context, linkID string, linkType string, links []string, col string, id string) error {
+// removeLinkFromList is a helper function to remove a link from an array of links. This works by setting all the links except the specified linkID.
+func removeLinkFromList(linkID string, links []string) []string {
 	var newLinks []string
 	for _, link := range links {
 		if link != linkID {
 			newLinks = append(newLinks, link)
 		}
 	}
-
-	switch col {
-	case "packages":
-		var pkg pkgVersion
-		if err := c.kv.Get(ctx, col, id, &pkg); err != nil {
-			return fmt.Errorf("error getting package version from keyvalue: %w", err)
-		}
-		switch linkType {
-		case certifyVulnLinkType:
-			pkg.CertifyVulnLinks = newLinks
-		case hasSBOMLinkType:
-			pkg.HasSBOMs = newLinks
-		}
-		return setkv(ctx, col, &pkg, c)
-	case "vulnerabilities":
-		var vuln vulnTypeStruct
-		if err := c.kv.Get(ctx, col, id, &vuln); err != nil {
-			return fmt.Errorf("error getting vulnerability from keyvalue: %w", err)
-		}
-		switch linkType {
-		case certifyVulnLinkType:
-			vuln.VulnIDs = newLinks
-		}
-		return setkv(ctx, col, &vuln, c)
-	case "builders":
-		var builder builderStruct
-		if err := c.kv.Get(ctx, col, id, &builder); err != nil {
-			return fmt.Errorf("error getting builder from keyvalue: %w", err)
-		}
-		switch linkType {
-		case hasSLSALinkType:
-			builder.HasSLSAs = newLinks
-		}
-		return setkv(ctx, col, &builder, c)
-	case "artifacts":
-		var artifact artStruct
-		if err := c.kv.Get(ctx, col, id, &artifact); err != nil {
-			return fmt.Errorf("error getting artifact from keyvalue: %w", err)
-		}
-		switch linkType {
-		case hasSBOMLinkType:
-			artifact.HasSBOMs = newLinks
-		case hasSLSALinkType:
-			artifact.HasSLSAs = newLinks
-		}
-		return setkv(ctx, col, &artifact, c)
-	default:
-		return errors.New("unsupported entity type")
-	}
+	return newLinks
 }
 
 // DeleteCertifyVuln deletes a specified certifyVuln node along with all associated relationships.
@@ -137,23 +82,29 @@ func (c *demoClient) DeleteCertifyVuln(ctx context.Context, id string) (bool, er
 		if errors.Is(err, kv.NotFoundError) {
 			return false, nil // Not found, nothing to delete
 		}
-		return false, gqlerror.Errorf("%v :: %s", funcName, err) // TODO: Improve error messages
+		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
 	// Remove backlinks from associated package and vulnerability
-	foundPackage, err := c.returnFoundPkgVersion(ctx, &model.IDorPkgInput{PackageVersionID: &link.PackageID})
+	foundPkgNode, err := c.returnFoundPkgVersion(ctx, &model.IDorPkgInput{PackageVersionID: &link.PackageID})
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, certifyVulnLinkType, foundPackage.CertifyVulnLinks, "packages", foundPackage.ID()); err != nil {
+
+	foundPkgNode.CertifyVulnLinks = removeLinkFromList(link.ThisID, foundPkgNode.CertifyVulnLinks)
+	err = setkv(ctx, pkgVerCol, foundPkgNode, c)
+	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
 	foundVulnNode, err := c.returnFoundVulnerability(ctx, &model.IDorVulnerabilityInput{VulnerabilityNodeID: &link.VulnerabilityID})
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, certifyVulnLinkType, foundVulnNode.CertifyVulnLinks, "vulnerabilities", foundVulnNode.ID()); err != nil {
+
+	foundVulnNode.CertifyVulnLinks = removeLinkFromList(link.ThisID, foundPkgNode.CertifyVulnLinks)
+	err = setkv(ctx, vulnIDCol, foundPkgNode, c)
+	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 

pkg/assembler/backends/keyvalue/hasSBOM.go

@@ -100,15 +100,21 @@ func (c *demoClient) DeleteHasSBOM(ctx context.Context, id string) (bool, error)
 		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
-		if err := c.removeLinks(ctx, link.ThisID, hasSBOMLinkType, foundPkg.HasSBOMs, "packages", foundPkg.ID()); err != nil {
+
+		foundPkg.CertifyVulnLinks = removeLinkFromList(link.ThisID, foundPkg.HasSBOMs)
+		err = setkv(ctx, pkgVerCol, foundPkg, c)
+		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
 	} else if link.Artifact != "" {
 		foundArtifact, err := c.returnFoundArtifact(ctx, &model.IDorArtifactInput{ArtifactID: &link.Artifact})
 		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
-		if err := c.removeLinks(ctx, link.ThisID, hasSBOMLinkType, foundArtifact.HasSBOMs, "artifacts", foundArtifact.ID()); err != nil {
+
+		foundArtifact.HasSBOMs = removeLinkFromList(link.ThisID, foundArtifact.HasSBOMs)
+		err = setkv(ctx, artCol, foundArtifact, c)
+		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
 	}

pkg/assembler/backends/keyvalue/hasSLSA.go

@@ -90,7 +90,10 @@ func (c *demoClient) DeleteHasSLSA(ctx context.Context, id string) (bool, error)
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, hasSLSALinkType, foundSubject.HasSLSAs, "artifacts", foundSubject.ID()); err != nil {
+
+	foundSubject.HasSLSAs = removeLinkFromList(link.ThisID, foundSubject.HasSLSAs)
+	err = setkv(ctx, artCol, foundSubject, c)
+	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
@@ -99,7 +102,10 @@ func (c *demoClient) DeleteHasSLSA(ctx context.Context, id string) (bool, error)
 	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
-	if err := c.removeLinks(ctx, link.ThisID, hasSLSALinkType, foundBuiltBy.HasSLSAs, "builders", foundBuiltBy.ID()); err != nil {
+
+	foundBuiltBy.HasSLSAs = removeLinkFromList(link.ThisID, foundBuiltBy.HasSLSAs)
+	err = setkv(ctx, builderCol, foundBuiltBy, c)
+	if err != nil {
 		return false, gqlerror.Errorf("%v :: %s", funcName, err)
 	}
 
@@ -109,7 +115,10 @@ func (c *demoClient) DeleteHasSLSA(ctx context.Context, id string) (bool, error)
 		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
-		if err := c.removeLinks(ctx, link.ThisID, hasSLSALinkType, foundBuiltFrom.HasSLSAs, "artifacts", foundBuiltFrom.ID()); err != nil {
+
+		foundBuiltFrom.HasSLSAs = removeLinkFromList(link.ThisID, foundBuiltFrom.HasSLSAs)
+		err = setkv(ctx, artCol, foundBuiltFrom, c)
+		if err != nil {
 			return false, gqlerror.Errorf("%v :: %s", funcName, err)
 		}
 	}