Improve Azure implmentation

Author: pcallewaert

pkg/postgres/azure.go

@@ -5,24 +5,36 @@ import (
 	"strings"
 
 	"github.com/go-logr/logr"
-	"github.com/lib/pq"
+)
+
+type AzureType string
+
+const (
+	// Azure Database for PostgreSQL Flexible Server uses default convention for login
+	FLEXIBLE AzureType = "flexible"
+	// Azure Database for PostgreSQL Single Server uses <username>@<servername> convention
+	SINGLE AzureType = "single"
 )
 
 type azurepg struct {
 	serverName string
+	azureType  AzureType
 	pg
 }
 
 func newAzurePG(postgres *pg) PG {
 	splitUser := strings.Split(postgres.user, "@")
 	serverName := ""
-	// We need to know the server name for Azure Database for PostgreSQL Single Server
+	azureType := FLEXIBLE
 	if len(splitUser) > 1 {
+		// If a servername is found, we are using Azure Database for PostgreSQL Single Server
 		serverName = splitUser[1]
+		azureType = SINGLE
 	}
 	return &azurepg{
-		serverName,
-		*postgres,
+		serverName: serverName,
+		azureType:  azureType,
+		pg:         *postgres,
 	}
 }
 
@@ -31,24 +43,31 @@ func (azpg *azurepg) CreateUserRole(role, password string) (string, error) {
 	if err != nil {
 		return "", err
 	}
-	if azpg.serverName == "" {
+
+	// For Flexible Server, just return the role name as-is
+	if azpg.azureType == FLEXIBLE {
 		return returnedRole, nil
 	}
-	// Azure Database for PostgreSQL Single Server offering uses <username>@<servername> convention
+
+	// For Single Server, format as <username>@<servername>
 	return fmt.Sprintf("%s@%s", returnedRole, azpg.serverName), nil
 }
 
 func (azpg *azurepg) GetRoleForLogin(login string) string {
-	splitUser := strings.Split(azpg.user, "@")
-	if len(splitUser) > 1 {
-		return splitUser[0]
+	// For Azure Flexible Server, the login name is the same as the role name
+	if azpg.azureType == FLEXIBLE {
+		return login
 	}
-	return login
+
+	// For Azure Single Server, extract the username part before the '@' symbol
+	splitUser := strings.Split(azpg.user, "@")
+	return splitUser[0]
 }
 
 func (azpg *azurepg) CreateDB(dbname, role string) error {
-	// Have to add the master role to the group role before we can transfer the database owner
-	err := azpg.GrantRole(role, azpg.GetRoleForLogin(azpg.user))
+	// Grant admin privileges to the group role to enable database ownership transfer
+	// This step is necessary before we can set the specified role as the database owner
+	err := azpg.grantAdminRole(role, azpg.GetRoleForLogin(azpg.user))
 	if err != nil {
 		return err
 	}
@@ -57,32 +76,12 @@ func (azpg *azurepg) CreateDB(dbname, role string) error {
 }
 
 func (azpg *azurepg) DropRole(role, newOwner, database string, logger logr.Logger) error {
-	if azpg.serverName != "" {
-		// Logic for Single Server
-		azNewOwner := azpg.GetRoleForLogin(newOwner)
-		return azpg.pg.DropRole(role, azNewOwner, database, logger)
-	} else {
-		// Logic for Flexible Server (same as AWS)
-		// to REASSIGN OWNED BY unless he belongs to both roles
-		err := azpg.pg.GrantRole(role, azpg.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
-		}
-		err = azpg.pg.GrantRole(newOwner, azpg.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 azpg.pg.RevokeRole(newOwner, azpg.pg.user)
-
+	// For Azure Flexible Server, we can use the standard pg implementation
+	if azpg.azureType == FLEXIBLE {
 		return azpg.pg.DropRole(role, newOwner, database, logger)
 	}
+
+	// For Azure Single Server, we need to get the correct format for the newOwner
+	azNewOwner := azpg.GetRoleForLogin(newOwner)
+	return azpg.pg.DropRole(role, azNewOwner, database, logger)
 }

pkg/postgres/role.go

@@ -11,6 +11,7 @@ const (
 	CREATE_GROUP_ROLE   = `CREATE ROLE "%s"`
 	CREATE_USER_ROLE    = `CREATE ROLE "%s" WITH LOGIN PASSWORD '%s'`
 	GRANT_ROLE          = `GRANT "%s" TO "%s"`
+	GRANT_ADMIN_ROLE    = `GRANT "%s" TO "%s" WITH ADMIN OPTION`
 	ALTER_USER_SET_ROLE = `ALTER USER "%s" SET ROLE "%s"`
 	REVOKE_ROLE         = `REVOKE "%s" FROM "%s"`
 	UPDATE_PASSWORD     = `ALTER ROLE "%s" WITH PASSWORD '%s'`
@@ -44,6 +45,14 @@ func (c *pg) GrantRole(role, grantee string) error {
 	return nil
 }
 
+func (c *pg) grantAdminRole(role, grantee string) error {
+	_, err := c.db.Exec(fmt.Sprintf(GRANT_ADMIN_ROLE, role, grantee))
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (c *pg) AlterDefaultLoginRole(role, setRole string) error {
 	_, err := c.db.Exec(fmt.Sprintf(ALTER_USER_SET_ROLE, role, setRole))
 	if err != nil {