Add mitigation with task and mapping import

Signed-off-by: gearnode <bryan@frimin.fr>

pkg/coredata/control.go

@@ -159,6 +159,50 @@ WHERE
 	return nil
 }
 
+func (c *Control) LoadByFrameworkIDAndReferenceID(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	frameworkID gid.GID,
+	referenceID string,
+) error {
+	q := `
+SELECT
+    id,
+    reference_id,
+    framework_id,
+    tenant_id,
+    name,
+    description,
+    created_at,
+    updated_at
+FROM
+    controls
+WHERE
+    %s
+    AND framework_id = @framework_id
+    AND reference_id = @reference_id
+LIMIT 1;
+`
+	q = fmt.Sprintf(q, scope.SQLFragment())
+
+	args := pgx.StrictNamedArgs{"framework_id": frameworkID, "reference_id": referenceID}
+	maps.Copy(args, scope.SQLArguments())
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot query controls: %w", err)
+	}
+
+	control, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Control])
+	if err != nil {
+		return fmt.Errorf("cannot collect control: %w", err)
+	}
+
+	*c = control
+
+	return nil
+}
+
 func (c *Control) LoadByID(
 	ctx context.Context,
 	conn pg.Conn,

pkg/coredata/control_mitigation.go

@@ -36,6 +36,38 @@ type (
 	ControlMitigations []*ControlMitigation
 )
 
+func (cm ControlMitigation) Upsert(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+) error {
+	q := `
+INSERT INTO
+    controls_mitigations (
+        control_id,
+        mitigation_id,
+        tenant_id,
+        created_at
+    )
+VALUES (
+    @control_id,
+    @mitigation_id,
+    @tenant_id,
+    @created_at
+)
+ON CONFLICT (control_id, mitigation_id) DO NOTHING;
+`
+
+	args := pgx.StrictNamedArgs{
+		"control_id":    cm.ControlID,
+		"mitigation_id": cm.MitigationID,
+		"tenant_id":     scope.GetTenantID(),
+		"created_at":    cm.CreatedAt,
+	}
+	_, err := conn.Exec(ctx, q, args)
+	return err
+}
+
 func (cm ControlMitigation) Insert(
 	ctx context.Context,
 	conn pg.Conn,

pkg/coredata/framework.go

@@ -104,6 +104,48 @@ WHERE
 	return nil
 }
 
+func (f *Framework) LoadByReferenceID(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+	referenceID string,
+) error {
+	q := `
+SELECT
+    id,
+    organization_id,
+    reference_id,
+    name,
+    description,
+    created_at,
+    updated_at
+FROM
+    frameworks
+WHERE
+    %s
+    AND reference_id = @reference_id
+LIMIT 1;
+`
+
+	q = fmt.Sprintf(q, scope.SQLFragment())
+
+	args := pgx.StrictNamedArgs{"reference_id": referenceID}
+	maps.Copy(args, scope.SQLArguments())
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot query frameworks: %w", err)
+	}
+
+	framework, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Framework])
+	if err != nil {
+		return fmt.Errorf("cannot collect framework: %w", err)
+	}
+
+	*f = framework
+
+	return nil
+}
+
 func (f *Framework) LoadByID(
 	ctx context.Context,
 	conn pg.Conn,

pkg/coredata/migrations/20250401T144700Z.sql

@@ -0,0 +1,5 @@
+ALTER TABLE tasks ADD COLUMN reference_id TEXT DEFAULT '';
+ALTER TABLE tasks ALTER COLUMN reference_id SET NOT NULL;
+ALTER TABLE tasks ALTER COLUMN reference_id DROP DEFAULT;
+
+ALTER TABLE tasks ADD CONSTRAINT tasks_reference_id_unique UNIQUE (id, reference_id);

pkg/coredata/migrations/20250401T165000Z.sql

@@ -0,0 +1,3 @@
+ALTER TABLE mitigations RENAME COLUMN content_ref TO reference_id;
+ALTER TABLE mitigations ADD CONSTRAINT mitigations_org_ref_unique UNIQUE (organization_id, reference_id);
+ALTER TABLE mitigations DROP COLUMN standards;

pkg/coredata/migrations/20250401T171500Z.sql

@@ -0,0 +1 @@
+ALTER TABLE mitigations DROP COLUMN version;

pkg/coredata/migrations/20250401T173800Z.sql

@@ -0,0 +1,11 @@
+ALTER TABLE controls_mitigations
+    ADD CONSTRAINT fk_controls_mitigations_control_id
+    FOREIGN KEY (control_id)
+    REFERENCES controls(id)
+    ON DELETE CASCADE;
+
+ALTER TABLE controls_mitigations
+    ADD CONSTRAINT fk_controls_mitigations_mitigation_id
+    FOREIGN KEY (mitigation_id)
+    REFERENCES mitigations(id)
+    ON DELETE CASCADE;

pkg/coredata/migrations/20250401T175300Z.sql

@@ -0,0 +1,2 @@
+ALTER TABLE tasks DROP CONSTRAINT tasks_reference_id_unique;
+ALTER TABLE tasks ADD CONSTRAINT tasks_reference_id_unique UNIQUE (mitigation_id, reference_id);

pkg/coredata/migrations/20250401T175400Z.sql

@@ -0,0 +1,3 @@
+ALTER TABLE controls
+  DROP CONSTRAINT IF EXISTS controls_framework_ref_unique,
+  ADD CONSTRAINT controls_framework_ref_unique UNIQUE (framework_id, reference_id);

pkg/coredata/mitigation.go

@@ -37,10 +37,9 @@ type (
 		Description    string               `db:"description"`
 		Importance     MitigationImportance `db:"importance"`
 		State          MitigationState      `db:"state"`
-		ContentRef     string               `db:"content_ref"`
+		ReferenceID    string               `db:"reference_id"`
 		CreatedAt      time.Time            `db:"created_at"`
 		UpdatedAt      time.Time            `db:"updated_at"`
-		Standards      []string             `db:"standards"`
 	}
 
 	Mitigations []*Mitigation
@@ -73,10 +72,9 @@ WITH mtgtns AS (
 		m.description,
 		m.state,
 		m.importance,
-		m.content_ref,
+		m.reference_id,
 		m.created_at,
-		m.updated_at,
+		m.updated_at
-		m.standards
 	FROM
 		mitigations m
 	INNER JOIN
@@ -93,10 +91,9 @@ SELECT
 	description,
 	state,
 	importance,
-	content_ref,
+	reference_id,
 	created_at,
-	updated_at,
+	updated_at
-	standards
 FROM
 	mtgtns
 WHERE %s
@@ -141,10 +138,9 @@ WITH mtgtns AS (
 		m.description,
 		m.state,
 		m.importance,
-		m.content_ref,
+		m.reference_id,
 		m.created_at,
-		m.updated_at,
+		m.updated_at
-		m.standards
 	FROM
 		mitigations m
 	INNER JOIN
@@ -161,10 +157,9 @@ SELECT
 	description,
 	state,
 	importance,
-	content_ref,
+	reference_id,
 	created_at,
-	updated_at,
+	updated_at
-	standards
 FROM
 	mtgtns
 WHERE %s
@@ -208,10 +203,9 @@ SELECT
     description,
     state,
 	importance,
-    content_ref,
+    reference_id,
     created_at,
-    updated_at,
+    updated_at
-	standards
 FROM
     mitigations
 WHERE
@@ -256,10 +250,9 @@ SELECT
     description,
     state,
 	importance,
-    content_ref,
+    reference_id,
     created_at,
-    updated_at,
+    updated_at
-	standards
 FROM
     mitigations
 WHERE
@@ -288,6 +281,87 @@ LIMIT 1;
 	return nil
 }
 
+func (c *Mitigation) Upsert(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+) error {
+	q := `
+INSERT INTO
+    mitigations (
+        tenant_id,
+        id,
+        organization_id,
+		category,
+        name,
+		importance,
+		state,
+        description,
+        reference_id,
+        created_at,
+        updated_at
+	)
+VALUES (
+    @tenant_id,
+    @mitigation_id,
+    @organization_id,
+	@category,
+    @name,
+	@importance,
+	@state,
+    @description,
+    @reference_id,
+    @created_at,
+    @updated_at
+)
+ON CONFLICT (organization_id, reference_id) DO UPDATE SET
+    name = @name,
+    description = @description,
+    category = @category,
+    updated_at = @updated_at
+RETURNING
+	tenant_id,
+    id,
+    organization_id,
+	category,
+    name,
+	importance,
+	state,
+    description,
+	reference_id,
+    created_at,
+    updated_at
+`
+
+	args := pgx.StrictNamedArgs{
+		"tenant_id":       scope.GetTenantID(),
+		"mitigation_id":   c.ID,
+		"organization_id": c.OrganizationID,
+		"category":        c.Category,
+		"name":            c.Name,
+		"importance":      c.Importance,
+		"state":           c.State,
+		"description":     c.Description,
+		"reference_id":    c.ReferenceID,
+		"created_at":      c.CreatedAt,
+		"updated_at":      c.UpdatedAt,
+	}
+
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot query mitigations: %w", err)
+	}
+
+	mitigation, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Mitigation])
+	if err != nil {
+		return fmt.Errorf("cannot collect mitigations: %w", err)
+	}
+
+	*c = mitigation
+
+	return nil
+}
+
 func (c Mitigation) Insert(
 	ctx context.Context,
 	conn pg.Conn,
@@ -304,10 +378,9 @@ INSERT INTO
 		importance,
 		state,
         description,
-        content_ref,
+        reference_id,
         created_at,
-        updated_at,
+        updated_at
-		standards
     )
 VALUES (
     @tenant_id,
@@ -318,10 +391,9 @@ VALUES (
 	@importance,
 	@state,
     @description,
-    @content_ref,
+    @reference_id,
     @created_at,
-    @updated_at,
+    @updated_at
-	@standards
 );
 `
 
@@ -332,12 +404,11 @@ VALUES (
 		"category":        c.Category,
 		"name":            c.Name,
 		"description":     c.Description,
-		"content_ref":     c.ContentRef,
+		"reference_id":    c.ReferenceID,
 		"created_at":      c.CreatedAt,
 		"updated_at":      c.UpdatedAt,
 		"state":           c.State,
 		"importance":      c.Importance,
-		"standards":       c.Standards,
 	}
 	_, err := conn.Exec(ctx, q, args)
 	return err

pkg/coredata/task.go

@@ -34,11 +34,11 @@ type (
 		Name         string         `db:"name"`
 		Description  string         `db:"description"`
 		State        TaskState      `db:"state"`
+		ReferenceID  string         `db:"reference_id"`
 		TimeEstimate *time.Duration `db:"time_estimate"`
 		AssignedToID *gid.GID       `db:"assigned_to"`
 		CreatedAt    time.Time      `db:"created_at"`
 		UpdatedAt    time.Time      `db:"updated_at"`
-		Version      int            `db:"version"`
 	}
 
 	Tasks []*Task
@@ -66,11 +66,11 @@ SELECT
     name,
     description,
     state,
+    reference_id,
     time_estimate,
     assigned_to,
     created_at,
-    updated_at,
+    updated_at
-    version
 FROM
     tasks
 WHERE
@@ -112,12 +112,12 @@ INSERT INTO
         mitigation_id,
         name,
         description,
+        reference_id,
         state,
         time_estimate,
         assigned_to,
         created_at,
-        updated_at,
+        updated_at
-        version
     )
 VALUES (
     @tenant_id,
@@ -125,12 +125,12 @@ VALUES (
     @mitigation_id,
     @name,
     @description,
+    @reference_id,
     @state,
     @time_estimate,
     @assigned_to,
     @created_at,
-    @updated_at,
+    @updated_at
-    @version
 );
 `
 
@@ -140,17 +140,95 @@ VALUES (
 		"mitigation_id": c.MitigationID,
 		"name":          c.Name,
 		"description":   c.Description,
+		"reference_id":  c.ReferenceID,
 		"state":         c.State,
 		"time_estimate": c.TimeEstimate,
 		"assigned_to":   c.AssignedToID,
 		"created_at":    c.CreatedAt,
 		"updated_at":    c.UpdatedAt,
-		"version":       0,
 	}
 	_, err := conn.Exec(ctx, q, args)
 	return err
 }
 
+func (c *Task) Upsert(
+	ctx context.Context,
+	conn pg.Conn,
+	scope Scoper,
+) error {
+	q := `
+INSERT INTO
+    tasks (
+        tenant_id,
+        id,
+        mitigation_id,
+        name,
+        description,
+        reference_id,
+        state,
+        time_estimate,
+        assigned_to,
+        created_at,
+        updated_at
+    )
+VALUES (
+    @tenant_id,
+    @task_id,
+    @mitigation_id,
+    @name,
+    @description,
+    @reference_id,
+    @state,
+    @time_estimate,
+    @assigned_to,
+    @created_at,
+    @updated_at
+)
+ON CONFLICT (mitigation_id, reference_id) DO UPDATE SET
+    name = @name,
+    description = @description,
+    updated_at = @updated_at
+RETURNING
+    id,
+    mitigation_id,
+    name,
+    description,
+    reference_id,
+    state,
+    time_estimate,
+    assigned_to,
+    created_at,
+    updated_at
+`
+
+	args := pgx.StrictNamedArgs{
+		"tenant_id":     scope.GetTenantID(),
+		"task_id":       c.ID,
+		"mitigation_id": c.MitigationID,
+		"name":          c.Name,
+		"description":   c.Description,
+		"reference_id":  c.ReferenceID,
+		"state":         c.State,
+		"time_estimate": c.TimeEstimate,
+		"assigned_to":   c.AssignedToID,
+		"created_at":    c.CreatedAt,
+		"updated_at":    c.UpdatedAt,
+	}
+	rows, err := conn.Query(ctx, q, args)
+	if err != nil {
+		return fmt.Errorf("cannot upsert task: %w", err)
+	}
+
+	task, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Task])
+	if err != nil {
+		return fmt.Errorf("cannot collect tasks: %w", err)
+	}
+
+	*c = task
+
+	return nil
+}
+
 func (c *Tasks) LoadByMitigationID(
 	ctx context.Context,
 	conn pg.Conn,
@@ -165,11 +243,11 @@ SELECT
     name,
     description,
     state,
+    reference_id,
     time_estimate,
     assigned_to,
     created_at,
-    updated_at,
+    updated_at
-    version
 FROM
     tasks
 WHERE
@@ -240,8 +318,7 @@ func (c *Task) AssignTo(
 	q := `
 UPDATE tasks SET
     assigned_to = @assigned_to,
-    updated_at = @updated_at,
+    updated_at = @updated_at
-    version = version + 1
 WHERE %s
     AND id = @task_id
 RETURNING 
@@ -249,12 +326,12 @@ RETURNING
     mitigation_id,
     name,
     description,
+    reference_id,
     state,
     time_estimate,
     assigned_to,
     created_at,
-    updated_at,
+    updated_at
-    version
 `
 	q = fmt.Sprintf(q, scope.SQLFragment())
 
@@ -289,8 +366,7 @@ func (c *Task) Unassign(
 	q := `
 UPDATE tasks SET
     assigned_to = NULL,
-    updated_at = @updated_at,
+    updated_at = @updated_at
-    version = version + 1
 WHERE %s
     AND id = @task_id
 RETURNING 
@@ -298,12 +374,12 @@ RETURNING
     mitigation_id,
     name,
     description,
+    reference_id,
     state,
     time_estimate,
     assigned_to,
     created_at,
-    updated_at,
+    updated_at
-    version
 `
 	q = fmt.Sprintf(q, scope.SQLFragment())
 

pkg/probo/mitigation_service.go

@@ -53,6 +53,16 @@ type (
 			Description string                        `json:"description"`
 			Category    string                        `json:"category"`
 			Importance  coredata.MitigationImportance `json:"importance"`
+			ReferenceID string                        `json:"reference-id"`
+			Standards   []struct {
+				Framework string `json:"framework"`
+				Control   string `json:"control"`
+			} `json:"standards"`
+			Tasks []struct {
+				Name        string `json:"name"`
+				Description string `json:"description"`
+				ReferenceID string `json:"reference-id"`
+			} `json:"tasks"`
 		} `json:"mitigations"`
 	}
 )
@@ -124,36 +134,80 @@ func (s MitigationService) Import(
 	organizationID gid.GID,
 	req ImportMitigationRequest,
 ) (*page.Page[*coredata.Mitigation, coredata.MitigationOrderField], error) {
-
 	importedMitigations := coredata.Mitigations{}
-	for _, mitigation := range req.Mitigations {
-		now := time.Now()
-
-		mitigationID, err := gid.NewGID(organizationID.TenantID(), coredata.MitigationEntityType)
-		if err != nil {
-			return nil, fmt.Errorf("cannot create global id: %w", err)
-		}
-
-		importedMitigations = append(importedMitigations, &coredata.Mitigation{
-			ID:             mitigationID,
-			OrganizationID: organizationID,
-			Name:           mitigation.Name,
-			Description:    mitigation.Description,
-			Category:       mitigation.Category,
-			State:          coredata.MitigationStateNotStarted,
-			Standards:      []string{},
-			Importance:     mitigation.Importance,
-			CreatedAt:      now,
-			UpdatedAt:      now,
-		})
-	}
 
 	err := s.svc.pg.WithTx(
 		ctx,
 		func(tx pg.Conn) error {
-			for _, mitigation := range importedMitigations {
+			for i := range req.Mitigations {
-				if err := mitigation.Insert(ctx, tx, s.svc.scope); err != nil {
+				now := time.Now()
-					return fmt.Errorf("cannot insert mitigation: %w", err)
+
+				mitigationID, err := gid.NewGID(organizationID.TenantID(), coredata.MitigationEntityType)
+				if err != nil {
+					return fmt.Errorf("cannot create global id: %w", err)
+				}
+
+				mitigation := &coredata.Mitigation{
+					ID:             mitigationID,
+					OrganizationID: organizationID,
+					Name:           req.Mitigations[i].Name,
+					Description:    req.Mitigations[i].Description,
+					Category:       req.Mitigations[i].Category,
+					State:          coredata.MitigationStateNotStarted,
+					ReferenceID:    req.Mitigations[i].ReferenceID,
+					Importance:     req.Mitigations[i].Importance,
+					CreatedAt:      now,
+					UpdatedAt:      now,
+				}
+
+				importedMitigations = append(importedMitigations, mitigation)
+
+				if err := mitigation.Upsert(ctx, tx, s.svc.scope); err != nil {
+					return fmt.Errorf("cannot upsert mitigation: %w", err)
+				}
+
+				for j := range req.Mitigations[i].Tasks {
+					taskID, err := gid.NewGID(organizationID.TenantID(), coredata.TaskEntityType)
+					if err != nil {
+						return fmt.Errorf("cannot create global id: %w", err)
+					}
+
+					task := &coredata.Task{
+						ID:           taskID,
+						MitigationID: mitigation.ID,
+						Name:         req.Mitigations[i].Tasks[j].Name,
+						Description:  req.Mitigations[i].Tasks[j].Description,
+						ReferenceID:  req.Mitigations[i].Tasks[j].ReferenceID,
+						State:        coredata.TaskStateTodo,
+						CreatedAt:    now,
+						UpdatedAt:    now,
+					}
+
+					if err := task.Upsert(ctx, tx, s.svc.scope); err != nil {
+						return fmt.Errorf("cannot upsert task: %w", err)
+					}
+				}
+
+				for _, standard := range req.Mitigations[i].Standards {
+					framework := &coredata.Framework{}
+					if err := framework.LoadByReferenceID(ctx, tx, s.svc.scope, standard.Framework); err != nil {
+						continue
+					}
+
+					control := &coredata.Control{}
+					if err := control.LoadByFrameworkIDAndReferenceID(ctx, tx, s.svc.scope, framework.ID, standard.Control); err != nil {
+						continue
+					}
+
+					controlMitigation := &coredata.ControlMitigation{
+						ControlID:    control.ID,
+						MitigationID: mitigation.ID,
+						CreatedAt:    now,
+					}
+
+					if err := controlMitigation.Upsert(ctx, tx, s.svc.scope); err != nil {
+						return fmt.Errorf("cannot insert control mitigation: %w", err)
+					}
 				}
 			}
 
@@ -271,7 +325,6 @@ func (s MitigationService) Create(
 		Description:    req.Description,
 		Category:       req.Category,
 		State:          coredata.MitigationStateNotStarted,
-		Standards:      []string{},
 		Importance:     req.Importance,
 		CreatedAt:      now,
 		UpdatedAt:      now,

pkg/probo/task_service.go

@@ -22,6 +22,7 @@ import (
 	"github.com/getprobo/probo/pkg/coredata"
 	"github.com/getprobo/probo/pkg/gid"
 	"github.com/getprobo/probo/pkg/page"
+	"github.com/getprobo/probo/pkg/slug"
 	"go.gearno.de/kit/pg"
 )
 
@@ -65,6 +66,7 @@ func (s TaskService) Create(
 		TimeEstimate: req.TimeEstimate,
 		AssignedToID: req.AssignedToID,
 		State:        coredata.TaskStateTodo,
+		ReferenceID:  slug.Make(req.Name),
 		CreatedAt:    now,
 		UpdatedAt:    now,
 	}