Support of ALTER ROLE ... SET ROLE ... (#209)

Author: kostiantyn-nemchenko

postgresql/resource_postgresql_role.go

@@ -34,6 +34,7 @@ const (
 	roleRolesAttr                           = "roles"
 	roleSearchPathAttr                      = "search_path"
 	roleStatementTimeoutAttr                = "statement_timeout"
+	roleAssumeRoleAttr                      = "assume_role"
 
 	// Deprecated options
 	roleDepEncryptedAttr = "encrypted"
@@ -167,6 +168,11 @@ func resourcePostgreSQLRole() *schema.Resource {
 				Description:  "Abort any statement that takes more than the specified number of milliseconds",
 				ValidateFunc: validation.IntAtLeast(0),
 			},
+			roleAssumeRoleAttr: {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Description: "Role to switch to at login",
+			},
 		},
 	}
 }
@@ -301,6 +307,10 @@ func resourcePostgreSQLRoleCreate(db *DBConnection, d *schema.ResourceData) erro
 		return err
 	}
 
+	if err = setAssumeRole(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return fmt.Errorf("could not commit transaction: %w", err)
 	}
@@ -446,6 +456,7 @@ func resourcePostgreSQLRoleReadImpl(db *DBConnection, d *schema.ResourceData) er
 	d.Set(roleBypassRLSAttr, roleBypassRLS)
 	d.Set(roleRolesAttr, pgArrayToSet(roleRoles))
 	d.Set(roleSearchPathAttr, readSearchPath(roleConfig))
+	d.Set(roleAssumeRoleAttr, readAssumeRole(roleConfig))
 
 	statementTimeout, err := readStatementTimeout(roleConfig)
 	if err != nil {
@@ -522,6 +533,20 @@ func readStatementTimeout(roleConfig pq.ByteaArray) (int, error) {
 	return 0, nil
 }
 
+// readAssumeRole searches for a role entry in the rolconfig array.
+// In case no such value is present, it returns empty string.
+func readAssumeRole(roleConfig pq.ByteaArray) string {
+	var res string
+	var assumeRoleAttr = "role"
+	for _, v := range roleConfig {
+		config := string(v)
+		if strings.HasPrefix(config, assumeRoleAttr) {
+			res = strings.TrimPrefix(config, assumeRoleAttr+"=")
+		}
+	}
+	return res
+}
+
 // readRolePassword reads password either from Postgres if admin user is a superuser
 // or only from Terraform state.
 func readRolePassword(db *DBConnection, d *schema.ResourceData, roleCanLogin bool) (string, error) {
@@ -660,6 +685,10 @@ func resourcePostgreSQLRoleUpdate(db *DBConnection, d *schema.ResourceData) erro
 		return err
 	}
 
+	if err = setAssumeRole(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return fmt.Errorf("could not commit transaction: %w", err)
 	}
@@ -1008,3 +1037,28 @@ func setIdleInTransactionSessionTimeout(txn *sql.Tx, d *schema.ResourceData) err
 	}
 	return nil
 }
+
+func setAssumeRole(txn *sql.Tx, d *schema.ResourceData) error {
+	if !d.HasChange(roleAssumeRoleAttr) {
+		return nil
+	}
+
+	roleName := d.Get(roleNameAttr).(string)
+	assumeRole := d.Get(roleAssumeRoleAttr).(string)
+	if assumeRole != "" {
+		sql := fmt.Sprintf(
+			"ALTER ROLE %s SET ROLE TO %s", pq.QuoteIdentifier(roleName), pq.QuoteIdentifier(assumeRole),
+		)
+		if _, err := txn.Exec(sql); err != nil {
+			return fmt.Errorf("could not set role %s for %s: %w", assumeRole, roleName, err)
+		}
+	} else {
+		sql := fmt.Sprintf(
+			"ALTER ROLE %s RESET ROLE", pq.QuoteIdentifier(roleName),
+		)
+		if _, err := txn.Exec(sql); err != nil {
+			return fmt.Errorf("could not reset role for %s: %w", roleName, err)
+		}
+	}
+	return nil
+}

postgresql/resource_postgresql_role_test.go

@@ -46,6 +46,7 @@ func TestAccPostgresqlRole_Basic(t *testing.T) {
 					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "skip_reassign_owned", "false"),
 					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "statement_timeout", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "idle_in_transaction_session_timeout", "0"),
+					resource.TestCheckResourceAttr("postgresql_role.role_with_defaults", "assume_role", ""),
 
 					resource.TestCheckResourceAttr("postgresql_role.role_with_create_database", "name", "role_with_create_database"),
 					resource.TestCheckResourceAttr("postgresql_role.role_with_create_database", "create_database", "true"),
@@ -120,6 +121,7 @@ resource "postgresql_role" "update_role" {
   search_path = ["mysearchpath"]
   statement_timeout = 30000
   idle_in_transaction_session_timeout = 60000
+  assume_role = "${postgresql_role.group_role.name}"
 }
 `
 	resource.Test(t, resource.TestCase{
@@ -143,6 +145,7 @@ resource "postgresql_role" "update_role" {
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.#", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "statement_timeout", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "idle_in_transaction_session_timeout", "0"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "assume_role", ""),
 					testAccCheckRoleCanLogin(t, "update_role", "toto"),
 				),
 			},
@@ -163,6 +166,7 @@ resource "postgresql_role" "update_role" {
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.0", "mysearchpath"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "statement_timeout", "30000"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "idle_in_transaction_session_timeout", "60000"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "assume_role", "group_role"),
 					testAccCheckRoleCanLogin(t, "update_role2", "titi"),
 				),
 			},
@@ -180,6 +184,7 @@ resource "postgresql_role" "update_role" {
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "search_path.#", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "statement_timeout", "0"),
 					resource.TestCheckResourceAttr("postgresql_role.update_role", "idle_in_transaction_session_timeout", "0"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "assume_role", ""),
 					testAccCheckRoleCanLogin(t, "update_role", "toto"),
 				),
 			},
@@ -412,6 +417,7 @@ resource "postgresql_role" "role_with_defaults" {
   valid_until = "infinity"
   statement_timeout = 0
   idle_in_transaction_session_timeout = 0
+  assume_role = ""
 }
 
 resource "postgresql_role" "role_with_create_database" {

website/docs/r/postgresql_role.html.markdown

@@ -116,6 +116,8 @@ resource "postgresql_role" "my_replication_role" {
 
 * `statement_timeout` - (Optional) Defines [`statement_timeout`](https://www.postgresql.org/docs/current/runtime-config-client.html#RUNTIME-CONFIG-CLIENT-STATEMENT) setting for this role which allows to abort any statement that takes more than the specified amount of time.
 
+* `assume_role` - (Optional) Defines the role to switch to at login via [`SET ROLE`](https://www.postgresql.org/docs/current/sql-set-role.html).
+
 ## Import Example
 
 `postgresql_role` supports importing resources.  Supposing the following