From 7aaf7b7f2fac8b793557cdb4439911db3965caf3 Mon Sep 17 00:00:00 2001
From: kankou-aliaksei <aliaksei.kankou@gmail.com>
Date: Sun, 29 Sep 2024 19:47:26 +0400
Subject: [PATCH] Add support for adding comments to PostgreSQL roles

---
 postgresql/resource_postgresql_role.go      | 37 +++++++++++++++++++++
 postgresql/resource_postgresql_role_test.go | 14 ++++++++
 2 files changed, 51 insertions(+)

diff --git a/postgresql/resource_postgresql_role.go b/postgresql/resource_postgresql_role.go
index b7cb0fab..348327cd 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"
+	roleCommentAttr                         = "comment"
 
 	// Deprecated options
 	roleDepEncryptedAttr = "encrypted"
@@ -173,6 +174,11 @@ func resourcePostgreSQLRole() *schema.Resource {
 				Optional:    true,
 				Description: "Role to switch to at login",
 			},
+			roleCommentAttr: {
+				Type:        schema.TypeString,
+				Optional:    true,
+				Description: "The comment to set for the role",
+			},
 		},
 	}
 }
@@ -311,6 +317,10 @@ func resourcePostgreSQLRoleCreate(db *DBConnection, d *schema.ResourceData) erro
 		return err
 	}
 
+	if err = setRoleComment(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return fmt.Errorf("could not commit transaction: %w", err)
 	}
@@ -689,6 +699,10 @@ func resourcePostgreSQLRoleUpdate(db *DBConnection, d *schema.ResourceData) erro
 		return err
 	}
 
+	if err := setRoleComment(txn, d); err != nil {
+		return err
+	}
+
 	if err = txn.Commit(); err != nil {
 		return fmt.Errorf("could not commit transaction: %w", err)
 	}
@@ -1062,3 +1076,26 @@ func setAssumeRole(txn *sql.Tx, d *schema.ResourceData) error {
 	}
 	return nil
 }
+
+func setRoleComment(txn *sql.Tx, d *schema.ResourceData) error {
+	if !d.HasChange(roleCommentAttr) {
+		return nil
+	}
+
+	comment := d.Get(roleCommentAttr).(string)
+	roleName := d.Get(roleNameAttr).(string)
+
+	if comment != "" {
+		sql := fmt.Sprintf("COMMENT ON ROLE %s IS '%s'", pq.QuoteIdentifier(roleName), pqQuoteLiteral(comment))
+		if _, err := txn.Exec(sql); err != nil {
+			return fmt.Errorf("Error setting comment on role %s: %w", roleName, err)
+		}
+	} else {
+		sql := fmt.Sprintf("COMMENT ON ROLE %s IS NULL", pq.QuoteIdentifier(roleName))
+		if _, err := txn.Exec(sql); err != nil {
+			return fmt.Errorf("Error clearing comment on role %s: %w", roleName, err)
+		}
+	}
+
+	return nil
+}
diff --git a/postgresql/resource_postgresql_role_test.go b/postgresql/resource_postgresql_role_test.go
index ef502f00..814e600a 100644
--- a/postgresql/resource_postgresql_role_test.go
+++ b/postgresql/resource_postgresql_role_test.go
@@ -58,6 +58,10 @@ func TestAccPostgresqlRole_Basic(t *testing.T) {
 					resource.TestCheckResourceAttr("postgresql_role.sub_role", "roles.1", "role_simple"),
 
 					testAccCheckPostgresqlRoleExists("role_with_search_path", nil, []string{"bar", "foo-with-hyphen"}),
+
+					testAccCheckPostgresqlRoleExists("role_with_comment", nil, nil),
+					resource.TestCheckResourceAttr("postgresql_role.role_with_comment", "name", "role_with_comment"),
+					resource.TestCheckResourceAttr("postgresql_role.role_with_comment", "comment", "This is a test comment"),
 				),
 			},
 		},
@@ -104,6 +108,7 @@ resource "postgresql_role" "update_role" {
   login = true
   password = "toto"
   valid_until = "2099-05-04 12:00:00+00"
+  comment = "Initial comment"
 }
 `
 
@@ -122,6 +127,7 @@ resource "postgresql_role" "update_role" {
   statement_timeout = 30000
   idle_in_transaction_session_timeout = 60000
   assume_role = "${postgresql_role.group_role.name}"
+  comment = "Updated comment"
 }
 `
 	resource.Test(t, resource.TestCase{
@@ -146,6 +152,7 @@ resource "postgresql_role" "update_role" {
 					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", ""),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "comment", "Initial comment"),
 					testAccCheckRoleCanLogin(t, "update_role", "toto"),
 				),
 			},
@@ -167,6 +174,7 @@ resource "postgresql_role" "update_role" {
 					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"),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "comment", "Updated comment"),
 					testAccCheckRoleCanLogin(t, "update_role2", "titi"),
 				),
 			},
@@ -185,6 +193,7 @@ resource "postgresql_role" "update_role" {
 					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", ""),
+					resource.TestCheckResourceAttr("postgresql_role.update_role", "comment", "Initial comment"),
 					testAccCheckRoleCanLogin(t, "update_role", "toto"),
 				),
 			},
@@ -437,4 +446,9 @@ resource "postgresql_role" "role_with_search_path" {
   name = "role_with_search_path"
   search_path = ["bar", "foo-with-hyphen"]
 }
+
+resource "postgresql_role" "role_with_comment" {
+  name = "role_with_comment"
+  comment = "This is a test comment"
+}
 `