Add support for subscriptions to data sources (#5271)

database/migrations/000111_data_sources_subscription_id.down.sql

@@ -0,0 +1,8 @@
+-- SPDX-FileCopyrightText: Copyright 2025 The Minder Authors
+-- SPDX-License-Identifier: Apache-2.0
+
+BEGIN;
+
+ALTER TABLE rule_type DROP COLUMN subscription_id;
+
+COMMIT;

database/migrations/000111_data_sources_subscription_id.up.sql

@@ -0,0 +1,10 @@
+-- SPDX-FileCopyrightText: Copyright 2025 The Minder Authors
+-- SPDX-License-Identifier: Apache-2.0
+
+BEGIN;
+
+ALTER TABLE data_sources
+    ADD COLUMN subscription_id UUID DEFAULT NULL
+        REFERENCES subscriptions(id);
+
+COMMIT;

database/query/datasources.sql

@@ -1,8 +1,8 @@
 -- CreateDataSource creates a new datasource in a given project.
 
 -- name: CreateDataSource :one
-INSERT INTO data_sources (project_id, name, display_name)
-VALUES ($1, $2, $3) RETURNING *;
+INSERT INTO data_sources (project_id, name, display_name, subscription_id)
+VALUES ($1, $2, $3, sqlc.narg(subscription_id)) RETURNING *;
 
 -- AddDataSourceFunction adds a function to a datasource.
 

internal/controlplane/handlers_datasource.go

@@ -36,7 +36,7 @@ func (s *Server) CreateDataSource(ctx context.Context,
 	}
 
 	// Process the request
-	ret, err := s.dataSourcesService.Create(ctx, dsReq, nil)
+	ret, err := s.dataSourcesService.Create(ctx, uuid.Nil, dsReq, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -167,7 +167,7 @@ func (s *Server) UpdateDataSource(ctx context.Context,
 	}
 
 	// Process the request
-	ret, err := s.dataSourcesService.Update(ctx, dsReq, nil)
+	ret, err := s.dataSourcesService.Update(ctx, uuid.Nil, dsReq, nil)
 	if err != nil {
 		return nil, err
 	}

internal/controlplane/handlers_datasource_test.go

@@ -37,7 +37,7 @@ func TestCreateDataSource(t *testing.T) {
 			setupMocks: func(dsService *mock_service.MockDataSourcesService, featureClient *flags.FakeClient) {
 				featureClient.Data = map[string]interface{}{"data_sources": true}
 				dsService.EXPECT().
-					Create(gomock.Any(), gomock.Any(), gomock.Any()).
+					Create(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
 					Return(&minderv1.DataSource{Name: "test-ds"}, nil)
 			},
 			request: &minderv1.CreateDataSourceRequest{
@@ -417,7 +417,7 @@ func TestUpdateDataSource(t *testing.T) {
 			setupMocks: func(dsService *mock_service.MockDataSourcesService, featureClient *flags.FakeClient, _ *mockdb.MockStore) {
 				featureClient.Data = map[string]interface{}{"data_sources": true}
 				dsService.EXPECT().
-					Update(gomock.Any(), gomock.Any(), gomock.Any()).
+					Update(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
 					Return(&minderv1.DataSource{Id: dsIDStr, Name: "updated-ds"}, nil)
 			},
 			request: &minderv1.UpdateDataSourceRequest{

internal/datasources/service/helpers.go

@@ -5,6 +5,7 @@ package service
 
 import (
 	"context"
+	"database/sql"
 	"errors"
 	"fmt"
 
@@ -17,6 +18,39 @@ import (
 	minderv1 "github.com/mindersec/minder/pkg/api/protobuf/go/minder/v1"
 )
 
+var (
+	getByNameQuery = func(ctx context.Context, tx db.ExtendQuerier, projs []uuid.UUID, name string) (db.DataSource, error) {
+		ds, err := tx.GetDataSourceByName(ctx, db.GetDataSourceByNameParams{
+			Name:     name,
+			Projects: projs,
+		})
+		if err != nil {
+			if errors.Is(err, sql.ErrNoRows) {
+				return db.DataSource{}, util.UserVisibleError(codes.NotFound,
+					"data source of name %s not found", name)
+			}
+			return db.DataSource{}, fmt.Errorf("failed to get data source by name: %w", err)
+		}
+
+		return ds, nil
+	}
+	getByIDQuery = func(ctx context.Context, tx db.ExtendQuerier, projs []uuid.UUID, id uuid.UUID) (db.DataSource, error) {
+		ds, err := tx.GetDataSource(ctx, db.GetDataSourceParams{
+			ID:       id,
+			Projects: projs,
+		})
+		if errors.Is(err, sql.ErrNoRows) {
+			return db.DataSource{}, util.UserVisibleError(codes.NotFound,
+				"data source of id %s not found", id.String())
+		}
+		if err != nil {
+			return db.DataSource{}, fmt.Errorf("failed to get data source by name: %w", err)
+		}
+
+		return ds, nil
+	}
+)
+
 func (d *dataSourceService) getDataSourceSomehow(
 	ctx context.Context,
 	project uuid.UUID,
@@ -33,23 +67,55 @@ func (d *dataSourceService) getDataSourceSomehow(
 
 	tx := stx.Q()
 
+	ds, err := getDataSourceFromDb(ctx, project, opts, tx, theSomehow)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get data source from DB: %w", err)
+	}
+
+	dsfuncs, err := getDataSourceFunctions(ctx, tx, ds)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get data source functions: %w", err)
+	}
+
+	if err := stx.Commit(); err != nil {
+		return nil, fmt.Errorf("failed to commit transaction: %w", err)
+	}
+
+	return dataSourceDBToProtobuf(*ds, dsfuncs)
+}
+
+func getDataSourceFromDb(
+	ctx context.Context,
+	project uuid.UUID,
+	opts *ReadOptions,
+	qtx db.ExtendQuerier,
+	dbQuery func(ctx context.Context, qtx db.ExtendQuerier, projs []uuid.UUID) (db.DataSource, error),
+) (*db.DataSource, error) {
 	var projs []uuid.UUID
 	if len(opts.hierarchy) > 0 {
 		projs = opts.hierarchy
 	} else {
-		prjs, err := listRelevantProjects(ctx, tx, project, opts.canSearchHierarchical())
+		prjs, err := listRelevantProjects(ctx, qtx, project, opts.canSearchHierarchical())
 		if err != nil {
 			return nil, fmt.Errorf("failed to list relevant projects: %w", err)
 		}
 
 		projs = prjs
 	}
 
-	ds, err := theSomehow(ctx, tx, projs)
+	ds, err := dbQuery(ctx, qtx, projs)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get data source by name: %w", err)
+		return nil, fmt.Errorf("failed to get data source from DB: %w", err)
 	}
 
+	return &ds, nil
+}
+
+func getDataSourceFunctions(
+	ctx context.Context,
+	tx db.ExtendQuerier,
+	ds *db.DataSource,
+) ([]db.DataSourcesFunction, error) {
 	dsfuncs, err := tx.ListDataSourceFunctions(ctx, db.ListDataSourceFunctionsParams{
 		DataSourceID: ds.ID,
 		ProjectID:    ds.ProjectID,
@@ -66,11 +132,7 @@ func (d *dataSourceService) getDataSourceSomehow(
 		return nil, errors.New("data source has no functions")
 	}
 
-	if err := stx.Commit(); err != nil {
-		return nil, fmt.Errorf("failed to commit transaction: %w", err)
-	}
-
-	return dataSourceDBToProtobuf(ds, dsfuncs)
+	return dsfuncs, nil
 }
 
 func (d *dataSourceService) instantiateDataSource(
@@ -110,9 +172,15 @@ func listRelevantProjects(
 }
 
 func validateDataSourceFunctionsUpdate(
-	existingDS, newDS *minderv1.DataSource,
+	existingDS *db.DataSource, existingFunctions []db.DataSourcesFunction, newDS *minderv1.DataSource,
 ) error {
-	existingImpl, err := datasources.BuildFromProtobuf(existingDS)
+	existingDsProto, err := dataSourceDBToProtobuf(*existingDS, existingFunctions)
+	if err != nil {
+		// If we got here, it means the existing data source is invalid.
+		return fmt.Errorf("failed to convert data source to protobuf: %w", err)
+	}
+
+	existingImpl, err := datasources.BuildFromProtobuf(existingDsProto)
 	if err != nil {
 		// If we got here, it means the existing data source is invalid.
 		return fmt.Errorf("failed to build data source from protobuf: %w", err)