Improve framework data isolation

Signed-off-by: Bryan Frimin <bryan@getprobo.com>

pkg/coredata/framework.go

@@ -250,6 +250,7 @@ func (f Framework) Delete(
 	ctx context.Context,
 	conn pg.Conn,
 	scope Scoper,
+	frameworkID gid.GID,
 ) error {
 	q := `
 DELETE
@@ -260,7 +261,7 @@ WHERE
     AND id = @framework_id;
 `
 
-	args := pgx.StrictNamedArgs{"framework_id": f.ID}
+	args := pgx.StrictNamedArgs{"framework_id": frameworkID}
 	maps.Copy(args, scope.SQLArguments())
 	q = fmt.Sprintf(q, scope.SQLFragment())
 

pkg/probo/framework_service.go

@@ -69,22 +69,31 @@ func (s FrameworkService) Create(
 	req CreateFrameworkRequest,
 ) (*coredata.Framework, error) {
 	now := time.Now()
-	frameworkID := gid.New(s.svc.scope.GetTenantID(), coredata.FrameworkEntityType)
+	organization := &coredata.Organization{}
 
 	framework := &coredata.Framework{
-		ID:             frameworkID,
+		ID:          gid.New(s.svc.scope.GetTenantID(), coredata.FrameworkEntityType),
-		OrganizationID: req.OrganizationID,
+		Name:        req.Name,
-		Name:           req.Name,
+		Description: req.Description,
-		Description:    req.Description,
+		ReferenceID: slug.Make(req.Name),
-		ReferenceID:    slug.Make(req.Name),
+		CreatedAt:   now,
-		CreatedAt:      now,
+		UpdatedAt:   now,
-		UpdatedAt:      now,
 	}
 
-	err := s.svc.pg.WithConn(
+	err := s.svc.pg.WithTx(
 		ctx,
 		func(conn pg.Conn) error {
-			return framework.Insert(ctx, conn, s.svc.scope)
+			if err := organization.LoadByID(ctx, conn, s.svc.scope, req.OrganizationID); err != nil {
+				return fmt.Errorf("cannot load organization: %w", err)
+			}
+
+			framework.OrganizationID = organization.ID
+
+			if err := framework.Insert(ctx, conn, s.svc.scope); err != nil {
+				return fmt.Errorf("cannot insert framework: %w", err)
+			}
+
+			return nil
 		},
 	)
 
@@ -101,17 +110,27 @@ func (s FrameworkService) ListForOrganizationID(
 	cursor *page.Cursor[coredata.FrameworkOrderField],
 ) (*page.Page[*coredata.Framework, coredata.FrameworkOrderField], error) {
 	var frameworks coredata.Frameworks
+	organization := &coredata.Organization{}
 
 	err := s.svc.pg.WithConn(
 		ctx,
 		func(conn pg.Conn) error {
-			return frameworks.LoadByOrganizationID(
+			if err := organization.LoadByID(ctx, conn, s.svc.scope, organizationID); err != nil {
+				return fmt.Errorf("cannot load organization: %w", err)
+			}
+
+			err := frameworks.LoadByOrganizationID(
 				ctx,
 				conn,
 				s.svc.scope,
-				organizationID,
+				organization.ID,
 				cursor,
 			)
+			if err != nil {
+				return fmt.Errorf("cannot load frameworks: %w", err)
+			}
+
+			return nil
 		},
 	)
 
@@ -177,12 +196,12 @@ func (s FrameworkService) Delete(
 	ctx context.Context,
 	frameworkID gid.GID,
 ) error {
-	framework := &coredata.Framework{ID: frameworkID}
+	framework := &coredata.Framework{}
 
 	return s.svc.pg.WithConn(
 		ctx,
 		func(conn pg.Conn) error {
-			return framework.Delete(ctx, conn, s.svc.scope)
+			return framework.Delete(ctx, conn, s.svc.scope, frameworkID)
 		},
 	)
 }
@@ -192,46 +211,47 @@ func (s FrameworkService) Import(
 	organizationID gid.GID,
 	req ImportFrameworkRequest,
 ) (*coredata.Framework, error) {
+	var framework *coredata.Framework
 	frameworkID := gid.New(organizationID.TenantID(), coredata.FrameworkEntityType)
-
 	now := time.Now()
-	framework := &coredata.Framework{
-		ID:             frameworkID,
-		OrganizationID: organizationID,
-		ReferenceID:    req.Framework.ID,
-		Name:           req.Framework.Name,
-		CreatedAt:      now,
-		UpdatedAt:      now,
-	}
-
-	importedControls := coredata.Controls{}
-	for _, control := range req.Framework.Controls {
-		controlID := gid.New(organizationID.TenantID(), coredata.ControlEntityType)
-
-		now := time.Now()
-		control := &coredata.Control{
-			ID:          controlID,
-			TenantID:    organizationID.TenantID(),
-			FrameworkID: frameworkID,
-			ReferenceID: control.ID,
-			Name:        control.Name,
-			Description: control.Description,
-			CreatedAt:   now,
-			UpdatedAt:   now,
-		}
-
-		importedControls = append(importedControls, control)
-	}
 
 	err := s.svc.pg.WithTx(
 		ctx,
 		func(tx pg.Conn) error {
+			organization := &coredata.Organization{}
+			if err := organization.LoadByID(ctx, tx, s.svc.scope, organizationID); err != nil {
+				return fmt.Errorf("cannot load organization: %w", err)
+			}
+
+			framework = &coredata.Framework{
+				ID:             frameworkID,
+				OrganizationID: organization.ID,
+				ReferenceID:    req.Framework.ID,
+				Name:           req.Framework.Name,
+				CreatedAt:      now,
+				UpdatedAt:      now,
+			}
+
 			if err := framework.Insert(ctx, tx, s.svc.scope); err != nil {
 				return fmt.Errorf("cannot insert framework: %w", err)
 			}
 
-			for _, importedControl := range importedControls {
+			for _, control := range req.Framework.Controls {
-				if err := importedControl.Insert(ctx, tx, s.svc.scope); err != nil {
+				controlID := gid.New(organization.ID.TenantID(), coredata.ControlEntityType)
+
+				now := time.Now()
+				control := &coredata.Control{
+					ID:          controlID,
+					TenantID:    organizationID.TenantID(),
+					FrameworkID: frameworkID,
+					ReferenceID: control.ID,
+					Name:        control.Name,
+					Description: control.Description,
+					CreatedAt:   now,
+					UpdatedAt:   now,
+				}
+
+				if err := control.Insert(ctx, tx, s.svc.scope); err != nil {
 					return fmt.Errorf("cannot insert control: %w", err)
 				}
 			}