chore: warehouse backend config refactoring

warehouse/backend-config.go → warehouse/backend_config.go

@@ -5,6 +5,8 @@ import (
 	"encoding/json"
 	"sync"
 
+	"github.com/samber/lo"
+
 	"golang.org/x/exp/slices"
 
 	"github.com/rudderlabs/rudder-go-kit/config"
@@ -20,32 +22,32 @@ import (
 
 // TODO: add tests
 func newBackendConfigManager(
-	c *config.Config, // TODO possibly use this to get all the needed variables
-	wrappedDB *sqlquerywrapper.DB,
+	c *config.Config,
+	db *sqlquerywrapper.DB,
 	bc backendconfig.BackendConfig,
-	l logger.Logger,
+	log logger.Logger,
 ) *backendConfigManager {
 	if c == nil {
 		c = config.Default
 	}
 	if bc == nil {
 		bc = backendconfig.DefaultBackendConfig
 	}
-	if l == nil {
-		l = logger.NOP
+	if log == nil {
+		log = logger.NOP
 	}
 	bcm := &backendConfigManager{
 		conf:                 c,
-		db:                   wrappedDB,
-		schema:               repo.NewWHSchemas(wrappedDB),
+		db:                   db,
+		schema:               repo.NewWHSchemas(db),
 		backendConfig:        bc,
-		logger:               l,
+		logger:               log,
 		initialConfigFetched: make(chan struct{}),
 		connectionsMap:       make(map[string]map[string]model.Warehouse),
 	}
-	if config.GetBool("ENABLE_TUNNELLING", true) {
+	if c.GetBool("ENABLE_TUNNELLING", true) {
 		bcm.internalControlPlaneClient = cpclient.NewInternalClientWithCache(
-			backendconfig.GetConfigBackendURL(),
+			c.GetString("CONFIG_BACKEND_URL", "https://api.rudderstack.com"),
 			cpclient.BasicAuth{
 				Username: c.GetString("CP_INTERNAL_API_USERNAME", ""),
 				Password: c.GetString("CP_INTERNAL_API_PASSWORD", ""),
@@ -66,8 +68,9 @@ type backendConfigManager struct {
 
 	initialConfigFetched          chan struct{}
 	closeInitialConfigFetchedOnce sync.Once
-	subscriptions                 []chan []model.Warehouse
-	subscriptionsMu               sync.Mutex
+
+	subscriptions   []chan []model.Warehouse
+	subscriptionsMu sync.Mutex
 
 	// variables to store the backend configuration
 	warehouses   []model.Warehouse
@@ -86,7 +89,10 @@ func (s *backendConfigManager) Start(ctx context.Context) {
 		select {
 		case <-ctx.Done():
 			return
-		case data := <-ch:
+		case data, ok := <-ch:
+			if !ok {
+				return
+			}
 			s.processData(ctx, data.Data.(map[string]backendconfig.ConfigT))
 		}
 	}
@@ -246,26 +252,20 @@ func (s *backendConfigManager) SourceIDsByWorkspace() map[string][]string {
 func (s *backendConfigManager) WarehousesBySourceID(sourceID string) []model.Warehouse {
 	s.warehousesMu.RLock()
 	defer s.warehousesMu.RUnlock()
-	var warehouses []model.Warehouse
-	for _, wh := range s.warehouses {
-		if wh.Source.ID == sourceID {
-			warehouses = append(warehouses, wh)
-		}
-	}
-	return warehouses
+
+	return lo.Filter(s.warehouses, func(w model.Warehouse, _ int) bool {
+		return w.Source.ID == sourceID
+	})
 }
 
 // WarehousesByDestID gets all WHs for the given destination ID
 func (s *backendConfigManager) WarehousesByDestID(destID string) []model.Warehouse {
 	s.warehousesMu.RLock()
 	defer s.warehousesMu.RUnlock()
-	var warehouses []model.Warehouse
-	for _, wh := range s.warehouses {
-		if wh.Destination.ID == destID {
-			warehouses = append(warehouses, wh)
-		}
-	}
-	return warehouses
+
+	return lo.Filter(s.warehouses, func(w model.Warehouse, _ int) bool {
+		return w.Destination.ID == destID
+	})
 }
 
 func (s *backendConfigManager) attachSSHTunnellingInfo(
@@ -277,7 +277,8 @@ func (s *backendConfigManager) attachSSHTunnellingInfo(
 		return upstream
 	}
 
-	pkgLogger.Debugf("Fetching ssh keys for destination: %s", upstream.ID)
+	s.logger.Debugf("Fetching ssh keys for destination: %s", upstream.ID)
+
 	keys, err := s.internalControlPlaneClient.GetDestinationSSHKeys(ctx, upstream.ID)
 	if err != nil {
 		s.logger.Errorf("fetching ssh keys for destination: %s", err.Error())

warehouse/backend_config_test.go

@@ -0,0 +1,247 @@
+package warehouse
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/golang/mock/gomock"
+	"github.com/ory/dockertest/v3"
+	"github.com/stretchr/testify/require"
+
+	"github.com/rudderlabs/rudder-go-kit/config"
+	"github.com/rudderlabs/rudder-go-kit/logger"
+	"github.com/rudderlabs/rudder-go-kit/testhelper/docker/resource"
+	backendconfig "github.com/rudderlabs/rudder-server/backend-config"
+	mocksBackendConfig "github.com/rudderlabs/rudder-server/mocks/backend-config"
+	"github.com/rudderlabs/rudder-server/utils/pubsub"
+	"github.com/rudderlabs/rudder-server/warehouse/integrations/middleware/sqlquerywrapper"
+	"github.com/rudderlabs/rudder-server/warehouse/internal/model"
+	warehouseutils "github.com/rudderlabs/rudder-server/warehouse/utils"
+)
+
+func TestBackendConfigManager(t *testing.T) {
+	pool, err := dockertest.NewPool("")
+	require.NoError(t, err)
+
+	pgResource, err := resource.SetupPostgres(pool, t)
+	require.NoError(t, err)
+
+	workspaceID := "test-workspace-id"
+	sourceID := "test-source-id"
+	destinationID := "test-destination-id"
+	namespace := "test-namespace"
+
+	unknownSourceID := "unknown-source-id"
+	unknownDestinationID := "unknown-destination-id"
+	unknownDestinationName := "unknown-destination-name"
+
+	ctrl := gomock.NewController(t)
+	mockBackendConfig := mocksBackendConfig.NewMockBackendConfig(ctrl)
+	mockBackendConfig.EXPECT().Subscribe(gomock.Any(), backendconfig.TopicBackendConfig).DoAndReturn(func(ctx context.Context, topic backendconfig.Topic) pubsub.DataChannel {
+		ch := make(chan pubsub.DataEvent, 1)
+		ch <- pubsub.DataEvent{
+			Data: map[string]backendconfig.ConfigT{
+				workspaceID: {
+					WorkspaceID: workspaceID,
+					Sources: []backendconfig.SourceT{
+						{
+							ID:      sourceID,
+							Enabled: true,
+							Destinations: []backendconfig.DestinationT{
+								{
+									ID:      destinationID,
+									Enabled: true,
+									DestinationDefinition: backendconfig.DestinationDefinitionT{
+										Name: warehouseutils.RS,
+									},
+									Config: map[string]interface{}{
+										"namespace": namespace,
+									},
+								},
+							},
+						},
+						{
+							ID:      unknownSourceID,
+							Enabled: true,
+							Destinations: []backendconfig.DestinationT{
+								{
+									ID:      unknownDestinationID,
+									Enabled: true,
+									DestinationDefinition: backendconfig.DestinationDefinitionT{
+										Name: unknownDestinationName,
+									},
+									Config: map[string]interface{}{},
+								},
+							},
+						},
+					},
+				},
+			},
+			Topic: string(backendconfig.TopicBackendConfig),
+		}
+		close(ch)
+		return ch
+	}).AnyTimes()
+
+	svc := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusOK)
+		_, _ = w.Write([]byte(`{"publicKey": "public_key", "privateKey": "private_key"}`))
+	}))
+	defer svc.Close()
+
+	c := config.New()
+	c.Set("CONFIG_BACKEND_URL", svc.URL)
+	c.Set("CP_INTERNAL_API_USERNAME", "username")
+	c.Set("CP_INTERNAL_API_PASSWORD", "password")
+
+	db := sqlquerywrapper.New(
+		pgResource.DB,
+		sqlquerywrapper.WithLogger(logger.NOP),
+	)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	bcm := newBackendConfigManager(c, db, mockBackendConfig, logger.NOP)
+
+	t.Run("Subscriptions", func(t *testing.T) {
+		require.False(t, bcm.IsInitialized())
+		require.Equal(t, bcm.Connections(), map[string]map[string]model.Warehouse{})
+
+		csm, ok := bcm.ConnectionSourcesMap(destinationID)
+		require.False(t, ok)
+
+		require.Nil(t, csm)
+		require.Nil(t, bcm.SourceIDsByWorkspace())
+		require.Equal(t, bcm.WarehousesBySourceID(sourceID), []model.Warehouse{})
+		require.Equal(t, bcm.WarehousesByDestID(destinationID), []model.Warehouse{})
+
+		subscriptionsCh := bcm.Subscribe(ctx)
+
+		go func() {
+			bcm.Start(ctx)
+		}()
+
+		expectedWarehouse := model.Warehouse{
+			WorkspaceID: workspaceID,
+			Source: backendconfig.SourceT{
+				ID:      sourceID,
+				Enabled: true,
+				Destinations: []backendconfig.DestinationT{
+					{
+						ID:      destinationID,
+						Enabled: true,
+						DestinationDefinition: backendconfig.DestinationDefinitionT{
+							Name: warehouseutils.RS,
+						},
+						Config: map[string]interface{}{
+							"namespace": namespace,
+						},
+					},
+				},
+			},
+			Destination: backendconfig.DestinationT{
+				ID:      destinationID,
+				Enabled: true,
+				DestinationDefinition: backendconfig.DestinationDefinitionT{
+					Name: warehouseutils.RS,
+				},
+				Config: map[string]interface{}{
+					"namespace": namespace,
+				},
+			},
+			Namespace:  "test_namespace",
+			Type:       warehouseutils.RS,
+			Identifier: "RS:test-source-id:test-destination-id",
+		}
+
+		require.Equal(t, <-subscriptionsCh, []model.Warehouse{expectedWarehouse})
+		require.True(t, bcm.IsInitialized())
+		require.Equal(t, bcm.Connections(), map[string]map[string]model.Warehouse{
+			destinationID: {
+				sourceID: expectedWarehouse,
+			},
+		})
+
+		csm, ok = bcm.ConnectionSourcesMap(destinationID)
+		require.True(t, ok)
+		require.Equal(t, csm, map[string]model.Warehouse{
+			sourceID: expectedWarehouse,
+		})
+
+		require.Equal(t, bcm.SourceIDsByWorkspace(), map[string][]string{
+			workspaceID: {sourceID, unknownSourceID},
+		})
+		require.Equal(t, bcm.WarehousesBySourceID(sourceID), []model.Warehouse{expectedWarehouse})
+		require.Equal(t, bcm.WarehousesByDestID(destinationID), []model.Warehouse{expectedWarehouse})
+		require.Equal(t, bcm.WarehousesBySourceID(unknownSourceID), []model.Warehouse{})
+		require.Equal(t, bcm.WarehousesByDestID(unknownDestinationID), []model.Warehouse{})
+	})
+
+	t.Run("Tunnelling", func(t *testing.T) {
+		testCases := []struct {
+			name     string
+			input    backendconfig.DestinationT
+			expected backendconfig.DestinationT
+		}{
+			{
+				name: "tunnelling disabled",
+				input: backendconfig.DestinationT{
+					ID:      destinationID,
+					Enabled: true,
+					DestinationDefinition: backendconfig.DestinationDefinitionT{
+						Name: warehouseutils.RS,
+					},
+					Config: map[string]interface{}{
+						"namespace": namespace,
+					},
+				},
+				expected: backendconfig.DestinationT{
+					ID:      destinationID,
+					Enabled: true,
+					DestinationDefinition: backendconfig.DestinationDefinitionT{
+						Name: warehouseutils.RS,
+					},
+					Config: map[string]interface{}{
+						"namespace": namespace,
+					},
+				},
+			},
+			{
+				name: "tunnelling enabled",
+				input: backendconfig.DestinationT{
+					ID:      destinationID,
+					Enabled: true,
+					DestinationDefinition: backendconfig.DestinationDefinitionT{
+						Name: warehouseutils.RS,
+					},
+					Config: map[string]interface{}{
+						"namespace": namespace,
+						"useSSH":    true,
+					},
+				},
+				expected: backendconfig.DestinationT{
+					ID:      destinationID,
+					Enabled: true,
+					DestinationDefinition: backendconfig.DestinationDefinitionT{
+						Name: warehouseutils.RS,
+					},
+					Config: map[string]interface{}{
+						"namespace":     namespace,
+						"useSSH":        true,
+						"sshPrivateKey": "private_key",
+					},
+				},
+			},
+		}
+
+		for _, tc := range testCases {
+			tc := tc
+			t.Run(tc.name, func(t *testing.T) {
+				require.Equal(t, bcm.attachSSHTunnellingInfo(ctx, tc.input), tc.expected)
+			})
+		}
+	})
+}