Remove duplicate database code in provider logic

Author: pcallewaert

pkg/postgres/aws.go

@@ -29,7 +29,14 @@ const (
 )
 
 func (c *pg) CreateDB(dbname, role string) error {
-	_, err := c.db.Exec(fmt.Sprintf(CREATE_DB, dbname))
+	// Have to add the master role to the group role before we can transfer the database owner
+	err := c.GrantRole(role, c.user)
+	if err != nil {
+		return err
+	}
+
+	// Create database
+	_, err = c.db.Exec(fmt.Sprintf(CREATE_DB, dbname))
 	if err != nil {
 		// eat DUPLICATE DATABASE ERROR
 		if err.(*pq.Error).Code != "42P04" {

pkg/postgres/azure.go

@@ -71,20 +71,7 @@ func NewPG(cfg *config.Cfg, logger logr.Logger) (PG, error) {
 		default_database: cfg.PostgresDefaultDb,
 	}
 
-	switch cfg.CloudProvider {
-	case "AWS":
-		logger.Info("Using AWS wrapper")
-		return newAWSPG(postgres), nil
-	case "Azure":
-		logger.Info("Using Azure wrapper")
-		return newAzurePG(postgres), nil
-	case "GCP":
-		logger.Info("Using GCP wrapper")
-		return newGCPPG(postgres), nil
-	default:
-		logger.Info("Using default postgres implementation")
-		return postgres, nil
-	}
+	return postgres, nil
 }
 
 func (c *pg) GetUser() string {

pkg/postgres/database.go

@@ -20,6 +20,13 @@ const (
 	REASIGN_OBJECTS     = `REASSIGN OWNED BY "%s" TO "%s"`
 )
 
+var reservedRoles = map[string]struct{}{
+	"alloydbsuperuser":  {}, // GCP AlloyDB
+	"cloudsqlsuperuser": {}, // GCP Cloud SQL
+	"rdsadmin":          {}, // AWS RDS
+	"azuresu":           {}, // Azure Database for PostgreSQL
+}
+
 func (c *pg) CreateGroupRole(role string) error {
 	// Error code 42710 is duplicate_object (role already exists)
 	_, err := c.db.Exec(fmt.Sprintf(CREATE_GROUP_ROLE, role))
@@ -48,6 +55,11 @@ func (c *pg) CreateUserRole(role, password string) (string, error) {
 	if err != nil {
 		return "", err
 	}
+
+	err = c.GrantRole(role, c.user)
+	if err != nil {
+		return "", err
+	}
 	return role, nil
 }
 
@@ -76,6 +88,33 @@ func (c *pg) RevokeRole(role, revoked string) error {
 }
 
 func (c *pg) DropRole(role, newOwner, database string, logger logr.Logger) error {
+	if _, reserved := reservedRoles[role]; reserved || role == c.user {
+		logger.Info(fmt.Sprintf("not dropping %s as it is a reserved role", role))
+		return nil
+	}
+
+	err := c.GrantRole(role, c.user)
+	if err != nil && err.(*pq.Error).Code != "0LP01" {
+		if err.(*pq.Error).Code == "42704" {
+			// The group role does not exist, no point in continuing
+			return nil
+		}
+		return err
+	}
+	defer c.RevokeRole(role, c.user)
+	if newOwner != c.user {
+		err = c.GrantRole(newOwner, c.user)
+		if err != nil && err.(*pq.Error).Code != "0LP01" {
+			if err.(*pq.Error).Code == "42704" {
+				// The group role does not exist, no point of granting roles
+				logger.Info(fmt.Sprintf("not granting %s to %s as %s does not exist", role, newOwner, newOwner))
+				return nil
+			}
+			return err
+		}
+		defer c.RevokeRole(newOwner, c.user)
+	}
+
 	// REASSIGN OWNED BY only works if the correct database is selected
 	tmpDb, err := GetConnection(c.user, c.pass, c.host, database, c.args, logger)
 	if err != nil {