chore: warehouse backend config refactoring

runner/runner.go

@@ -342,7 +342,6 @@ func runAllInit() {
 	warehouse.Init2()
 	warehouse.Init3()
 	warehouse.Init4()
-	warehouse.Init5()
 	warehouse.Init6()
 	warehousearchiver.Init()
 	validations.Init()

warehouse/admin.go

@@ -20,7 +20,10 @@ import (
 	warehouseutils "github.com/rudderlabs/rudder-server/warehouse/utils"
 )
 
-type WarehouseAdmin struct{}
+type Admin struct {
+	bcManager *backendConfigManager
+	logger    logger.Logger
+}
 
 type QueryInput struct {
 	DestID       string
@@ -37,12 +40,15 @@ type ConfigurationTestOutput struct {
 	Error string
 }
 
-func Init5() {
-	admin.RegisterAdminHandler("Warehouse", &WarehouseAdmin{})
+func RegisterAdmin(bcManager *backendConfigManager, logger logger.Logger) {
+	admin.RegisterAdminHandler("Warehouse", &Admin{
+		bcManager: bcManager,
+		logger:    logger.Child("admin"),
+	})
 }
 
 // TriggerUpload sets uploads to start without delay
-func (*WarehouseAdmin) TriggerUpload(off bool, reply *string) error {
+func (*Admin) TriggerUpload(off bool, reply *string) error {
 	startUploadAlways = !off
 	if off {
 		*reply = "Turned off explicit warehouse upload triggers.\nWarehouse uploads will continue to be done as per schedule in control plane."
@@ -53,12 +59,12 @@ func (*WarehouseAdmin) TriggerUpload(off bool, reply *string) error {
 }
 
 // Query the underlying warehouse
-func (*WarehouseAdmin) Query(s QueryInput, reply *warehouseutils.QueryResult) error {
+func (a *Admin) Query(s QueryInput, reply *warehouseutils.QueryResult) error {
 	if strings.TrimSpace(s.DestID) == "" {
 		return errors.New("please specify the destination ID to query the warehouse")
 	}
 
-	srcMap, ok := bcManager.ConnectionSourcesMap(s.DestID) // TODO remove global variable
+	srcMap, ok := a.bcManager.ConnectionSourcesMap(s.DestID)
 	if !ok {
 		return errors.New("please specify a valid and existing destination ID")
 	}
@@ -92,19 +98,19 @@ func (*WarehouseAdmin) Query(s QueryInput, reply *warehouseutils.QueryResult) er
 	}
 	defer client.Close()
 
-	pkgLogger.Infof(`[WH Admin]: Querying warehouse: %s:%s`, warehouse.Type, warehouse.Destination.ID)
+	a.logger.Infof(`[WH Admin]: Querying warehouse: %s:%s`, warehouse.Type, warehouse.Destination.ID)
 	*reply, err = client.Query(s.SQLStatement)
 	return err
 }
 
 // ConfigurationTest test the underlying warehouse destination
-func (*WarehouseAdmin) ConfigurationTest(s ConfigurationTestInput, reply *ConfigurationTestOutput) error {
+func (a *Admin) ConfigurationTest(s ConfigurationTestInput, reply *ConfigurationTestOutput) error {
 	if strings.TrimSpace(s.DestID) == "" {
 		return errors.New("please specify the destination ID to query the warehouse")
 	}
 
 	var warehouse model.Warehouse
-	srcMap, ok := bcManager.ConnectionSourcesMap(s.DestID) // TODO remove global variable
+	srcMap, ok := a.bcManager.ConnectionSourcesMap(s.DestID)
 	if !ok {
 		return fmt.Errorf("please specify a valid and existing destinationID: %s", s.DestID)
 	}
@@ -114,7 +120,7 @@ func (*WarehouseAdmin) ConfigurationTest(s ConfigurationTestInput, reply *Config
 		break
 	}
 
-	pkgLogger.Infof(`[WH Admin]: Validating warehouse destination: %s:%s`, warehouse.Type, warehouse.Destination.ID)
+	a.logger.Infof(`[WH Admin]: Validating warehouse destination: %s:%s`, warehouse.Type, warehouse.Destination.ID)
 
 	destinationValidator := validations.NewDestinationValidator()
 	res := destinationValidator.Validate(context.TODO(), &warehouse.Destination)

warehouse/api.go

@@ -107,6 +107,7 @@ type UploadAPIT struct {
 	warehouseDBHandle *DB
 	log               logger.Logger
 	connectionManager *controlplane.ConnectionManager
+	bcManager         *backendConfigManager
 	isMultiWorkspace  bool
 }
 
@@ -119,7 +120,7 @@ const (
 	NoSuchSync              = "No such sync exist"
 )
 
-func InitWarehouseAPI(dbHandle *sql.DB, log logger.Logger) error {
+func InitWarehouseAPI(dbHandle *sql.DB, bcManager *backendConfigManager, log logger.Logger) error {
 	connectionToken, tokenType, isMultiWorkspace, err := deployment.GetConnectionToken()
 	if err != nil {
 		return err
@@ -161,6 +162,7 @@ func InitWarehouseAPI(dbHandle *sql.DB, log logger.Logger) error {
 				})
 			},
 		},
+		bcManager: bcManager,
 	}
 	return nil
 }
@@ -528,8 +530,8 @@ func (uploadReq *UploadReq) validateReq() error {
 }
 
 func (uploadReq *UploadReq) authorizeSource(sourceID string) bool {
-	currentList := bcManager.SourceIDsByWorkspace()
-	authorizedSourceIDs, ok := currentList[uploadReq.WorkspaceID] // TODO remove global variable
+	currentList := uploadReq.API.bcManager.SourceIDsByWorkspace()
+	authorizedSourceIDs, ok := currentList[uploadReq.WorkspaceID]
 	if !ok {
 		pkgLogger.Errorf(`Could not find sourceID in workspace %q: %v`, uploadReq.WorkspaceID, currentList)
 		return false
@@ -541,7 +543,7 @@ func (uploadReq *UploadReq) authorizeSource(sourceID string) bool {
 
 func (uploadsReq *UploadsReq) authorizedSources() []string {
 	pkgLogger.Debugf(`Getting authorizedSourceIDs for workspace:%s`, uploadsReq.WorkspaceID)
-	return bcManager.SourceIDsByWorkspace()[uploadsReq.WorkspaceID] // TODO remove global variable
+	return uploadsReq.API.bcManager.SourceIDsByWorkspace()[uploadsReq.WorkspaceID]
 }
 
 func (uploadsReq *UploadsReq) getUploadsFromDB(ctx context.Context, isMultiWorkspace bool, query string) ([]*proto.WHUploadResponse, int32, error) {

warehouse/backend_config.go

@@ -20,7 +20,6 @@ import (
 	whutils "github.com/rudderlabs/rudder-server/warehouse/utils"
 )
 
-// TODO: add tests
 func newBackendConfigManager(
 	c *config.Config,
 	db *sqlquerywrapper.DB,
@@ -106,6 +105,12 @@ func (s *backendConfigManager) Subscribe(ctx context.Context) <-chan []model.War
 	ch := make(chan []model.Warehouse, 10)
 	s.subscriptions = append(s.subscriptions, ch)
 
+	s.warehousesMu.Lock()
+	if len(s.warehouses) > 0 {
+		ch <- s.warehouses
+	}
+	s.warehousesMu.Unlock()
+
 	go func() {
 		<-ctx.Done()
 
@@ -126,26 +131,23 @@ func (s *backendConfigManager) processData(ctx context.Context, data map[string]
 	})
 
 	var (
-		warehouses               []model.Warehouse
-		connectionFlags          backendconfig.ConnectionFlags
-		sourceIDsByWorkspaceTemp = make(map[string][]string)
+		warehouses           []model.Warehouse
+		connectionFlags      backendconfig.ConnectionFlags
+		sourceIDsByWorkspace = make(map[string][]string)
 	)
+
 	for workspaceID, wConfig := range data {
 		// the last connection flags should be enough, since they are all the same in multi-workspace environments
 		connectionFlags = wConfig.ConnectionFlags
-		// map source IDs to workspace IDs
-		workspaceBySourceIDs := make(map[string]string)
 
 		for _, source := range wConfig.Sources {
-			workspaceBySourceIDs[source.ID] = workspaceID
-
-			if _, ok := sourceIDsByWorkspaceTemp[workspaceID]; !ok {
-				sourceIDsByWorkspaceTemp[workspaceID] = make([]string, 0, len(wConfig.Sources))
+			if _, ok := sourceIDsByWorkspace[workspaceID]; !ok {
+				sourceIDsByWorkspace[workspaceID] = make([]string, 0, len(wConfig.Sources))
 			}
-			sourceIDsByWorkspaceTemp[workspaceID] = append(sourceIDsByWorkspaceTemp[workspaceID], source.ID)
+			sourceIDsByWorkspace[workspaceID] = append(sourceIDsByWorkspace[workspaceID], source.ID)
 
 			for _, destination := range source.Destinations {
-				if !slices.Contains(whutils.WarehouseDestinations, destination.DestinationDefinition.Name) {
+				if _, ok := warehouseutils.WarehouseDestinationMap[destination.DestinationDefinition.Name]; !ok {
 					s.logger.Debugf("Not a warehouse destination, skipping %s", destination.DestinationDefinition.Name)
 					continue
 				}
@@ -204,7 +206,7 @@ func (s *backendConfigManager) processData(ctx context.Context, data map[string]
 	s.warehousesMu.Unlock()
 
 	s.sourceIDsByWorkspaceMu.Lock()
-	s.sourceIDsByWorkspace = sourceIDsByWorkspaceTemp
+	s.sourceIDsByWorkspace = sourceIDsByWorkspace
 	s.sourceIDsByWorkspaceMu.Unlock()
 
 	s.subscriptionsMu.Lock()

warehouse/retry.go

@@ -132,7 +132,7 @@ func (retryReq *RetryRequest) UploadsToRetry(ctx context.Context) (response Retr
 }
 
 func (retryReq *RetryRequest) getSourceIDs() []string {
-	return bcManager.SourceIDsByWorkspace()[retryReq.WorkspaceID] // TODO remove global variable
+	return retryReq.API.bcManager.SourceIDsByWorkspace()[retryReq.WorkspaceID]
 }
 
 func (retryReq *RetryRequest) clausesQuery(sourceIDs []string) []FilterClause {

warehouse/scheduling_test.go

@@ -20,7 +20,6 @@ func initWarehouse() {
 	Init2()
 	Init3()
 	Init4()
-	Init5()
 }
 
 var _ = Describe("Scheduling", func() {

warehouse/warehouse.go

@@ -88,7 +88,6 @@ var (
 	runningMode                         string
 	uploadStatusTrackFrequency          time.Duration
 	uploadAllocatorSleep                time.Duration
-	waitForConfig                       time.Duration
 	waitForWorkerSleep                  time.Duration
 	ShouldForceSetLowerVersion          bool
 	maxParallelJobCreation              int
@@ -181,7 +180,6 @@ func loadConfig() {
 	runningMode = config.GetString("Warehouse.runningMode", "")
 	config.RegisterDurationConfigVariable(30, &uploadStatusTrackFrequency, false, time.Minute, []string{"Warehouse.uploadStatusTrackFrequency", "Warehouse.uploadStatusTrackFrequencyInMin"}...)
 	config.RegisterDurationConfigVariable(5, &uploadAllocatorSleep, false, time.Second, []string{"Warehouse.uploadAllocatorSleep", "Warehouse.uploadAllocatorSleepInS"}...)
-	config.RegisterDurationConfigVariable(5, &waitForConfig, false, time.Second, []string{"Warehouse.waitForConfig", "Warehouse.waitForConfigInS"}...)
 	config.RegisterDurationConfigVariable(5, &waitForWorkerSleep, false, time.Second, []string{"Warehouse.waitForWorkerSleep", "Warehouse.waitForWorkerSleepInS"}...)
 	config.RegisterBoolConfigVariable(true, &ShouldForceSetLowerVersion, false, "SQLMigrator.forceSetLowerVersion")
 	config.RegisterIntConfigVariable(8, &maxParallelJobCreation, true, 1, "Warehouse.maxParallelJobCreation")
@@ -258,6 +256,10 @@ func (wh *HandleT) backendConfigSubscriber(ctx context.Context) {
 	for warehouses := range bcManager.Subscribe(ctx) {
 		wh.Logger.Info(`Received updated workspace config`)
 
+		warehouses = lo.Filter(warehouses, func(warehouse model.Warehouse, _ int) bool {
+			return warehouse.Destination.DestinationDefinition.Name == wh.destType
+		})
+
 		wh.configSubscriberLock.Lock()
 		wh.warehouses = warehouses
 		if wh.workspaceBySourceIDs == nil {
@@ -1453,11 +1455,13 @@ func Start(ctx context.Context, app app.App) error {
 		bcManager.Start(ctx)
 	})
 
+	RegisterAdmin(bcManager, pkgLogger)
+
 	runningMode := config.GetString("Warehouse.runningMode", "")
 	if runningMode == DegradedMode {
 		pkgLogger.Infof("WH: Running warehouse service in degraded mode...")
 		if isMaster() {
-			err := InitWarehouseAPI(dbHandle, pkgLogger.Child("upload_api"))
+			err := InitWarehouseAPI(dbHandle, bcManager, pkgLogger.Child("upload_api"))
 			if err != nil {
 				pkgLogger.Errorf("WH: Failed to start warehouse api: %v", err)
 				return err
@@ -1554,7 +1558,7 @@ func Start(ctx context.Context, app app.App) error {
 			return nil
 		}))
 
-		err := InitWarehouseAPI(dbHandle, pkgLogger.Child("upload_api"))
+		err := InitWarehouseAPI(dbHandle, bcManager, pkgLogger.Child("upload_api"))
 		if err != nil {
 			pkgLogger.Errorf("WH: Failed to start warehouse api: %v", err)
 			return err

warehouse/warehouse_test.go

@@ -63,7 +63,6 @@ func initWarehouse() {
 	Init2()
 	Init3()
 	Init4()
-	Init5()
 	validations.Init()
 	misc.Init()
 }

warehouse/warehousegrpc_test.go

@@ -660,7 +660,7 @@ func setupWarehouseGRPCTest(
 	}
 	resetBackendConfigManager()
 
-	require.NoError(t, InitWarehouseAPI(pgResource.DB, logger.NOP))
+	require.NoError(t, InitWarehouseAPI(pgResource.DB, bcManager, logger.NOP))
 
 	return pgResource, resetBackendConfigManager
 }