From 4beb865cc90095f0339a91806d2baad3001a403f Mon Sep 17 00:00:00 2001
From: meethil15 <meethil.yadav@sentilink.com>
Date: Sat, 17 May 2025 20:54:10 -0500
Subject: [PATCH] Add capabitlity to set pgaudit logging for role

---
 postgresql/resource_postgresql_role.go      | 63 +++++++++++++++++++++
 postgresql/resource_postgresql_role_test.go | 24 ++++++++
 2 files changed, 87 insertions(+)

diff --git a/postgresql/resource_postgresql_role.go b/postgresql/resource_postgresql_role.go
index 5b346f3b..90ef6636 100644
--- a/postgresql/resource_postgresql_role.go
+++ b/postgresql/resource_postgresql_role.go
@@ -35,6 +35,7 @@ const (
 	roleSearchPathAttr                      = "search_path"
 	roleStatementTimeoutAttr                = "statement_timeout"
 	roleAssumeRoleAttr                      = "assume_role"
+	roleAuditAttr                           = "audit"
 
 	// Deprecated options
 	roleDepEncryptedAttr = "encrypted"
@@ -173,6 +174,15 @@ func resourcePostgreSQLRole() *schema.Resource {
 				Optional:    true,
 				Description: "Role to switch to at login",
 			},
+			roleAuditAttr: {
+				Type: schema.TypeList,
+				Elem: &schema.Schema{
+					Type:         schema.TypeString,
+					ValidateFunc: validation.StringInSlice([]string{"ALL", "READ", "WRITE", "FUNCTION", "ROLE", "DDL", "MISC", "MISC_SET"}, false),
+				},
+				Optional:    true,
+				Description: "List of audit classes to apply to the role. Allowed values: ALL, READ, WRITE, FUNCTION, ROLE, DDL, MISC, MISC_SET",
+			},
 		},
 	}
 }
@@ -311,6 +321,10 @@ func resourcePostgreSQLRoleCreate(db *DBConnection, d *schema.ResourceData) erro
 		return err
 	}
 
+	if err = setAudit(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return fmt.Errorf("could not commit transaction: %w", err)
 	}
@@ -472,6 +486,20 @@ func resourcePostgreSQLRoleReadImpl(db *DBConnection, d *schema.ResourceData) er
 
 	d.Set(roleIdleInTransactionSessionTimeoutAttr, idleInTransactionSessionTimeout)
 
+	var auditClasses string
+	auditSQL := fmt.Sprintf(`SELECT setting FROM pg_roles
+        JOIN pg_settings ON rolname = $1 AND name = 'pgaudit.log'`)
+	err = db.QueryRow(auditSQL, roleID).Scan(&auditClasses)
+	if err != nil && err != sql.ErrNoRows {
+		return fmt.Errorf("Error reading audit classes: %w", err)
+	}
+
+	if auditClasses != "" {
+		d.Set(roleAuditAttr, strings.Split(auditClasses, ","))
+	} else {
+		d.Set(roleAuditAttr, nil)
+	}
+
 	d.SetId(roleName)
 
 	password, err := readRolePassword(db, d, roleCanLogin)
@@ -689,6 +717,10 @@ func resourcePostgreSQLRoleUpdate(db *DBConnection, d *schema.ResourceData) erro
 		return err
 	}
 
+	if err := setAudit(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return fmt.Errorf("could not commit transaction: %w", err)
 	}
@@ -1062,3 +1094,34 @@ func setAssumeRole(txn *sql.Tx, d *schema.ResourceData) error {
 	}
 	return nil
 }
+
+func setAudit(txn *sql.Tx, d *schema.ResourceData) error {
+	if !d.HasChange(roleAuditAttr) {
+		return nil
+	}
+
+	roleName := d.Get(roleNameAttr).(string)
+	auditClasses := d.Get(roleAuditAttr).([]interface{})
+
+	if len(auditClasses) > 0 {
+		classes := strings.Join(func() []string {
+			result := make([]string, len(auditClasses))
+			for i, class := range auditClasses {
+				result[i] = pq.QuoteLiteral(class.(string))
+			}
+			return result
+		}(), ", ")
+
+		sql := fmt.Sprintf("ALTER ROLE %s SET pgaudit.log TO %s", pq.QuoteIdentifier(roleName), classes)
+		if _, err := txn.Exec(sql); err != nil {
+			return fmt.Errorf("could not set audit classes %s for role %s: %w", classes, roleName, err)
+		}
+	} else {
+		sql := fmt.Sprintf("ALTER ROLE %s RESET pgaudit.log", pq.QuoteIdentifier(roleName))
+		if _, err := txn.Exec(sql); err != nil {
+			return fmt.Errorf("could not reset audit classes for role %s: %w", roleName, err)
+		}
+	}
+
+	return nil
+}
diff --git a/postgresql/resource_postgresql_role_test.go b/postgresql/resource_postgresql_role_test.go
index fc958840..efd82bc0 100644
--- a/postgresql/resource_postgresql_role_test.go
+++ b/postgresql/resource_postgresql_role_test.go
@@ -237,6 +237,30 @@ resource "postgresql_role" "test_role" {
 	})
 }
 
+func TestAccPostgresqlRole_Audit(t *testing.T) {
+	resource.Test(t, resource.TestCase{
+		PreCheck: func() {
+			testAccPreCheck(t)
+		},
+		Providers: testAccProviders,
+		Steps: []resource.TestStep{
+			{
+				Config: `
+resource "postgresql_role" "test" {
+  name  = "test_role"
+  audit = ["READ", "WRITE"]
+}
+`,
+				Check: resource.ComposeTestCheckFunc(
+					resource.TestCheckResourceAttr("postgresql_role.test", "audit.#", "2"),
+					resource.TestCheckResourceAttr("postgresql_role.test", "audit.0", "READ"),
+					resource.TestCheckResourceAttr("postgresql_role.test", "audit.1", "WRITE"),
+				),
+			},
+		},
+	})
+}
+
 func testAccCheckPostgresqlRoleDestroy(s *terraform.State) error {
 	client := testAccProvider.Meta().(*Client)