Skip to content

Commit 2b4da80

Browse files
author
Chris Hundt
committed
[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.
1 parent b59bf13 commit 2b4da80

1 file changed

Lines changed: 114 additions & 0 deletions

File tree

migrations/user.go

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
package migrations
2+
3+
import (
4+
"fmt"
5+
"log"
6+
7+
"github.com/jmoiron/sqlx"
8+
"github.com/lib/pq"
9+
)
10+
11+
type PostgreSQLUser struct {
12+
Username string
13+
GrantRoles []string
14+
}
15+
16+
// Make sure that the given users exist in database cluster and have only the
17+
// role memberships specified. If withPasswords is true, set each user's password
18+
// to its username. Otherwise remove each user's password.
19+
// All operations are done in a single transaction.
20+
func EnsureUsersWithRoles(db *sqlx.DB, users []PostgreSQLUser, withPasswords bool) error {
21+
tx, err := db.Begin()
22+
if err != nil {
23+
return fmt.Errorf("Error starting transaction: %w", err)
24+
}
25+
committed := false
26+
defer func() {
27+
if !committed {
28+
err := tx.Rollback()
29+
if err != nil {
30+
log.Printf("Error rollink back: %s", err)
31+
}
32+
}
33+
}()
34+
35+
for _, user := range users {
36+
createUserSQL := fmt.Sprintf(`
37+
DO $$
38+
DECLARE
39+
username text := %s;
40+
BEGIN
41+
IF NOT EXISTS (
42+
SELECT FROM pg_catalog.pg_user WHERE usename = username
43+
) THEN
44+
EXECUTE format('CREATE USER %%I', username);
45+
END IF;
46+
END
47+
$$`, pq.QuoteLiteral(user.Username))
48+
_, err := tx.Exec(createUserSQL)
49+
if err != nil {
50+
return fmt.Errorf("Failed to create user %q: %w", user.Username, err)
51+
}
52+
53+
// Drop all existing roles
54+
dropRolesSQL := fmt.Sprintf(`
55+
DO $$
56+
DECLARE
57+
r RECORD;
58+
BEGIN
59+
FOR r IN
60+
SELECT roleid::regrole AS granted_role
61+
FROM pg_catalog.pg_auth_members
62+
WHERE member = %s::regrole
63+
LOOP
64+
EXECUTE format('REVOKE %%I FROM %s', r.granted_role);
65+
END LOOP;
66+
END
67+
$$;`, pq.QuoteLiteral(user.Username), pq.QuoteIdentifier(user.Username))
68+
_, err = tx.Exec(dropRolesSQL)
69+
if err != nil {
70+
return fmt.Errorf("Failed to drop roles for user %q: %w", user.Username, err)
71+
}
72+
73+
// There could be privileges on a variety of different objects.
74+
// See https://www.postgresql.org/docs/current/sql-revoke.html
75+
// But we will just worry about roles.
76+
77+
// Add roles
78+
for _, role := range user.GrantRoles {
79+
grantSQL := fmt.Sprintf("GRANT %s TO %s", pq.QuoteIdentifier(role), pq.QuoteIdentifier(user.Username))
80+
_, err = tx.Exec(grantSQL)
81+
if err != nil {
82+
return fmt.Errorf("Failed to give role %q to user %q: %w", role, user.Username, err)
83+
}
84+
}
85+
86+
// Set or remove password
87+
if withPasswords {
88+
_, err = tx.Exec(
89+
fmt.Sprintf("ALTER USER %s WITH PASSWORD %s",
90+
pq.QuoteIdentifier(user.Username),
91+
pq.QuoteLiteral(user.Username)),
92+
)
93+
if err != nil {
94+
return fmt.Errorf("Failed to set password for user %q: %w", user.Username, err)
95+
}
96+
} else {
97+
_, err = tx.Exec(
98+
fmt.Sprintf("ALTER USER %s WITH PASSWORD NULL",
99+
pq.QuoteIdentifier(user.Username)),
100+
)
101+
if err != nil {
102+
return fmt.Errorf("Failed to remove password for user %q: %w", user.Username, err)
103+
}
104+
}
105+
}
106+
107+
committed = true
108+
err = tx.Commit()
109+
if err != nil {
110+
return fmt.Errorf("Error committing transaction: %w", err)
111+
}
112+
113+
return nil
114+
}

0 commit comments

Comments
 (0)