refactor: Use context on all SQL calls

Author: oxyno-zeta

internal/controller/postgresql/postgres/aws.go

@@ -1,6 +1,7 @@
 package postgres
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/lib/pq"
@@ -21,32 +22,32 @@ func newAWSPG(postgres *pg) PG {
 	}
 }
 
-func (c *awspg) AlterDefaultLoginRole(role, setRole string) error {
+func (c *awspg) AlterDefaultLoginRole(ctx context.Context, role, setRole string) error {
 	// On AWS RDS the postgres user isn't really superuser so he doesn't have permissions
 	// to ALTER USER unless he belongs to both roles
-	err := c.GrantRole(role, c.user, false)
+	err := c.GrantRole(ctx, role, c.user, false)
 	if err != nil {
 		return err
 	}
 
 	defer func() {
-		err := c.RevokeRole(role, c.user)
+		err := c.RevokeRole(ctx, role, c.user)
 		// Check error
 		if err != nil {
 			c.log.Error(err, "error in revoke role")
 		}
 	}()
 
-	return c.pg.AlterDefaultLoginRole(role, setRole)
+	return c.pg.AlterDefaultLoginRole(ctx, role, setRole)
 }
 
-func (c *awspg) CreateDB(dbname, role string) error {
+func (c *awspg) CreateDB(ctx context.Context, dbname, role string) error {
 	err := c.connect(c.defaultDatabase)
 	if err != nil {
 		return err
 	}
 
-	_, err = c.db.Exec(fmt.Sprintf(CreateDBWithoutOwnerSQLTemplate, dbname))
+	_, err = c.db.ExecContext(ctx, fmt.Sprintf(CreateDBWithoutOwnerSQLTemplate, dbname))
 	if err != nil {
 		// eat DUPLICATE DATABASE ERROR
 		// Try to cast error
@@ -56,18 +57,18 @@ func (c *awspg) CreateDB(dbname, role string) error {
 		}
 	}
 
-	_, err = c.db.Exec(fmt.Sprintf(AlterDBOwnerSQLTemplate, dbname, role))
+	_, err = c.db.ExecContext(ctx, fmt.Sprintf(AlterDBOwnerSQLTemplate, dbname, role))
 	if err != nil {
 		return err
 	}
 
 	return nil
 }
 
-func (c *awspg) DropRoleAndDropAndChangeOwnedBy(role, newOwner, database string) error {
+func (c *awspg) DropRoleAndDropAndChangeOwnedBy(ctx context.Context, role, newOwner, database string) error {
 	// On AWS RDS the postgres user isn't really superuser so he doesn't have permissions
 	// to REASSIGN OWNED BY unless he belongs to both roles
-	err := c.GrantRole(role, c.user, false)
+	err := c.GrantRole(ctx, role, c.user, false)
 	// Check error
 	if err != nil {
 		// Try to cast error
@@ -85,7 +86,7 @@ func (c *awspg) DropRoleAndDropAndChangeOwnedBy(role, newOwner, database string)
 		}
 	}
 
-	err = c.GrantRole(newOwner, c.user, false)
+	err = c.GrantRole(ctx, newOwner, c.user, false)
 	// Check error
 	if err != nil {
 		// Try to cast error
@@ -107,19 +108,19 @@ func (c *awspg) DropRoleAndDropAndChangeOwnedBy(role, newOwner, database string)
 	}
 
 	defer func() {
-		err := c.RevokeRole(newOwner, c.user)
+		err := c.RevokeRole(ctx, newOwner, c.user)
 		if err != nil {
 			c.log.Error(err, "error in revoke role")
 		}
 	}()
 
-	return c.pg.DropRoleAndDropAndChangeOwnedBy(role, newOwner, database)
+	return c.pg.DropRoleAndDropAndChangeOwnedBy(ctx, role, newOwner, database)
 }
 
-func (c *awspg) ChangeAndDropOwnedBy(role, newOwner, database string) error {
+func (c *awspg) ChangeAndDropOwnedBy(ctx context.Context, role, newOwner, database string) error {
 	// On AWS RDS the postgres user isn't really superuser so he doesn't have permissions
 	// to REASSIGN OWNED BY unless he belongs to both roles
-	err := c.GrantRole(role, c.user, false)
+	err := c.GrantRole(ctx, role, c.user, false)
 	// Check error
 	if err != nil {
 		// Try to cast error
@@ -137,7 +138,7 @@ func (c *awspg) ChangeAndDropOwnedBy(role, newOwner, database string) error {
 		}
 	}
 
-	err = c.GrantRole(newOwner, c.user, false)
+	err = c.GrantRole(ctx, newOwner, c.user, false)
 	// Check error
 	if err != nil {
 		// Try to cast error
@@ -159,11 +160,11 @@ func (c *awspg) ChangeAndDropOwnedBy(role, newOwner, database string) error {
 	}
 
 	defer func() {
-		err := c.RevokeRole(newOwner, c.user)
+		err := c.RevokeRole(ctx, newOwner, c.user)
 		if err != nil {
 			c.log.Error(err, "error in revoke role")
 		}
 	}()
 
-	return c.pg.ChangeAndDropOwnedBy(role, newOwner, database)
+	return c.pg.ChangeAndDropOwnedBy(ctx, role, newOwner, database)
 }

internal/controller/postgresql/postgres/azure.go

@@ -1,6 +1,7 @@
 package postgres
 
 import (
+	"context"
 	"fmt"
 	"strings"
 )
@@ -26,16 +27,16 @@ func newAzurePG(postgres *pg) PG {
 	}
 }
 
-func (azpg *azurepg) CreateUserRole(role, password string, attributes *RoleAttributes) (string, error) {
-	returnedRole, err := azpg.pg.CreateUserRole(role, password, attributes)
+func (azpg *azurepg) CreateUserRole(ctx context.Context, role, password string, attributes *RoleAttributes) (string, error) {
+	returnedRole, err := azpg.pg.CreateUserRole(ctx, role, password, attributes)
 	if err != nil {
 		return "", err
 	}
 
 	return fmt.Sprintf("%s@%s", returnedRole, azpg.serverName), nil
 }
 
-func (azpg *azurepg) GetRoleForLogin(login string) string {
+func (azpg *azurepg) GetRoleForLogin(ctx context.Context, login string) string {
 	splitUser := strings.Split(azpg.user, "@")
 	if len(splitUser) > MinUserSplit {
 		return splitUser[0]
@@ -44,12 +45,12 @@ func (azpg *azurepg) GetRoleForLogin(login string) string {
 	return login
 }
 
-func (azpg *azurepg) CreateDB(dbname, role string) error {
+func (azpg *azurepg) CreateDB(ctx context.Context, 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), false)
+	err := azpg.GrantRole(ctx, role, azpg.GetRoleForLogin(ctx, azpg.user), false)
 	if err != nil {
 		return err
 	}
 
-	return azpg.pg.CreateDB(dbname, role)
+	return azpg.pg.CreateDB(ctx, dbname, role)
 }