Add data source processing to rule update/create.

This change adds code for the maintenance of references to Data
Sources during the creation or update of rule types.

To keep things simple and limit the amount of added code, the code
path is the same for both create and update, in that it always try to
delete references to data sources even in the case of rule creation,
which is always a noop.

Fixes #5049

internal/controlplane/handlers_ruletype.go

@@ -176,8 +176,12 @@ func (s *Server) CreateRuleType(
 		return nil, util.UserVisibleError(codes.InvalidArgument, "%s", err)
 	}
 
+	ruleDS := crt.GetRuleType().GetDef().GetEval().GetDataSources()
+	if len(ruleDS) > 0 && !flags.Bool(ctx, s.featureFlags, flags.DataSources) {
+		return nil, util.UserVisibleError(codes.InvalidArgument, "DataSources feature is disabled")
+	}
+
 	newRuleType, err := db.WithTransaction(s.store, func(qtx db.ExtendQuerier) (*minderv1.RuleType, error) {
-		ruleDS := crt.GetRuleType().GetDef().GetEval().GetDataSources()
 		if err := s.validateDataSources(ctx, projectID, ruleDS, qtx); err != nil {
 			// We expect the error to be a user visible error
 			return nil, err
@@ -222,8 +226,12 @@ func (s *Server) UpdateRuleType(
 		return nil, util.UserVisibleError(codes.InvalidArgument, "%s", err)
 	}
 
+	ruleDS := urt.GetRuleType().GetDef().GetEval().GetDataSources()
+	if len(ruleDS) > 0 && !flags.Bool(ctx, s.featureFlags, flags.DataSources) {
+		return nil, util.UserVisibleError(codes.InvalidArgument, "DataSources feature is disabled")
+	}
+
 	updatedRuleType, err := db.WithTransaction(s.store, func(qtx db.ExtendQuerier) (*minderv1.RuleType, error) {
-		ruleDS := urt.GetRuleType().GetDef().GetEval().GetDataSources()
 		if err := s.validateDataSources(ctx, projectID, ruleDS, qtx); err != nil {
 			// We expect the error to be a user visible error
 			return nil, err
@@ -388,10 +396,6 @@ func (s *Server) validateDataSources(
 		return nil
 	}
 
-	if len(ruleDS) > 0 && !flags.Bool(ctx, s.featureFlags, flags.DataSources) {
-		return status.Errorf(codes.Unavailable, "DataSources feature is disabled")
-	}
-
 	opts := datasourcesvc.ReadBuilder().WithTransaction(qtx)
 	for _, requested := range ruleDS {
 		_, err := s.dataSourcesService.GetByName(

pkg/ruletypes/service.go

@@ -150,6 +150,14 @@ func (_ *ruleTypeService) CreateRuleType(
 		return nil, fmt.Errorf("failed to create rule type: %w", err)
 	}
 
+	// Data Sources reference update. Note that this step can be
+	// safely performed after updating the rule, as the only thing
+	// we need from the previous code is project id and rule id.
+	ds := ruleTypeDef.GetEval().GetDataSources()
+	if err := processDataSources(ctx, newDBRecord.ID, ds, projectID, projects, qtx); err != nil {
+		return nil, fmt.Errorf("failed updating references to data sources: %w", err)
+	}
+
 	logger.BusinessRecord(ctx).RuleType = logger.RuleType{Name: newDBRecord.Name, ID: newDBRecord.ID}
 
 	rt, err := RuleTypePBFromDB(&newDBRecord)
@@ -229,6 +237,19 @@ func (_ *ruleTypeService) UpdateRuleType(
 		return nil, fmt.Errorf("failed to update rule type: %w", err)
 	}
 
+	projects, err := qtx.GetParentProjects(ctx, projectID)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get parent projects: %w", err)
+	}
+
+	// Data Sources reference update. Note that this step can be
+	// safely performed after updating the rule, as the only thing
+	// we need from the previous code is project id and rule id.
+	ds := ruleTypeDef.GetEval().GetDataSources()
+	if err := processDataSources(ctx, oldRuleType.ID, ds, projectID, projects, qtx); err != nil {
+		return nil, fmt.Errorf("failed updating references to data sources: %w", err)
+	}
+
 	logger.BusinessRecord(ctx).RuleType = logger.RuleType{Name: oldRuleType.Name, ID: oldRuleType.ID}
 
 	result, err := RuleTypePBFromDB(&updatedRuleType)
@@ -300,3 +321,71 @@ func validateRuleUpdate(existingRecord *db.RuleType, newRuleType *pb.RuleType) e
 
 	return nil
 }
+
+func processDataSources(
+	ctx context.Context,
+	ruleID uuid.UUID,
+	ds []*pb.DataSourceReference,
+	projectID uuid.UUID,
+	projectHierarchy []uuid.UUID,
+	qtx db.Querier,
+) error {
+	// We first verify that the data sources required are
+	// available within the project hierarchy.
+	datasources, err := getAvailableDataSources(ctx, ds, projectHierarchy, qtx)
+	if err != nil {
+		return fmt.Errorf("data source not available: %w", err)
+	}
+
+	// Then, we proceed to delete any data source reference we
+	// have for the old definition of the rule type.
+	deleteArgs := db.DeleteRuleTypeDataSourceParams{
+		Ruleid:    ruleID,
+		Projectid: projectID,
+	}
+	if err := qtx.DeleteRuleTypeDataSource(ctx, deleteArgs); err != nil {
+		return fmt.Errorf("error deleting references to data source: %w", err)
+	}
+
+	// Finally, we add references to the required data source.
+	fmt.Printf("WAAAAAAGH! %+v\n", len(datasources))
+	for _, datasource := range datasources {
+		insertArgs := db.AddRuleTypeDataSourceReferenceParams{
+			Ruletypeid:   ruleID,
+			Datasourceid: datasource.ID,
+			Projectid:    projectID,
+		}
+		if _, err := qtx.AddRuleTypeDataSourceReference(ctx, insertArgs); err != nil {
+			return fmt.Errorf("error adding references to data source: %w", err)
+		}
+	}
+
+	return nil
+}
+
+func getAvailableDataSources(
+	ctx context.Context,
+	requiredDataSources []*pb.DataSourceReference,
+	projects []uuid.UUID,
+	qtx db.Querier,
+) ([]db.DataSource, error) {
+	datasources := make([]db.DataSource, 0)
+
+	for _, datasource := range requiredDataSources {
+		qarg := db.GetDataSourceByNameParams{
+			Name:     datasource.Name,
+			Projects: projects,
+		}
+		dbDataSource, err := qtx.GetDataSourceByName(ctx, qarg)
+		if err != nil && errors.Is(err, sql.ErrNoRows) {
+			return nil, fmt.Errorf("data source of name %s not found", datasource.Name)
+		}
+		if err != nil {
+			return nil, fmt.Errorf("failed getting data sources: %w", err)
+		}
+
+		datasources = append(datasources, dbDataSource)
+	}
+
+	return datasources, nil
+}

pkg/ruletypes/service_test.go

@@ -105,13 +105,13 @@ func TestRuleTypeService(t *testing.T) {
 		{
 			Name:       "CreateRuleType successfully creates a new rule type",
 			RuleType:   newRuleType(withBasicStructure),
-			DBSetup:    dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withSuccessfulDeleteRuleTypeDataSource),
 			TestMethod: create,
 		},
 		{
 			Name:           "CreateRuleType successfully creates a new namespaced rule type",
 			RuleType:       newRuleType(withBasicStructure, withRuleName(namespacedRuleName)),
-			DBSetup:        dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulNamespaceCreate),
+			DBSetup:        dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulNamespaceCreate, withSuccessfulDeleteRuleTypeDataSource),
 			SubscriptionID: subscriptionID,
 			TestMethod:     create,
 		},
@@ -182,40 +182,120 @@ func TestRuleTypeService(t *testing.T) {
 		{
 			Name:       "UpdateRuleType successfully updates an existing rule",
 			RuleType:   newRuleType(withBasicStructure),
-			DBSetup:    dbf.NewDBMock(withSuccessfulGet, withSuccessfulUpdate),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withSuccessfulDeleteRuleTypeDataSource),
 			TestMethod: update,
 		},
 		{
-			Name:           "UpdateRuleType successfully updates an existing rule",
+			Name:           "UpdateRuleType successfully updates an existing rule with subscription",
 			RuleType:       newRuleType(withBasicStructure),
-			DBSetup:        dbf.NewDBMock(withSuccessfulNamespaceGet, withSuccessfulUpdate),
+			DBSetup:        dbf.NewDBMock(withHierarchyGet, withSuccessfulNamespaceGet, withSuccessfulUpdate, withSuccessfulDeleteRuleTypeDataSource),
 			TestMethod:     update,
 			SubscriptionID: subscriptionID,
 		},
 		{
 			Name:           "UpsertRuleType successfully creates a new namespaced rule type",
 			RuleType:       newRuleType(withBasicStructure, withRuleName(namespacedRuleName)),
-			DBSetup:        dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulNamespaceCreate),
+			DBSetup:        dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulNamespaceCreate, withSuccessfulDeleteRuleTypeDataSource),
 			SubscriptionID: subscriptionID,
 			TestMethod:     upsert,
 		},
 		{
 			Name:           "UpsertRuleType successfully updates an existing rule",
 			RuleType:       newRuleType(withBasicStructure, withRuleName(namespacedRuleName)),
-			DBSetup:        dbf.NewDBMock(withHierarchyGet, withSuccessfulNamespaceGet, withSuccessfulUpdate),
+			DBSetup:        dbf.NewDBMock(withHierarchyGet, withHierarchyGet, withSuccessfulNamespaceGet, withSuccessfulUpdate, withSuccessfulDeleteRuleTypeDataSource),
 			TestMethod:     upsert,
 			SubscriptionID: subscriptionID,
 		},
 		{
 			Name:       "CreateRuleType with EvaluationFailureMessage",
 			RuleType:   newRuleType(withBasicStructure, withEvaluationFailureMessage(shortFailureMessage)),
-			DBSetup:    dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreateWithEvaluationFailureMessage),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreateWithEvaluationFailureMessage, withSuccessfulDeleteRuleTypeDataSource),
 			TestMethod: create,
 		},
 		{
 			Name:       "UpdateRuleType with EvaluationFailureMessage",
 			RuleType:   newRuleType(withBasicStructure, withEvaluationFailureMessage(shortFailureMessage)),
-			DBSetup:    dbf.NewDBMock(withSuccessfulGet, withSuccessfulUpdateWithEvaluationFailureMessage),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdateWithEvaluationFailureMessage, withSuccessfulDeleteRuleTypeDataSource),
+			TestMethod: update,
+		},
+		{
+			Name:       "CreateRuleType with Data Sources",
+			RuleType:   newRuleType(withBasicStructure, withDataSources),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withSuccessfulGetDataSourcesByName(1), withSuccessfulDeleteRuleTypeDataSource, withSuccessfulAddRuleTypeDataSourceReference),
+			TestMethod: create,
+		},
+		{
+			Name:       "UpdateRuleType with Data Sources",
+			RuleType:   newRuleType(withBasicStructure, withDataSources),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withSuccessfulGetDataSourcesByName(1), withSuccessfulDeleteRuleTypeDataSource, withSuccessfulAddRuleTypeDataSourceReference),
+			TestMethod: update,
+		},
+		{
+			Name:          "CreateRuleType with Data Sources not found",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "data source not available",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withNotFoundGetDataSourcesByName),
+			TestMethod:    create,
+		},
+		{
+			Name:          "UpdateRuleType with Data Sources not found",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "data source not available",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withNotFoundGetDataSourcesByName),
+			TestMethod:    update,
+		},
+		{
+			Name:          "CreateRuleType with Data Sources failed get",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "failed getting data sources",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withFailedGetDataSourcesByName),
+			TestMethod:    create,
+		},
+		{
+			Name:          "UpdateRuleType with Data Sources failed get",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "failed getting data sources",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withFailedGetDataSourcesByName),
+			TestMethod:    update,
+		},
+		{
+			Name:          "CreateRuleType with Data Sources failed delete",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "error deleting references to data source",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withSuccessfulGetDataSourcesByName(1), withFailedDeleteRuleTypeDataSource),
+			TestMethod:    create,
+		},
+		{
+			Name:          "UpdateRuleType with Data Sources failed delete",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "error adding references to data source",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withSuccessfulGetDataSourcesByName(1), withSuccessfulDeleteRuleTypeDataSource, withFailedAddRuleTypeDataSourceReference),
+			TestMethod:    update,
+		},
+		{
+			Name:          "CreateRuleType with Data Sources failed add",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "error adding references to data source",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withSuccessfulGetDataSourcesByName(1), withSuccessfulDeleteRuleTypeDataSource, withFailedAddRuleTypeDataSourceReference),
+			TestMethod:    create,
+		},
+		{
+			Name:          "UpdateRuleType with Data Sources failed add",
+			RuleType:      newRuleType(withBasicStructure, withDataSources),
+			ExpectedError: "error adding references to data source",
+			DBSetup:       dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withSuccessfulGetDataSourcesByName(1), withSuccessfulDeleteRuleTypeDataSource, withFailedAddRuleTypeDataSourceReference),
+			TestMethod:    update,
+		},
+		{
+			Name:       "CreateRuleType with Data Sources multiple adds",
+			RuleType:   newRuleType(withBasicStructure, withDataSources, withDataSources),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withNotFoundGet, withSuccessfulCreate, withSuccessfulGetDataSourcesByName(2), withSuccessfulDeleteRuleTypeDataSource, withSuccessfulAddRuleTypeDataSourceReference),
+			TestMethod: create,
+		},
+		{
+			Name:       "UpdateRuleType with Data Sources multiple add",
+			RuleType:   newRuleType(withBasicStructure, withDataSources, withDataSources),
+			DBSetup:    dbf.NewDBMock(withHierarchyGet, withSuccessfulGet, withSuccessfulUpdate, withSuccessfulGetDataSourcesByName(2), withSuccessfulDeleteRuleTypeDataSource, withSuccessfulAddRuleTypeDataSourceReference),
 			TestMethod: update,
 		},
 	}
@@ -338,6 +418,16 @@ func withBasicStructure(ruleType *pb.RuleType) {
 	ruleType.Severity = &pb.Severity{Value: pb.Severity_VALUE_HIGH}
 }
 
+func withDataSources(ruleType *pb.RuleType) {
+	datasources := []*pb.DataSourceReference{
+		{
+			// We just need a random string
+			Name: fmt.Sprintf("foo-%s", uuid.New().String()),
+		},
+	}
+	ruleType.Def.Eval.DataSources = datasources
+}
+
 func withEvaluationFailureMessage(message string) func(ruleType *pb.RuleType) {
 	return func(ruleType *pb.RuleType) {
 		ruleType.ShortFailureMessage = message
@@ -431,6 +521,52 @@ func withFailedUpdate(mock dbf.DBMock) {
 		Return(db.RuleType{}, errDefault)
 }
 
+func withSuccessfulGetDataSourcesByName(times int) func(dbf.DBMock) {
+	return func(mock dbf.DBMock) {
+		call := mock.EXPECT().
+			GetDataSourceByName(gomock.Any(), gomock.Any())
+		for i := 0; i < times; i++ {
+			call = call.Return(db.DataSource{Name: fmt.Sprintf("foo-%d", i)}, nil)
+		}
+	}
+}
+
+func withNotFoundGetDataSourcesByName(mock dbf.DBMock) {
+	mock.EXPECT().
+		GetDataSourceByName(gomock.Any(), gomock.Any()).
+		Return(db.DataSource{}, sql.ErrNoRows)
+}
+
+func withFailedGetDataSourcesByName(mock dbf.DBMock) {
+	mock.EXPECT().
+		GetDataSourceByName(gomock.Any(), gomock.Any()).
+		Return(db.DataSource{}, errDefault)
+}
+
+func withSuccessfulDeleteRuleTypeDataSource(mock dbf.DBMock) {
+	mock.EXPECT().
+		DeleteRuleTypeDataSource(gomock.Any(), gomock.Any()).
+		Return(nil)
+}
+
+func withFailedDeleteRuleTypeDataSource(mock dbf.DBMock) {
+	mock.EXPECT().
+		DeleteRuleTypeDataSource(gomock.Any(), gomock.Any()).
+		Return(errDefault)
+}
+
+func withSuccessfulAddRuleTypeDataSourceReference(mock dbf.DBMock) {
+	mock.EXPECT().
+		AddRuleTypeDataSourceReference(gomock.Any(), gomock.Any()).
+		Return(db.RuleTypeDataSource{}, nil)
+}
+
+func withFailedAddRuleTypeDataSourceReference(mock dbf.DBMock) {
+	mock.EXPECT().
+		AddRuleTypeDataSourceReference(gomock.Any(), gomock.Any()).
+		Return(db.RuleTypeDataSource{}, errDefault)
+}
+
 func newDBRuleType(severity db.Severity, subscriptionID uuid.UUID, failureMessage string) db.RuleType {
 	name := ruleName
 	if subscriptionID != uuid.Nil {