PTEUDO-2418: move grant privileges logic outside create role check

pkg/dbclient/client.go

@@ -426,67 +426,67 @@ func (pc *client) CreateRole(dbName, rolename, schema string) (bool, error) {
 
 	start := time.Now()
 	var exists bool
-	created := false
-
-	err := pc.DB.QueryRow("SELECT EXISTS(SELECT 1 FROM pg_catalog.pg_roles where pg_roles.rolname = $1)", rolename).Scan(&exists)
+	err := pc.DB.QueryRow("SELECT EXISTS(SELECT 1 FROM pg_catalog.pg_roles WHERE rolname = $1)", rolename).Scan(&exists)
 	if err != nil {
 		pc.log.Error(err, "could not query for role")
 		metrics.UsersCreatedErrors.WithLabelValues("read error").Inc()
-		return created, err
+		return false, err
 	}
 
+	created := false
 	if !exists {
-		pc.log.Info("creating a ROLE", "role", rolename)
-
-		_, err = pc.DB.Exec(fmt.Sprintf("CREATE ROLE %s WITH NOLOGIN", pq.QuoteIdentifier(rolename)))
-		if err != nil {
-			pc.log.Error(err, "could not create role "+rolename)
-			metrics.UsersCreatedErrors.WithLabelValues("create error").Inc()
-			return created, err
+		if err := pc.createRole(rolename); err != nil {
+			return false, err
 		}
+		created = true
+	}
 
-		db, err := pc.getDB(dbName)
-		if err != nil {
-			pc.log.Error(err, "could not connect to db", "database", dbName)
-			return created, err
-		}
-		defer db.Close()
+	if err := pc.grantRolePrivileges(dbName, rolename, schema); err != nil {
+		return created, err
+	}
 
-		grantPrivileges := `
-			GRANT ALL PRIVILEGES ON DATABASE %s TO %s;
-			GRANT ALL ON SCHEMA %s TO %s;
-			GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA %s TO %s;
-			GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA %s TO %s;
-		`
+	metrics.UsersCreateTime.Observe(time.Since(start).Seconds())
+	return created, nil
+}
 
-		_, err = db.Exec(fmt.Sprintf(grantPrivileges,
-			pq.QuoteIdentifier(dbName), pq.QuoteIdentifier(rolename),
-			pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename),
-			pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename),
-			pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename),
-		))
-		if err != nil {
-			pc.log.Error(err, "could not set permissions to role "+rolename)
-			metrics.UsersCreatedErrors.WithLabelValues("grant error").Inc()
-			return created, err
-		}
+func (pc *client) createRole(rolename string) error {
+	pc.log.Info("creating a ROLE", "role", rolename)
+	_, err := pc.DB.Exec(fmt.Sprintf("CREATE ROLE %s WITH NOLOGIN", pq.QuoteIdentifier(rolename)))
+	if err != nil {
+		pc.log.Error(err, "could not create role", "role", rolename)
+		metrics.UsersCreatedErrors.WithLabelValues("create error").Inc()
+		return err
+	}
+	pc.log.Info("role has been created", "role", rolename)
+	metrics.UsersCreated.Inc()
+	return nil
+}
 
-		grantSchemaPrivileges := `GRANT ALL ON SCHEMA %s TO %s;`
-		_, err = db.Exec(fmt.Sprintf(grantSchemaPrivileges, pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename)))
-		if err != nil {
-			pc.log.Error(err, "could not set schema privileges to role "+rolename)
+func (pc *client) grantRolePrivileges(dbName, rolename, schema string) error {
+	pc.log.Info("granting privileges to role", "dbName", dbName, "role", rolename, "schema", schema)
+	db, err := pc.getDB(dbName)
+	if err != nil {
+		pc.log.Error(err, "could not connect to db", "database", dbName)
+		return err
+	}
+	defer db.Close()
+
+	privileges := []string{
+		fmt.Sprintf("GRANT ALL PRIVILEGES ON DATABASE %s TO %s;", pq.QuoteIdentifier(dbName), pq.QuoteIdentifier(rolename)),
+		fmt.Sprintf("GRANT ALL ON SCHEMA %s TO %s;", pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename)),
+		fmt.Sprintf("GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA %s TO %s;", pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename)),
+		fmt.Sprintf("GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA %s TO %s;", pq.QuoteIdentifier(schema), pq.QuoteIdentifier(rolename)),
+	}
+
+	for _, query := range privileges {
+		if _, err := db.Exec(query); err != nil {
+			pc.log.Error(err, "could not set permissions", "role", rolename, "query", query)
 			metrics.UsersCreatedErrors.WithLabelValues("grant error").Inc()
-			return created, err
+			return err
 		}
-
-		created = true
-		pc.log.Info("role has been created", "role", rolename)
-		metrics.UsersCreated.Inc()
-		duration := time.Since(start)
-		metrics.UsersCreateTime.Observe(duration.Seconds())
 	}
 
-	return created, nil
+	return nil
 }
 
 func (pc *client) CreateAdminRole(dbName, rolename, schema string) (bool, error) {

pkg/dbclient/client_test.go

@@ -1,6 +1,9 @@
 package dbclient
 
 import (
+	"fmt"
+	"github.com/google/uuid"
+	"github.com/lib/pq"
 	"testing"
 
 	"github.com/go-logr/logr"
@@ -18,6 +21,50 @@ func NewTestLogger(t *testing.T) logr.Logger {
 	return zapr.NewLogger(zapTest)
 }
 
+func TestPostgresGrantRolePrivileges(t *testing.T) {
+	roleName := uuid.New().String()
+	dbName := uuid.New().String()
+
+	db, dsn, dbClose := dockerdb.Run(dockerdb.Config{
+		Username: "testuser",
+		Password: "testpassword",
+		Database: dbName,
+	})
+	defer dbClose()
+
+	pc := &client{
+		dbType:  "postgres",
+		dbURL:   dsn,
+		DB:      db,
+		adminDB: db,
+		log:     NewTestLogger(t),
+	}
+
+	_, err := pc.CreateDatabase(dbName)
+	if err != nil {
+	}
+	_, err = pc.DB.Exec(fmt.Sprintf("CREATE ROLE %s WITH NOLOGIN", pq.QuoteIdentifier(roleName)))
+
+	err = pc.grantRolePrivileges(dbName, roleName, "public")
+	if err != nil {
+		t.Errorf("\t%s GrantRolePrivileges() error = %v", failed, err)
+	}
+
+	var dbPriv, schemaPriv, tablePriv bool
+	err = pc.DB.QueryRow(`
+	SELECT 
+		has_database_privilege($1, $2, 'CONNECT'),
+		has_schema_privilege($1, $3, 'USAGE'),
+		has_schema_privilege($1, $3, 'CREATE')
+	`, roleName, dbName, "public").Scan(&dbPriv, &schemaPriv, &tablePriv)
+
+	if err != nil {
+		t.Errorf("\t%s Error checking privileges: %v", failed, err)
+	} else if !dbPriv || !schemaPriv || !tablePriv {
+		t.Errorf("\t%s Role does not have expected privileges", failed)
+	}
+}
+
 func TestPostgresClientOperations(t *testing.T) {
 
 	db, dsn, close := dockerdb.Run(dockerdb.Config{