Merge branch 'rds-iam-instance-identifier-projection' into 'master'

Add RDS IAM instance identifier to config projection

See merge request postgres-ai/database-lab!1106

Author: agneum

engine/internal/rdsrefresh/dblab.go

@@ -16,6 +16,7 @@ import (
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/client/dblabapi"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/log"
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
+	"gitlab.com/postgres-ai/database-lab/v3/pkg/util/projection"
 )
 
 const (
@@ -170,19 +171,44 @@ func (c *DBLabClient) WaitForRefreshComplete(ctx context.Context, pollInterval,
 	}
 }
 
+// SourceConfigUpdate contains source database connection parameters for config update.
+type SourceConfigUpdate struct {
+	Host     string
+	Port     int
+	DBName   string
+	Username string
+	Password string
+	// RDSIAMDBInstance is the RDS DB instance identifier for IAM auth. When empty, this field is omitted from the config update.
+	RDSIAMDBInstance string
+}
+
 // UpdateSourceConfig updates the source database connection in DBLab config.
 // DBLab automatically reloads the configuration after the update.
-func (c *DBLabClient) UpdateSourceConfig(ctx context.Context, host string, port int, dbname, username, password string) error {
-	port64 := int64(port)
-	updateReq := models.ConfigProjection{
-		Host:     &host,
+func (c *DBLabClient) UpdateSourceConfig(ctx context.Context, update SourceConfigUpdate) error {
+	port64 := int64(update.Port)
+	proj := models.ConfigProjection{
+		Host:     &update.Host,
 		Port:     &port64,
-		DBName:   &dbname,
-		Username: &username,
-		Password: &password,
+		DBName:   &update.DBName,
+		Username: &update.Username,
+		Password: &update.Password,
+	}
+
+	if update.RDSIAMDBInstance != "" {
+		proj.RDSIAMDBInstance = &update.RDSIAMDBInstance
+	}
+
+	nested := map[string]interface{}{}
+
+	// defensive error check: StoreJSON only fails if target is not an addressable struct,
+	// which cannot happen here since proj is always a valid ConfigProjection value.
+	if err := projection.StoreJSON(&proj, nested, projection.StoreOptions{
+		Groups: []string{"default", "sensitive"},
+	}); err != nil {
+		return fmt.Errorf("failed to build config projection: %w", err)
 	}
 
-	bodyBytes, err := json.Marshal(updateReq)
+	bodyBytes, err := json.Marshal(nested)
 	if err != nil {
 		return fmt.Errorf("failed to marshal config update: %w", err)
 	}

engine/internal/rdsrefresh/dblab_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
+	"gitlab.com/postgres-ai/database-lab/v3/pkg/util/projection"
 )
 
 func TestDBLabClientHealth(t *testing.T) {
@@ -130,7 +131,13 @@ func TestDBLabClientUpdateSourceConfig(t *testing.T) {
 			assert.Equal(t, http.MethodPost, r.Method)
 			assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
 
-			err := json.NewDecoder(r.Body).Decode(&receivedConfig)
+			var nested map[string]interface{}
+			err := json.NewDecoder(r.Body).Decode(&nested)
+			require.NoError(t, err)
+
+			err = projection.LoadJSON(&receivedConfig, nested, projection.LoadOptions{
+				Groups: []string{"default", "sensitive"},
+			})
 			require.NoError(t, err)
 
 			w.WriteHeader(http.StatusOK)
@@ -140,14 +147,48 @@ func TestDBLabClientUpdateSourceConfig(t *testing.T) {
 		client, err := NewDBLabClient(&DBLabConfig{APIEndpoint: server.URL, Token: "test-token"})
 		require.NoError(t, err)
 
-		err = client.UpdateSourceConfig(context.Background(), "clone-host.rds.amazonaws.com", 5432, "postgres", "dbuser", "dbpass")
+		err = client.UpdateSourceConfig(context.Background(), SourceConfigUpdate{
+			Host: "clone-host.rds.amazonaws.com", Port: 5432, DBName: "postgres",
+			Username: "dbuser", Password: "dbpass", RDSIAMDBInstance: "my-rds-clone",
+		})
 		require.NoError(t, err)
 
 		assert.Equal(t, "clone-host.rds.amazonaws.com", *receivedConfig.Host)
 		assert.Equal(t, int64(5432), *receivedConfig.Port)
 		assert.Equal(t, "postgres", *receivedConfig.DBName)
 		assert.Equal(t, "dbuser", *receivedConfig.Username)
 		assert.Equal(t, "dbpass", *receivedConfig.Password)
+		assert.Equal(t, "my-rds-clone", *receivedConfig.RDSIAMDBInstance)
+	})
+
+	t.Run("successful without rds iam instance", func(t *testing.T) {
+		var receivedConfig models.ConfigProjection
+
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			var nested map[string]interface{}
+			err := json.NewDecoder(r.Body).Decode(&nested)
+			require.NoError(t, err)
+
+			err = projection.LoadJSON(&receivedConfig, nested, projection.LoadOptions{
+				Groups: []string{"default", "sensitive"},
+			})
+			require.NoError(t, err)
+
+			w.WriteHeader(http.StatusOK)
+		}))
+		defer server.Close()
+
+		client, err := NewDBLabClient(&DBLabConfig{APIEndpoint: server.URL, Token: "test-token"})
+		require.NoError(t, err)
+
+		err = client.UpdateSourceConfig(context.Background(), SourceConfigUpdate{
+			Host: "host.rds.amazonaws.com", Port: 5432, DBName: "postgres",
+			Username: "dbuser", Password: "dbpass",
+		})
+		require.NoError(t, err)
+
+		assert.Equal(t, "host.rds.amazonaws.com", *receivedConfig.Host)
+		assert.Nil(t, receivedConfig.RDSIAMDBInstance)
 	})
 
 	t.Run("error on non-2xx status", func(t *testing.T) {
@@ -160,7 +201,9 @@ func TestDBLabClientUpdateSourceConfig(t *testing.T) {
 		client, err := NewDBLabClient(&DBLabConfig{APIEndpoint: server.URL, Token: "test-token"})
 		require.NoError(t, err)
 
-		err = client.UpdateSourceConfig(context.Background(), "host", 5432, "db", "user", "pass")
+		err = client.UpdateSourceConfig(context.Background(), SourceConfigUpdate{
+			Host: "host", Port: 5432, DBName: "db", Username: "user", Password: "pass",
+		})
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "invalid configuration")
 	})
@@ -175,7 +218,9 @@ func TestDBLabClientUpdateSourceConfig(t *testing.T) {
 		client, err := NewDBLabClient(&DBLabConfig{APIEndpoint: server.URL, Token: "test-token"})
 		require.NoError(t, err)
 
-		err = client.UpdateSourceConfig(context.Background(), "host", 5432, "db", "user", "pass")
+		err = client.UpdateSourceConfig(context.Background(), SourceConfigUpdate{
+			Host: "host", Port: 5432, DBName: "db", Username: "user", Password: "pass",
+		})
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "internal server error")
 	})

engine/internal/rdsrefresh/refresher.go

@@ -180,14 +180,14 @@ func (r *Refresher) Run(ctx context.Context) *RefreshResult {
 	// step 6: update DBLab config with RDS clone endpoint
 	log.Msg("updating DBLab config...")
 
-	if err := r.dblab.UpdateSourceConfig(
-		ctx,
-		clone.Endpoint,
-		int(clone.Port),
-		r.cfg.Source.DBName,
-		r.cfg.Source.Username,
-		r.cfg.Source.Password,
-	); err != nil {
+	if err := r.dblab.UpdateSourceConfig(ctx, SourceConfigUpdate{
+		Host:             clone.Endpoint,
+		Port:             int(clone.Port),
+		DBName:           r.cfg.Source.DBName,
+		Username:         r.cfg.Source.Username,
+		Password:         r.cfg.Source.Password,
+		RDSIAMDBInstance: clone.Identifier,
+	}); err != nil {
 		result.Error = fmt.Errorf("failed to update DBLab config: %w", err)
 		return result
 	}

engine/internal/rdsrefresh/refresher_test.go

@@ -12,10 +12,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/aws"
+	awsrds "github.com/aws/aws-sdk-go-v2/service/rds"
+	"github.com/aws/aws-sdk-go-v2/service/rds/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"gitlab.com/postgres-ai/database-lab/v3/pkg/models"
+	"gitlab.com/postgres-ai/database-lab/v3/pkg/util/projection"
 )
 
 func TestRefreshResult(t *testing.T) {
@@ -205,3 +209,93 @@ func TestRefreshWorkflow(t *testing.T) {
 		assert.LessOrEqual(t, result.Duration, elapsed+100*time.Millisecond)
 	})
 }
+
+func TestRefreshWorkflowPassesCloneIdentifier(t *testing.T) {
+	t.Run("passes clone identifier as rds iam db instance", func(t *testing.T) {
+		var receivedConfig models.ConfigProjection
+
+		statusCallCount := 0
+		refreshTriggered := false
+
+		dblabServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			switch r.URL.Path {
+			case "/healthz":
+				w.WriteHeader(http.StatusOK)
+				_, _ = w.Write([]byte(`{"status":"ok"}`))
+			case "/status":
+				statusCallCount++
+
+				var status models.RetrievalStatus
+
+				switch {
+				case !refreshTriggered || statusCallCount <= 2:
+					status = models.Finished
+				case statusCallCount == 3:
+					status = models.Refreshing
+				default:
+					status = models.Finished
+				}
+
+				resp := &models.InstanceStatus{Retrieving: models.Retrieving{Status: status}}
+				_ = json.NewEncoder(w).Encode(resp)
+			case "/admin/config":
+				var nested map[string]interface{}
+				require.NoError(t, json.NewDecoder(r.Body).Decode(&nested))
+				require.NoError(t, projection.LoadJSON(&receivedConfig, nested, projection.LoadOptions{Groups: []string{"default", "sensitive"}}))
+				w.WriteHeader(http.StatusOK)
+			case "/full-refresh":
+				refreshTriggered = true
+				statusCallCount = 0
+				_, _ = w.Write([]byte(`{"status":"OK"}`))
+			default:
+				w.WriteHeader(http.StatusNotFound)
+			}
+		}))
+		defer dblabServer.Close()
+
+		mock := &mockRDSAPI{
+			describeDBInstancesFunc: func(ctx context.Context, params *awsrds.DescribeDBInstancesInput, optFns ...func(*awsrds.Options)) (*awsrds.DescribeDBInstancesOutput, error) {
+				return &awsrds.DescribeDBInstancesOutput{
+					DBInstances: []types.DBInstance{{
+						DBInstanceIdentifier: aws.String(aws.ToString(params.DBInstanceIdentifier)),
+						DBInstanceStatus:     aws.String("available"),
+						Engine:               aws.String("postgres"),
+						EngineVersion:        aws.String("15.4"),
+						Endpoint:             &types.Endpoint{Address: aws.String("clone.rds.amazonaws.com"), Port: aws.Int32(5432)},
+					}},
+				}, nil
+			},
+			restoreDBInstanceFunc: func(ctx context.Context, params *awsrds.RestoreDBInstanceFromDBSnapshotInput, optFns ...func(*awsrds.Options)) (*awsrds.RestoreDBInstanceFromDBSnapshotOutput, error) {
+				return &awsrds.RestoreDBInstanceFromDBSnapshotOutput{}, nil
+			},
+			modifyDBInstanceFunc: func(ctx context.Context, params *awsrds.ModifyDBInstanceInput, optFns ...func(*awsrds.Options)) (*awsrds.ModifyDBInstanceOutput, error) {
+				return &awsrds.ModifyDBInstanceOutput{}, nil
+			},
+			deleteDBInstanceFunc: func(ctx context.Context, params *awsrds.DeleteDBInstanceInput, optFns ...func(*awsrds.Options)) (*awsrds.DeleteDBInstanceOutput, error) {
+				return &awsrds.DeleteDBInstanceOutput{}, nil
+			},
+		}
+
+		cfg := &Config{
+			Source:   SourceConfig{Type: "rds", Identifier: "test-db", DBName: "postgres", Username: "dbuser", Password: "dbpass", SnapshotIdentifier: "snap-123"},
+			RDSClone: RDSCloneConfig{InstanceClass: "db.t3.medium"},
+			DBLab:    DBLabConfig{APIEndpoint: dblabServer.URL, Token: "test-token", PollInterval: Duration(100 * time.Millisecond), Timeout: Duration(5 * time.Second)},
+		}
+
+		rdsClient := NewRDSClientWithAPI(mock, cfg)
+		dblabClient, err := NewDBLabClient(&cfg.DBLab)
+		require.NoError(t, err)
+
+		refresher := &Refresher{cfg: cfg, rds: rdsClient, dblab: dblabClient}
+		result := refresher.Run(context.Background())
+
+		require.True(t, result.Success, "expected success but got error: %v", result.Error)
+		require.NotNil(t, receivedConfig.RDSIAMDBInstance, "RDSIAMDBInstance should be set in config update")
+		assert.Contains(t, *receivedConfig.RDSIAMDBInstance, cloneNamePrefix)
+		assert.Equal(t, "clone.rds.amazonaws.com", *receivedConfig.Host)
+		assert.Equal(t, int64(5432), *receivedConfig.Port)
+		assert.Equal(t, "postgres", *receivedConfig.DBName)
+		assert.Equal(t, "dbuser", *receivedConfig.Username)
+		assert.Equal(t, "dbpass", *receivedConfig.Password)
+	})
+}

engine/pkg/models/configuration.go

@@ -28,4 +28,5 @@ type ConfigProjection struct {
 	RestoreCustomOptions []interface{}          `proj:"retrieval.spec.logicalRestore.options.customOptions"`
 	IgnoreDumpErrors     *bool                  `proj:"retrieval.spec.logicalDump.options.ignoreErrors"`
 	IgnoreRestoreErrors  *bool                  `proj:"retrieval.spec.logicalRestore.options.ignoreErrors"`
+	RDSIAMDBInstance     *string                `proj:"retrieval.spec.logicalDump.options.source.rdsIam.dbInstanceIdentifier"`
 }