fix: tests

Author: tkcontiant

internal/controller/postgres_controller_test.go

@@ -71,6 +71,9 @@ var _ = Describe("PostgresReconciler", func() {
 		// Gomock
 		mockCtrl = gomock.NewController(GinkgoT())
 		pg = mockpg.NewMockPG(mockCtrl)
+		pg.EXPECT().AlterDatabaseOwner(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+		pg.EXPECT().ReassignDatabaseOwner(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
+		pg.EXPECT().ReassignDatabaseOwner(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
 		cl = k8sClient
 		// Create runtime scheme
 		sc = scheme.Scheme

pkg/postgres/aws.go

@@ -82,3 +82,7 @@ func (c *awspg) DropRole(role, newOwner, database string, logger logr.Logger) er
 func (c *awspg) AlterDatabaseOwner(dbName, owner string) error {
 	return c.pg.AlterDatabaseOwner(dbName, owner)
 }
+
+func (c *awspg) ReassignDatabaseOwner(dbName, currentOwner, newOwner string, logger logr.Logger) error {
+	return c.pg.ReassignDatabaseOwner(dbName, currentOwner, newOwner, logger)
+}

pkg/postgres/azure.go

@@ -52,3 +52,7 @@ func (azpg *azurepg) DropRole(role, newOwner, database string, logger logr.Logge
 func (azpg *azurepg) AlterDatabaseOwner(dbName, owner string) error {
 	return azpg.pg.AlterDatabaseOwner(dbName, owner)
 }
+
+func (azpg *azurepg) ReassignDatabaseOwner(dbName, currentOwner, newOwner string, logger logr.Logger) error {
+	return azpg.pg.ReassignDatabaseOwner(dbName, currentOwner, newOwner, logger)
+}

pkg/postgres/database.go

@@ -12,6 +12,7 @@ const (
 	CREATE_SCHEMA           = `CREATE SCHEMA IF NOT EXISTS "%s" AUTHORIZATION "%s"`
 	CREATE_EXTENSION        = `CREATE EXTENSION IF NOT EXISTS "%s"`
 	ALTER_DB_OWNER          = `ALTER DATABASE "%s" OWNER TO "%s"`
+	REASSIGN_DB_OWNER       = `REASSIGN OWNED BY "%s" TO "%s"`
 	DROP_DATABASE           = `DROP DATABASE "%s"`
 	GRANT_USAGE_SCHEMA      = `GRANT USAGE ON SCHEMA "%s" TO "%s"`
 	GRANT_CREATE_TABLE      = `GRANT CREATE ON SCHEMA "%s" TO "%s"`
@@ -50,10 +51,31 @@ func (c *pg) CreateDB(dbname, role string) error {
 
 // reconcile the desired owner of the database
 func (c *pg) AlterDatabaseOwner(dbname, owner string) error {
+	if owner == "" {
+		return nil
+	}
 	_, err := c.db.Exec(fmt.Sprintf(ALTER_DB_OWNER, dbname, owner))
+	return err
+}
+
+func (c *pg) ReassignDatabaseOwner(dbName, currentOwner, newOwner string, logger logr.Logger) error {
+	if currentOwner == "" || newOwner == "" || currentOwner == newOwner {
+		return nil
+	}
+
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, dbName, c.args, logger)
 	if err != nil {
 		return err
 	}
+	defer tmpDb.Close()
+
+	_, err = tmpDb.Exec(fmt.Sprintf(REASSIGN_DB_OWNER, currentOwner, newOwner))
+	if err != nil {
+		if pqErr, ok := err.(*pq.Error); ok && pqErr.Code == "42704" {
+			return nil
+		}
+		return err
+	}
 	return nil
 }
 

pkg/postgres/gcp.go

@@ -87,3 +87,7 @@ func (c *gcppg) DropRole(role, newOwner, database string, logger logr.Logger) er
 func (c *gcppg) AlterDatabaseOwner(dbName, owner string) error {
 	return c.pg.AlterDatabaseOwner(dbName, owner)
 }
+
+func (c *gcppg) ReassignDatabaseOwner(dbName, currentOwner, newOwner string, logger logr.Logger) error {
+	return c.pg.ReassignDatabaseOwner(dbName, currentOwner, newOwner, logger)
+}

pkg/postgres/mock/postgres.go

@@ -19,6 +19,7 @@ type PG interface {
 	UpdatePassword(role, password string) error
 	GrantRole(role, grantee string) error
 	AlterDatabaseOwner(dbName, owner string) error
+	ReassignDatabaseOwner(dbName, currentOwner, newOwner string, logger logr.Logger) error
 	SetSchemaPrivileges(schemaPrivileges PostgresSchemaPrivileges, logger logr.Logger) error
 	RevokeRole(role, revoked string) error
 	AlterDefaultLoginRole(role, setRole string) error