Allow setting owner on risk factors

mondoohq · Jan 29, 2025 · a278ccf · a278ccf
1 parent 76bc97d
commit a278ccf
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 4 deletions.
diff --git a/internal/datalakes/inmemory/policyhub.go b/internal/datalakes/inmemory/policyhub.go
@@ -324,7 +324,7 @@ func (db *Db) getRiskFactor(ctx context.Context, mrn string) (*policy.RiskFactor
 	return found.(*policy.RiskFactor), nil
 }
 
-func (db *Db) SetRiskFactor(ctx context.Context, riskFactor *policy.RiskFactor) error {
+func (db *Db) SetRiskFactor(ctx context.Context, ownerMrn string, riskFactor *policy.RiskFactor) error {
 	db.cache.Set(dbIDRiskFactor+riskFactor.Mrn, riskFactor, 1)
 	return nil
 }

diff --git a/internal/datalakes/inmemory/policyresolver.go b/internal/datalakes/inmemory/policyresolver.go
@@ -632,7 +632,7 @@ func (db *Db) SetResolvedPolicy(ctx context.Context, mrn string, resolvedPolicy
 		if _, err := db.getRiskFactor(ctx, rf.Mrn); err == nil {
 			continue
 		}
-		if err := db.SetRiskFactor(ctx, rf); err != nil {
+		if err := db.SetRiskFactor(ctx, "", rf); err != nil {
 			return err
 		}
 	}

diff --git a/policy/datalake.go b/policy/datalake.go
@@ -61,7 +61,7 @@ type DataLake interface {
 	// SetPolicy stores a given policy in the data lake
 	SetPolicy(ctx context.Context, policy *Policy, filters []*explorer.Mquery) error
 	// SetRiskFactor creates and stores a risk factor
-	SetRiskFactor(ctx context.Context, riskFactor *RiskFactor) error
+	SetRiskFactor(ctx context.Context, ownerMrn string, riskFactor *RiskFactor) error
 
 	// List all policies for a given owner
 	// Note: Owner MRN is required

diff --git a/policy/hub.go b/policy/hub.go
@@ -201,7 +201,7 @@ func (s *LocalServices) SetBundleMap(ctx context.Context, bundleMap *PolicyBundl
 	}
 
 	for _, risk := range bundleMap.RiskFactors {
-		if err := s.DataLake.SetRiskFactor(ctx, risk); err != nil {
+		if err := s.DataLake.SetRiskFactor(ctx, bundleMap.OwnerMrn, risk); err != nil {
 			return err
 		}
 	}