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

@@ -4,7 +4,6 @@ import (
 	"fmt"
 
 	"github.com/go-logr/logr"
-	"github.com/lib/pq"
 )
 
 type gcppg struct {
@@ -17,47 +16,11 @@ func newGCPPG(postgres *pg) PG {
 	}
 }
 
-func (c *gcppg) CreateDB(dbname, role string) error {
-
-	err := c.GrantRole(role, c.user)
-	if err != nil {
-		return err
-	}
-	err = c.pg.CreateDB(dbname, role)
-	if err != nil {
-		return err
-	}
-	return nil
-}
-
 func (c *gcppg) DropRole(role, newOwner, database string, logger logr.Logger) error {
 	if role == "alloydbsuperuser" || role == "postgres" {
 		logger.Info(fmt.Sprintf("not dropping %s as it is a reserved AlloyDB role", role))
 		return nil
 	}
-	// On AlloyDB the postgres or other alloydbsuperuser members, aren't really superusers so they don't have permissions
-	// to REASSIGN OWNED BY unless he belongs to both roles
-	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)
-	}
 
 	return c.pg.DropRole(role, newOwner, database, logger)
 }

pkg/postgres/database.go

@@ -74,10 +74,10 @@ func NewPG(cfg *config.Cfg, logger logr.Logger) (PG, error) {
 	switch cfg.CloudProvider {
 	case "AWS":
 		logger.Info("Using AWS wrapper")
-		return newAWSPG(postgres), nil
+		return postgres, nil
 	case "Azure":
 		logger.Info("Using Azure wrapper")
-		return newAzurePG(postgres), nil
+		return postgres, nil
 	case "GCP":
 		logger.Info("Using GCP wrapper")
 		return newGCPPG(postgres), nil

pkg/postgres/gcp.go

@@ -48,6 +48,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 +81,28 @@ func (c *pg) RevokeRole(role, revoked string) error {
 }
 
 func (c *pg) DropRole(role, newOwner, database string, logger logr.Logger) error {
+	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 {