Create roles with Azure managed identities in Azure PostgreSQL Flexible Servers

README.md

@@ -79,3 +79,11 @@ TF_LOG=INFO go test -v ./postgresql -run ^TestAccPostgresqlRole_Basic$
 # cleans the env and tears down the postgres container
 make testacc_cleanup 
 ```
+
+The Azure identity role test can only run against an actual Azure postgres instance, also only a Microsoft Entra administrator is allowed to manage Microsoft Entra roles
+In order to run the Azure identity test, run the following commands:
+```sh
+source tests/switch_azure_ad_user.sh
+
+TF_LOG=INFO go test -v ./postgresql -run ^TestAccPostgresqlRole_AzureIdentity
+```

postgresql/resource_postgresql_role.go

@@ -35,6 +35,8 @@ const (
 	roleSearchPathAttr                      = "search_path"
 	roleStatementTimeoutAttr                = "statement_timeout"
 	roleAssumeRoleAttr                      = "assume_role"
+	roleAzureIdentityIdAttr                 = "azure_identity_id"
+	roleAzureIdentityTypeAttr               = "azure_identity_type"
 
 	// Deprecated options
 	roleDepEncryptedAttr = "encrypted"
@@ -173,6 +175,17 @@ func resourcePostgreSQLRole() *schema.Resource {
 				Optional:    true,
 				Description: "Role to switch to at login",
 			},
+			roleAzureIdentityIdAttr: {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Description: "Allow usage of role via this Azure identity",
+			},
+			roleAzureIdentityTypeAttr: {
+				Type:         schema.TypeString,
+				Optional:     true,
+				Description:  "Type of Azure identity: user, group or service",
+				ValidateFunc: validation.StringInSlice([]string{"user", "group", "service"}, false),
+			},
 		},
 	}
 }
@@ -291,6 +304,10 @@ func resourcePostgreSQLRoleCreate(db *DBConnection, d *schema.ResourceData) erro
 		return fmt.Errorf("error creating role %s: %w", roleName, err)
 	}
 
+	if err = azureIdentity(txn, d); err != nil {
+		return err
+	}
+
 	if err = grantRoles(txn, d); err != nil {
 		return err
 	}
@@ -474,12 +491,14 @@ func resourcePostgreSQLRoleReadImpl(db *DBConnection, d *schema.ResourceData) er
 
 	d.SetId(roleName)
 
-	password, err := readRolePassword(db, d, roleCanLogin)
-	if err != nil {
-		return err
-	}
+	if d.Get(roleAzureIdentityIdAttr).(string) == "" {
+		password, err := readRolePassword(db, d, roleCanLogin)
+		if err != nil {
+			return err
+		}
 
-	d.Set(rolePasswordAttr, password)
+		d.Set(rolePasswordAttr, password)
+	}
 	return nil
 }
 
@@ -947,6 +966,21 @@ func revokeRoles(txn *sql.Tx, d *schema.ResourceData) error {
 	return nil
 }
 
+func azureIdentity(txn *sql.Tx, d *schema.ResourceData) error {
+	role := d.Get(roleNameAttr).(string)
+	azureIdentityId := d.Get(roleAzureIdentityIdAttr).(string)
+	azureIdentityType := d.Get(roleAzureIdentityTypeAttr).(string)
+	if azureIdentityId != "" {
+		query := fmt.Sprintf(
+			"SECURITY LABEL for \"pgaadauth\" on role %s is 'aadauth,oid=%s,type=%s'", pq.QuoteIdentifier(role), azureIdentityId, azureIdentityType,
+		)
+		if _, err := txn.Exec(query); err != nil {
+			return fmt.Errorf("could not set identity %s on role %s %w", azureIdentityId, role, err)
+		}
+	}
+	return nil
+}
+
 func grantRoles(txn *sql.Tx, d *schema.ResourceData) error {
 	role := d.Get(roleNameAttr).(string)
 

postgresql/resource_postgresql_role_test.go

@@ -228,6 +228,39 @@ resource "postgresql_role" "test_role" {
 	})
 }
 
+func TestAccPostgresqlRole_AzureIdentity(t *testing.T) {
+	// This test needs an Azure PostgreSQL Flexible instance and an existing identity
+	if os.Getenv("AZURE") == "" {
+		t.Skip("Skipping azure identity test")
+	}
+	roleConfig := `
+resource "postgresql_role" "id-test" {
+  name = "id-test"
+  azure_identity_id = "00000000-0000-0000-0000-000000000000"
+  azure_identity_type = "service"
+  skip_reassign_owned = true
+  login = true
+}`
+	resource.Test(t, resource.TestCase{
+		PreCheck: func() {
+			testAccPreCheck(t)
+			testCheckCompatibleVersion(t, featurePrivileges)
+		},
+		Providers:    testAccProviders,
+		CheckDestroy: testAccCheckPostgresqlRoleDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: roleConfig,
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckPostgresqlRoleAzureExists("id-test", nil, nil),
+					testAccCheckPostgresqlRoleExists("id-test", nil, nil),
+					resource.TestCheckResourceAttr("postgresql_role.id-test", "name", "id-test"),
+				),
+			},
+		},
+	})
+}
+
 func testAccCheckPostgresqlRoleDestroy(s *terraform.State) error {
 	client := testAccProvider.Meta().(*Client)
 
@@ -250,6 +283,23 @@ func testAccCheckPostgresqlRoleDestroy(s *terraform.State) error {
 	return nil
 }
 
+func testAccCheckPostgresqlRoleAzureExists(roleName string, grantedRoles []string, searchPath []string) resource.TestCheckFunc {
+	return func(s *terraform.State) error {
+		client := testAccProvider.Meta().(*Client)
+
+		exists, err := checkRoleAzureExists(client, roleName)
+		if err != nil {
+			return fmt.Errorf("Error checking role %s", err)
+		}
+
+		if !exists {
+			return fmt.Errorf("Azure role not found")
+		}
+
+		return nil
+	}
+}
+
 func testAccCheckPostgresqlRoleExists(roleName string, grantedRoles []string, searchPath []string) resource.TestCheckFunc {
 	return func(s *terraform.State) error {
 		client := testAccProvider.Meta().(*Client)
@@ -278,6 +328,22 @@ func testAccCheckPostgresqlRoleExists(roleName string, grantedRoles []string, se
 	}
 }
 
+func checkRoleAzureExists(client *Client, roleName string) (bool, error) {
+	db, err := client.Connect()
+	if err != nil {
+		return false, err
+	}
+	var _rez int
+	err = db.QueryRow("SELECT 1 FROM pgaadauth_list_principals(false) WHERE rolname=$1", roleName).Scan(&_rez)
+	switch {
+	case err == sql.ErrNoRows:
+		return false, nil
+	case err != nil:
+		return false, fmt.Errorf("Error reading info about azure role: %s", err)
+	}
+	return true, nil
+}
+
 func checkRoleExists(client *Client, roleName string) (bool, error) {
 	db, err := client.Connect()
 	if err != nil {
@@ -291,7 +357,6 @@ func checkRoleExists(client *Client, roleName string) (bool, error) {
 	case err != nil:
 		return false, fmt.Errorf("Error reading info about role: %s", err)
 	}
-
 	return true, nil
 }
 

tests/switch_azure_ad_user.sh

@@ -0,0 +1,10 @@
+#!/bin/bash
+
+export TF_ACC=true
+export PGHOST=<some-flexible-server>.postgres.database.azure.com
+export PGPORT=5432
+export PGUSER=azure
+export PGPASSWORD=$(az account get-access-token --resource-type oss-rdbms --query "[accessToken]" -o tsv)
+export PGSSLMODE=require
+export PGSUPERUSER=false
+export AZURE=true

website/docs/r/postgresql_role.html.markdown

@@ -118,6 +118,10 @@ resource "postgresql_role" "my_replication_role" {
 
 * `assume_role` - (Optional) Defines the role to switch to at login via [`SET ROLE`](https://www.postgresql.org/docs/current/sql-set-role.html).
 
+* `azure_identity_id` - (Optional) Unique object identifier of the Microsoft Entra object. See [how to manage azure ad users](https://learn.microsoft.com/en-us/azure/postgresql/flexible-server/how-to-manage-azure-ad-users) for more details.
+
+* `azure_identity_type` - (Optional) Type of the Microsoft Entra object to link to this role: service, user or group.
+
 ## Import Example
 
 `postgresql_role` supports importing resources.  Supposing the following