feat: AWS RDS pg_repack additional support

Author: tkcontiant

internal/controller/postgres_controller.go

@@ -181,8 +181,8 @@ func (r *PostgresReconciler) Reconcile(ctx context.Context, req ctrl.Request) (c
 		if err != nil {
 			return requeue(errors.NewInternalError(err))
 		}
-		// Alter database owner if the owner role was changed
-		err = r.pg.AlterDatabaseOwner(instance.Spec.Database, instance.Status.Roles.Owner)
+		// Alter database owner to desiredOwner if the owner role was changed
+		err = r.pg.AlterDatabaseOwner(instance.Spec.Database, desiredOwner)
 		if err != nil {
 			return requeue(errors.NewInternalError(err))
 		}

internal/controller/postgres_controller_test.go

@@ -71,8 +71,6 @@ 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()).Return(nil).AnyTimes()
 		cl = k8sClient
 		// Create runtime scheme
 		sc = scheme.Scheme
@@ -365,6 +363,32 @@ var _ = Describe("PostgresReconciler", func() {
 			})
 		})
 
+		Context("MasterRole is changed for existing database", func() {
+			BeforeEach(func() {
+				modPostgres := postgresCR.DeepCopy()
+				modPostgres.Spec.MasterRole = "new-master-role"
+				modPostgres.Status = v1alpha1.PostgresStatus{
+					Succeeded: true,
+					Roles: v1alpha1.PostgresRoles{
+						Owner: "old-master-role",
+					},
+				}
+				initClient(modPostgres, false)
+			})
+
+			It("should alter database owner to the desired role", func() {
+				pg.EXPECT().RenameGroupRole("old-master-role", "new-master-role").Return(nil).Times(1)
+				pg.EXPECT().AlterDatabaseOwner(name, "new-master-role").Return(nil).Times(1)
+
+				err := runReconcile(rp, ctx, req)
+				Expect(err).NotTo(HaveOccurred())
+
+				foundPostgres := &v1alpha1.Postgres{}
+				Expect(cl.Get(ctx, types.NamespacedName{Name: name, Namespace: namespace}, foundPostgres)).To(BeNil())
+				Expect(foundPostgres.Status.Roles.Owner).To(Equal("new-master-role"))
+			})
+		})
+
 		Context("Correct annotation filter is set", func() {
 			BeforeEach(func() {
 				// Create client

pkg/postgres/aws.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/lib/pq"
 )
@@ -10,6 +11,11 @@ type awspg struct {
 	pg
 }
 
+const (
+	AWS_ALTER_REPACK_DEFAULT_PRIVS_TABLES    = `ALTER DEFAULT PRIVILEGES FOR ROLE "%s" IN SCHEMA "repack" GRANT INSERT ON TABLES TO PUBLIC`
+	AWS_ALTER_REPACK_DEFAULT_PRIVS_SEQUENCES = `ALTER DEFAULT PRIVILEGES FOR ROLE "%s" IN SCHEMA "repack" GRANT USAGE, SELECT ON SEQUENCES TO PUBLIC`
+)
+
 func newAWSPG(postgres *pg) PG {
 	return &awspg{
 		*postgres,
@@ -38,6 +44,45 @@ func (c *awspg) CreateDB(dbname, role string) error {
 	return c.pg.CreateDB(dbname, role)
 }
 
+func (c *awspg) CreateExtension(dbname, extension string) error {
+	// Keep standard extension creation behavior for AWS as well.
+	err := c.pg.CreateExtension(dbname, extension)
+	if err != nil {
+		return err
+	}
+
+	// AWS-specific workaround is only required for pg_repack.
+	if !strings.EqualFold(extension, "pg_repack") {
+		return nil
+	}
+
+	var owner string
+	// Resolve current database owner role to target ALTER DEFAULT PRIVILEGES FOR ROLE.
+	err = c.db.QueryRow(fmt.Sprintf(GET_DB_OWNER, dbname)).Scan(&owner)
+	if err != nil {
+		return err
+	}
+
+	// Execute pg_repack privilege statements in the target database.
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, dbname, c.args)
+	if err != nil {
+		return err
+	}
+	defer tmpDb.Close()
+
+	_, err = tmpDb.Exec(fmt.Sprintf(AWS_ALTER_REPACK_DEFAULT_PRIVS_TABLES, owner))
+	if err != nil {
+		return err
+	}
+
+	_, err = tmpDb.Exec(fmt.Sprintf(AWS_ALTER_REPACK_DEFAULT_PRIVS_SEQUENCES, owner))
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (c *awspg) CreateUserRole(role, password string) (string, error) {
 	returnedRole, err := c.pg.CreateUserRole(role, password)
 	if err != nil {