refactor: register workload with upsert instead of crud

* add constraint on name, type, ns and cluster in workload table
* implement management in client

Co-authored-by: ybelmekk <[email protected]>

Author: tommytroen

cmd/api/main.go

@@ -3,7 +3,9 @@ package main
 import (
 	"context"
 	"fmt"
+	"github.com/nais/v13s/internal/api/grpcmgmt"
 	"github.com/nais/v13s/internal/api/grpcvulnerabilities"
+	"github.com/nais/v13s/pkg/api/vulnerabilities/management"
 	"net"
 	"os"
 	"os/signal"
@@ -51,7 +53,8 @@ func main() {
 		log.Fatalf("Failed to create DependencyTrack client: %v", err)
 	}
 
-	vulnerabilities.RegisterVulnerabilitiesServer(grpcServer, &grpcvulnerabilities.Server{DpClient: dpClient, Db: db})
+	vulnerabilities.RegisterVulnerabilitiesServer(grpcServer, &grpcvulnerabilities.Server{Db: db})
+	management.RegisterManagementServer(grpcServer, &grpcmgmt.Server{Db: db, DpClient: dpClient})
 
 	stop := make(chan os.Signal, 1)
 	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)

cmd/client/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"context"
 	"fmt"
+	"github.com/nais/v13s/pkg/api/vulnerabilities/management"
 
 	"github.com/nais/v13s/pkg/api/vulnerabilities"
 	"google.golang.org/grpc"
@@ -18,6 +19,8 @@ func main() {
 		panic(err)
 	}
 
+	workloadManagement(c)
+
 	resp, err := c.ListVulnerabilitySummaries(
 		context.Background(),
 		vulnerabilities.ClusterFilter("prod-gcp"),
@@ -60,6 +63,30 @@ func main() {
 	fmt.Printf("summary: %v\n", resp2.VulnerabilitySummary)
 }
 
+func workloadManagement(c vulnerabilities.Client) {
+	ctx := context.Background()
+	_, err := c.RegisterWorkload(
+		ctx,
+		&management.RegisterWorkloadRequest{
+			Cluster:      "dev-fss",
+			Namespace:    "nais-system",
+			Workload:     "aivenator",
+			WorkloadType: "app",
+			ImageName:    "europe-north1-docker.pkg.dev/nais-io/nais/images/aivenator",
+			ImageTag:     "2025-01-22-124805-9221180",
+			Metadata: &management.Metadata{
+				Labels: map[string]string{
+					"workflow": "deploy",
+				},
+			},
+		},
+	)
+	if err != nil {
+		panic(err)
+	}
+
+}
+
 func handle(resp *vulnerabilities.ListVulnerabilitySummariesResponse, err error) {
 	if err != nil {
 		panic(err)

internal/api/grpcmgmt/server.go

@@ -2,6 +2,8 @@ package grpcmgmt
 
 import (
 	"context"
+	"errors"
+	"github.com/jackc/pgx/v5"
 	"github.com/nais/v13s/internal/database/sql"
 	"github.com/nais/v13s/internal/dependencytrack"
 	"github.com/nais/v13s/pkg/api/vulnerabilities/management"
@@ -15,24 +17,31 @@ type Server struct {
 	Db       *sql.Queries
 }
 
-func (s *Server) CreateWorkload(ctx context.Context, request *management.CreateWorkloadRequest) (*management.CreateWorkloadResponse, error) {
+func (s *Server) RegisterWorkload(ctx context.Context, request *management.RegisterWorkloadRequest) (*management.RegisterWorkloadResponse, error) {
 	metadata := map[string]string{}
 	if request.Metadata != nil {
 		metadata = request.Metadata.Labels
 	}
 
-	imageParams := sql.CreateImageParams{
-		Name:     request.ImageName,
-		Tag:      request.ImageTag,
-		Metadata: metadata,
-	}
-
-	_, err := s.Db.CreateImage(ctx, imageParams)
-	if err != nil {
+	_, err := s.Db.GetImage(ctx, sql.GetImageParams{
+		Name: request.ImageName,
+		Tag:  request.ImageTag,
+	})
+
+	if errors.Is(err, pgx.ErrNoRows) {
+		_, err = s.Db.CreateImage(ctx, sql.CreateImageParams{
+			Name:     request.ImageName,
+			Tag:      request.ImageTag,
+			Metadata: metadata,
+		})
+		if err != nil {
+			return nil, err
+		}
+	} else if err != nil {
 		return nil, err
 	}
 
-	w := sql.CreateWorkloadParams{
+	w := sql.UpsertWorkloadParams{
 		Name:         request.Workload,
 		WorkloadType: request.WorkloadType,
 		Namespace:    request.Namespace,
@@ -41,7 +50,7 @@ func (s *Server) CreateWorkload(ctx context.Context, request *management.CreateW
 		ImageTag:     request.ImageTag,
 	}
 
-	_, err = s.Db.CreateWorkload(ctx, w)
+	err = s.Db.UpsertWorkload(ctx, w)
 	if err != nil {
 		return nil, err
 	}
@@ -51,8 +60,8 @@ func (s *Server) CreateWorkload(ctx context.Context, request *management.CreateW
 		return nil, err
 	}
 
-	response := &management.CreateWorkloadResponse{}
-	if p.Metrics == nil {
+	response := &management.RegisterWorkloadResponse{}
+	if p == nil || p.Metrics == nil {
 		return response, nil
 	}
 
@@ -77,8 +86,3 @@ func (s *Server) CreateWorkload(ctx context.Context, request *management.CreateW
 
 	return response, err
 }
-
-func (s *Server) UpdateWorkload(ctx context.Context, request *management.UpdateWorkloadRequest) (*management.UpdateWorkloadResponse, error) {
-	//TODO implement me
-	panic("implement me")
-}

internal/api/grpcvulnerabilities/server.go

@@ -70,6 +70,8 @@ func (s *Server) ListVulnerabilitySummaries(ctx context.Context, request *vulner
 	return response, nil
 }
 
+// TODO: if no summaries are found, handle this case by not returning the summary? and maybe handle it in the sql query, right now we return 0 on all fields
+// TLDR: make distinction between no summary found and summary found with 0 values
 func (s *Server) GetVulnerabilitySummary(ctx context.Context, request *vulnerabilities.GetVulnerabilitySummaryRequest) (*vulnerabilities.GetVulnerabilitySummaryResponse, error) {
 	sum, err := s.Db.GetVulnerabilitySummary(ctx, sql.GetVulnerabilitySummaryParams{
 		Cluster:      request.Filter.Cluster,

internal/database/migrations/0001_schema.sql

@@ -26,6 +26,7 @@ CREATE TABLE workloads
     workload_type TEXT                                               NOT NULL,
     namespace     TEXT                                               NOT NULL,
     cluster       TEXT                                               NOT NULL,
+    CONSTRAINT workload_id UNIQUE (name, workload_type, namespace, cluster),
     image_name    TEXT                                               NOT NULL,
     image_tag     TEXT                                               NOT NULL,
     created_at    TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP NOT NULL,

internal/database/queries/image.sql

@@ -7,3 +7,6 @@ RETURNING
         *
 ;
 
+-- name: GetImage :one
+SELECT * FROM images WHERE name = @name AND tag = @tag;
+

internal/database/queries/workloads.sql

@@ -1,3 +1,27 @@
+-- name: UpsertWorkload :exec
+INSERT INTO workloads(
+    name,
+    workload_type,
+    namespace,
+    cluster,
+    image_name,
+    image_tag
+)
+VALUES (
+    @name,
+    @workload_type,
+    @namespace,
+    @cluster,
+    @image_name,
+    @image_tag
+) ON CONFLICT
+    ON CONSTRAINT workload_id DO
+        UPDATE
+    SET
+        image_name = @image_name,
+        image_tag = @image_tag
+;
+
 -- name: CreateWorkload :one
 INSERT INTO
     workloads (name, workload_type, namespace, cluster, image_name, image_tag)
@@ -10,14 +34,13 @@ RETURNING
 -- name: UpdateWorkload :one
 UPDATE workloads
 SET
-    name = COALESCE(sqlc.narg(name), name),
-    workload_type = COALESCE(sqlc.narg(workload_type), workload_type),
-    namespace = COALESCE(sqlc.narg(namespace), namespace),
-    cluster = COALESCE(sqlc.narg(cluster), cluster),
-    image_name = COALESCE(sqlc.narg(image_name), image_name),
-    image_tag = COALESCE(sqlc.narg(image_tag), image_tag)
+    image_name = @image_name,
+    image_tag = @image_tag
 WHERE
-    workloads.id = @id
+    cluster = @cluster AND
+    namespace = @namespace AND
+    workload_type = @workload_type AND
+    name = @name
 RETURNING
     *
 ;