From 2b4da80dbbe44048f2c36b86c2cd543a80d05735 Mon Sep 17 00:00:00 2001 From: Chris Hundt Date: Wed, 24 Sep 2025 14:24:55 -0400 Subject: [PATCH 1/3] [CMCSMACD-4695] Add a function for syncing users and roles This will be used instead of migrations for making sure that a database has the right users and roles. --- migrations/user.go | 114 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 114 insertions(+) create mode 100644 migrations/user.go diff --git a/migrations/user.go b/migrations/user.go new file mode 100644 index 0000000..8553cf1 --- /dev/null +++ b/migrations/user.go @@ -0,0 +1,114 @@ +package migrations + +import ( + "fmt" + "log" + + "github.com/jmoiron/sqlx" + "github.com/lib/pq" +) + +type PostgreSQLUser struct { + Username string + GrantRoles []string +} + +// Make sure that the given users exist in database cluster and have only the +// role memberships specified. If withPasswords is true, set each user's password +// to its username. Otherwise remove each user's password. +// All operations are done in a single transaction. +func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords bool) error { + tx, err := db.Begin() + if err != nil { + return fmt.Errorf("Error starting transaction: %w", err) + } + committed := false + defer func() { + if !committed { + err := tx.Rollback() + if err != nil { + log.Printf("Error rollink back: %s", err) + } + } + }() + + for _, user := range users { + createUserSQL := fmt.Sprintf(` + DO $$ + DECLARE + username text := %s; + BEGIN + IF NOT EXISTS ( + SELECT FROM pg_catalog.pg_user WHERE usename = username + ) THEN + EXECUTE format('CREATE USER %%I', username); + END IF; + END + $$`, pq.QuoteLiteral(user.Username)) + _, err := tx.Exec(createUserSQL) + if err != nil { + return fmt.Errorf("Failed to create user %q: %w", user.Username, err) + } + + // Drop all existing roles + dropRolesSQL := fmt.Sprintf(` + DO $$ + DECLARE + r RECORD; + BEGIN + FOR r IN + SELECT roleid::regrole AS granted_role + FROM pg_catalog.pg_auth_members + WHERE member = %s::regrole + LOOP + EXECUTE format('REVOKE %%I FROM %s', r.granted_role); + END LOOP; + END + $$;`, pq.QuoteLiteral(user.Username), pq.QuoteIdentifier(user.Username)) + _, err = tx.Exec(dropRolesSQL) + if err != nil { + return fmt.Errorf("Failed to drop roles for user %q: %w", user.Username, err) + } + + // There could be privileges on a variety of different objects. + // See https://www.postgresql.org/docs/current/sql-revoke.html + // But we will just worry about roles. + + // Add roles + for _, role := range user.GrantRoles { + grantSQL := fmt.Sprintf("GRANT %s TO %s", pq.QuoteIdentifier(role), pq.QuoteIdentifier(user.Username)) + _, err = tx.Exec(grantSQL) + if err != nil { + return fmt.Errorf("Failed to give role %q to user %q: %w", role, user.Username, err) + } + } + + // Set or remove password + if withPasswords { + _, err = tx.Exec( + fmt.Sprintf("ALTER USER %s WITH PASSWORD %s", + pq.QuoteIdentifier(user.Username), + pq.QuoteLiteral(user.Username)), + ) + if err != nil { + return fmt.Errorf("Failed to set password for user %q: %w", user.Username, err) + } + } else { + _, err = tx.Exec( + fmt.Sprintf("ALTER USER %s WITH PASSWORD NULL", + pq.QuoteIdentifier(user.Username)), + ) + if err != nil { + return fmt.Errorf("Failed to remove password for user %q: %w", user.Username, err) + } + } + } + + committed = true + err = tx.Commit() + if err != nil { + return fmt.Errorf("Error committing transaction: %w", err) + } + + return nil +} From e033a2ea5081482a325689f24f58837c504b85ee Mon Sep 17 00:00:00 2001 From: Chris Hundt Date: Thu, 25 Sep 2025 09:56:45 -0400 Subject: [PATCH 2/3] Adjust interface --- migrations/user.go | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/migrations/user.go b/migrations/user.go index 8553cf1..0f34b3f 100644 --- a/migrations/user.go +++ b/migrations/user.go @@ -13,11 +13,19 @@ type PostgreSQLUser struct { GrantRoles []string } +type UserAuthenticationType string + +const ( + UserAuthenticationTypeIAM UserAuthenticationType = "iam" + UserAuthenticationTypePassword UserAuthenticationType = "password" +) + // Make sure that the given users exist in database cluster and have only the -// role memberships specified. If withPasswords is true, set each user's password -// to its username. Otherwise remove each user's password. +// role memberships specified. If authType is UserAuthenticationTypePassword, +// set each user's password to its username. Otherwise remove each user's password +// and also add the rds_iam role for each user. // All operations are done in a single transaction. -func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords bool) error { +func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, authType UserAuthenticationType) error { tx, err := db.Begin() if err != nil { return fmt.Errorf("Error starting transaction: %w", err) @@ -75,7 +83,11 @@ func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords boo // But we will just worry about roles. // Add roles - for _, role := range user.GrantRoles { + roles := user.GrantRoles + if authType == UserAuthenticationTypeIAM { + roles = append(roles, "rds_iam") + } + for _, role := range roles { grantSQL := fmt.Sprintf("GRANT %s TO %s", pq.QuoteIdentifier(role), pq.QuoteIdentifier(user.Username)) _, err = tx.Exec(grantSQL) if err != nil { @@ -84,7 +96,8 @@ func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords boo } // Set or remove password - if withPasswords { + switch authType { + case UserAuthenticationTypePassword: _, err = tx.Exec( fmt.Sprintf("ALTER USER %s WITH PASSWORD %s", pq.QuoteIdentifier(user.Username), @@ -93,7 +106,7 @@ func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords boo if err != nil { return fmt.Errorf("Failed to set password for user %q: %w", user.Username, err) } - } else { + case UserAuthenticationTypeIAM: _, err = tx.Exec( fmt.Sprintf("ALTER USER %s WITH PASSWORD NULL", pq.QuoteIdentifier(user.Username)), @@ -101,6 +114,8 @@ func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords boo if err != nil { return fmt.Errorf("Failed to remove password for user %q: %w", user.Username, err) } + default: + return fmt.Errorf("Invalid authType %q", authType) } } From 8adfcc4753c43beb059c5fc7b9b98e89f6bf46d3 Mon Sep 17 00:00:00 2001 From: Chris Hundt Date: Thu, 25 Sep 2025 10:29:43 -0400 Subject: [PATCH 3/3] Address review comments --- migrations/user.go | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/migrations/user.go b/migrations/user.go index 0f34b3f..5f2da0e 100644 --- a/migrations/user.go +++ b/migrations/user.go @@ -1,6 +1,7 @@ package migrations import ( + "database/sql" "fmt" "log" @@ -30,13 +31,10 @@ func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, authType UserAuth if err != nil { return fmt.Errorf("Error starting transaction: %w", err) } - committed := false defer func() { - if !committed { - err := tx.Rollback() - if err != nil { - log.Printf("Error rollink back: %s", err) - } + err := tx.Rollback() + if err != nil && err != sql.ErrTxDone { + log.Printf("Error rolling back: %s", err) } }() @@ -119,7 +117,6 @@ func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, authType UserAuth } } - committed = true err = tx.Commit() if err != nil { return fmt.Errorf("Error committing transaction: %w", err)