FIX: Add options for Azure AD and AWS IAM in configuration.postgresql/v1 (#3977)

Signed-off-by: aarondegroff <[email protected]>
Signed-off-by: Aaron DeGroff <[email protected]>
Co-authored-by: aarondegroff <[email protected]>
Co-authored-by: Yaron Schneider <[email protected]>

Author: 3 people

configuration/postgres/metadata.go

@@ -63,6 +63,11 @@ func (m *metadata) InitWithMetadata(meta map[string]string) error {
 		return fmt.Errorf("invalid table name '%s'. non-alphanumerics or upper cased table names are not supported", m.ConfigTable)
 	}
 
+	// Timeout
+	if m.Timeout < 1*time.Second {
+		return errors.New("invalid value for 'timeout': must be greater than 1s")
+	}
+
 	opts := pgauth.InitWithMetadataOpts{
 		AzureADEnabled: true,
 		AWSIAMEnabled:  true,

configuration/postgres/postgres.go

@@ -76,9 +76,24 @@ var (
 	allowedTableNameChars = regexp.MustCompile(`^[a-z0-9./_]*$`)
 )
 
+type Options struct {
+	// Disables support for authenticating with Azure AD
+	NoAzureAD bool
+
+	// Disables support for authenticating with AWS IAM
+	NoAWSIAM bool
+}
+
 func NewPostgresConfigurationStore(logger logger.Logger) configuration.Store {
+	return NewPostgresConfigurationStoreWithOptions(logger, Options{})
+}
+
+// NewPostgresConfigurationStoreWithOptions creates a new instance of PostgreSQL store with options.
+func NewPostgresConfigurationStoreWithOptions(logger logger.Logger, opts Options) configuration.Store {
 	return &ConfigurationStore{
 		logger: logger,
+		enableAzureAD: !opts.NoAzureAD,
+		enableAWSIAM:  !opts.NoAWSIAM,
 	}
 }
 
@@ -114,21 +129,20 @@ func (p *ConfigurationStore) Init(ctx context.Context, metadata configuration.Me
 		p.awsAuthProvider.UpdatePostgres(ctx, config)
 	}
 
-	pool, err := pgxpool.NewWithConfig(ctx, config)
+	connCtx, connCancel := context.WithTimeout(ctx, p.metadata.Timeout)
+	defer connCancel()
+	p.client, err = pgxpool.NewWithConfig(connCtx, config)
 	if err != nil {
 		return fmt.Errorf("PostgreSQL configuration store connection error: %w", err)
 	}
 
-	err = pool.Ping(ctx)
+	pingCtx, pingCancel := context.WithTimeout(ctx, p.metadata.Timeout)
+	defer pingCancel()
+	err = p.client.Ping(pingCtx)
 	if err != nil {
 		return fmt.Errorf("PostgreSQL configuration store ping error: %w", err)
 	}
-	p.client = pool
 
-	err = p.client.Ping(ctx)
-	if err != nil {
-		return fmt.Errorf("unable to connect to configuration store: '%w'", err)
-	}
 	// check if table exists
 	exists := false
 	err = p.client.QueryRow(ctx, QueryTableExists, p.metadata.ConfigTable).Scan(&exists)

state/postgresql/v2/metadata.go

@@ -60,7 +60,7 @@ func (m *pgMetadata) InitWithMetadata(meta state.Metadata, opts pgauth.InitWithM
 		return err
 	}
 
-	// Validate and sanitize inputq
+	// Validate and sanitize input
 	err = m.PostgresAuthMetadata.InitWithMetadata(meta.Properties, opts)
 	if err != nil {
 		return err