From 8d7ef54491603a2927a0548332ee86ab448e99e8 Mon Sep 17 00:00:00 2001
From: Balaji Jeevan <32107263+bjeevan-ib@users.noreply.github.com>
Date: Mon, 15 Apr 2024 10:44:50 -0700
Subject: [PATCH] create partman extension in partman schema and grant access
 (#233)

* create partman extension in partman schema and grant access

* update golang version

* pinned version of envtest to avoid https://github.com/kubernetes-sigs/controller-runtime/issues/2720
---
 Makefile                                |  6 +-
 controllers/databaseclaim_controller.go |  9 ++-
 go.mod                                  |  2 +-
 helm/db-controller/values.yaml          |  2 +-
 pkg/dbclient/client.go                  | 82 ++++++++++++++++++++++---
 pkg/dbclient/interface.go               |  1 +
 pkg/pgctl/pgdump.go                     | 13 ++--
 7 files changed, 94 insertions(+), 21 deletions(-)

diff --git a/Makefile b/Makefile
index b731229..d95a848 100644
--- a/Makefile
+++ b/Makefile
@@ -230,7 +230,11 @@ kustomize: ## Download kustomize locally if necessary.
 ENVTEST = $(shell pwd)/bin/setup-envtest
 .PHONY: envtest
 envtest: ## Download envtest-setup locally if necessary.
-	$(call go-get-tool,$(ENVTEST),sigs.k8s.io/controller-runtime/tools/setup-envtest@latest)
+#commenting out the latest setup-envttest to use elease-0.16- because of https://github.com/kubernetes-sigs/controller-runtime/issues/2720
+	#$(call go-get-tool,$(ENVTEST),sigs.k8s.io/controller-runtime/tools/setup-envtest@latest)
+#TODO- reuse the above command after fix from controller-runtime or after we go to go 1.22
+	$(call go-get-tool,$(ENVTEST),sigs.k8s.io/controller-runtime/tools/setup-envtest@release-0.16)
+
 
 # go-get-tool will 'go get' any package $2 and install it to $1.
 PROJECT_DIR := $(shell dirname $(abspath $(lastword $(MAKEFILE_LIST))))
diff --git a/controllers/databaseclaim_controller.go b/controllers/databaseclaim_controller.go
index 0a8870e..666fb9d 100644
--- a/controllers/databaseclaim_controller.go
+++ b/controllers/databaseclaim_controller.go
@@ -1587,10 +1587,17 @@ func (r *DatabaseClaimReconciler) manageUser(dbClient dbclient.Client, status *p
 	rotationTime := r.getPasswordRotationTime()
 
 	// create role
-	_, err := dbClient.CreateGroup(dbName, baseUsername)
+	roleCreated, err := dbClient.CreateGroup(dbName, baseUsername)
 	if err != nil {
 		return err
 	}
+	if roleCreated {
+		// take care of special extentions related to the user
+		err = dbClient.CreateSpecialExtentions(dbName, baseUsername)
+		if err != nil {
+			return err
+		}
+	}
 
 	if dbu.IsUserChanged(*status) {
 		oldUsername := dbu.TrimUserSuffix(status.ConnectionInfo.Username)
diff --git a/go.mod b/go.mod
index 586c72f..a75e0de 100644
--- a/go.mod
+++ b/go.mod
@@ -125,7 +125,7 @@ require (
 	golang.org/x/oauth2 v0.8.0 // indirect
 	golang.org/x/sys v0.14.0 // indirect
 	golang.org/x/term v0.13.0 // indirect
-	golang.org/x/text v0.13.0 // indirect
+	golang.org/x/text v0.13.0
 	golang.org/x/time v0.3.0 // indirect
 	golang.org/x/tools v0.14.0 // indirect
 	gomodules.xyz/jsonpatch/v2 v2.3.0 // indirect
diff --git a/helm/db-controller/values.yaml b/helm/db-controller/values.yaml
index fafb212..6068556 100644
--- a/helm/db-controller/values.yaml
+++ b/helm/db-controller/values.yaml
@@ -128,7 +128,7 @@ controllerConfig:
   authSource: secret
   # if aws authorization is used iam role must be provided
   #iamRole: rds-role
-  caCertificateIdentifier: "rds-ca-2019"
+  caCertificateIdentifier: "rds-ca-rsa2048-g1"
   dbMultiAZEnabled: false
   dbSubnetGroupNameRef:
   defaultBackupPolicyValue: Bronze
diff --git a/pkg/dbclient/client.go b/pkg/dbclient/client.go
index 722223a..f8b4af8 100644
--- a/pkg/dbclient/client.go
+++ b/pkg/dbclient/client.go
@@ -21,9 +21,14 @@ const (
 	RDSSuperUserRole   = "rds_superuser"
 )
 
-var extensions = []string{"citext", "uuid-ossp",
+var defaultExtensions = []string{"citext", "uuid-ossp",
 	"pgcrypto", "hstore", "pg_stat_statements",
-	"plpgsql", "pg_partman", "hll", "pg_cron"}
+	"plpgsql", "hll"}
+
+var specialExtensionsMap = map[string]func(*client, string, string) error{
+	"pg_partman": (*client).pg_partman,
+	"pg_cron":    (*client).pg_cron,
+}
 
 type client struct {
 	dbType string
@@ -110,8 +115,9 @@ func (pc *client) CreateDataBase(name string) (bool, error) {
 	return pc.CreateDatabase(name)
 }
 
+// unit test override
 var getDefaulExtensions = func() []string {
-	return extensions
+	return defaultExtensions
 }
 
 func (pc *client) CreateDatabase(dbName string) (bool, error) {
@@ -157,21 +163,77 @@ func (pc *client) CreateDefaultExtentions(dbName string) error {
 		}
 		pc.log.Info("created extension " + s)
 	}
-	var exists bool
-	err = db.QueryRow("SELECT EXISTS(SELECT 1 FROM pg_extension WHERE extname = $1)", "pg_cron").Scan(&exists)
+
+	return nil
+}
+
+func (pc *client) CreateSpecialExtentions(dbName string, role string) error {
+	db, err := pc.getDB(dbName)
 	if err != nil {
+		pc.log.Error(err, "could not connect to db", "database", dbName)
 		return err
 	}
-	if exists {
-		// pg_cron extension is enabled - grant permission
-		_, err = db.Exec("GRANT USAGE ON SCHEMA cron TO public")
-		if err != nil {
-			pc.log.Error(err, "could not GRANT USAGE ON SCHEMA cron TO public")
+	pc.log.Info("connected to " + dbName)
+	defer db.Close()
+	for functionName := range specialExtensionsMap {
+		if err := specialExtensionsMap[functionName](pc, dbName, role); err != nil {
+			pc.log.Error(err, "error creating extention %s", functionName)
 			return err
 		}
 	}
 	return nil
 }
+func (pc *client) pg_cron(dbName string, role string) error {
+	// create extension pg_cron and grant usage to public
+	db, err := pc.getDB(dbName)
+	if err != nil {
+		pc.log.Error(err, "Failed to connect to database", "database", dbName)
+		return err
+	}
+	defer db.Close()
+
+	pc.log.Info("Connected to database", "database", dbName)
+
+	_, err = db.Exec(fmt.Sprintf(`
+		CREATE EXTENSION IF NOT EXISTS pg_cron;
+		GRANT USAGE ON SCHEMA cron TO %s;
+	`, pq.QuoteIdentifier(role)))
+	if err != nil {
+		pc.log.Error(err, "Failed to create extension pg_cron and grant usage on schema cron to public", "database", dbName)
+		return fmt.Errorf("failed to create extension pg_cron and grant usage on schema cron to public: %w", err)
+	}
+
+	pc.log.Info("Created pg_cron extension and granted usage on schema cron", "role", role)
+
+	return nil
+}
+func (pc *client) pg_partman(dbName string, role string) error {
+
+	db, err := pc.getDB(dbName)
+	if err != nil {
+		pc.log.Error(err, "could not connect to db", "database", dbName)
+		return err
+	}
+	defer db.Close()
+	createPartman := fmt.Sprintf(`
+		CREATE SCHEMA IF NOT EXISTS partman;
+		CREATE EXTENSION IF NOT EXISTS pg_partman WITH SCHEMA partman;
+		GRANT ALL ON SCHEMA partman TO %s;
+		GRANT ALL ON ALL TABLES IN SCHEMA partman TO %s;
+		GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA partman TO %s;
+		GRANT EXECUTE ON ALL PROCEDURES IN SCHEMA partman TO %s;
+	`, pq.QuoteIdentifier(role), pq.QuoteIdentifier(role),
+		pq.QuoteIdentifier(role), pq.QuoteIdentifier(role))
+
+	_, err = db.Exec(createPartman)
+	if err != nil {
+		pc.log.Error(err, "could not create extension pg_partman")
+		return err
+	}
+	pc.log.Info("Created pg_partmann extension and granted usage on schema partman", "role", role)
+
+	return nil
+}
 
 func (pc *client) ManageSystemFunctions(dbName string, functions map[string]string) error {
 	db, err := pc.getDB(dbName)
diff --git a/pkg/dbclient/interface.go b/pkg/dbclient/interface.go
index 1dd8674..40a1c14 100644
--- a/pkg/dbclient/interface.go
+++ b/pkg/dbclient/interface.go
@@ -7,6 +7,7 @@ type Client interface {
 	CreateUser(username, role, userPassword string) (bool, error)
 	CreateGroup(dbName, username string) (bool, error)
 	CreateDefaultExtentions(dbName string) error
+	CreateSpecialExtentions(dbName string, role string) error
 	RenameUser(oldUsername string, newUsername string) error
 	UpdateUser(oldUsername, newUsername, rolename, password string) error
 	UpdatePassword(username string, userPassword string) error
diff --git a/pkg/pgctl/pgdump.go b/pkg/pgctl/pgdump.go
index 4d50389..3bacda1 100644
--- a/pkg/pgctl/pgdump.go
+++ b/pkg/pgctl/pgdump.go
@@ -133,13 +133,12 @@ func (x *Dump) modifyPgDumpInfo() error {
 		return fmt.Errorf("error running sed command to comment create policy: %w", err)
 	}
 
-	// If pg_cron is installed, grant usage on schema cron to public
-	grantCmd := exec.Command("sed", "-i", "/^CREATE EXTENSION IF NOT EXISTS pg_cron/a GRANT USAGE ON SCHEMA cron TO public;", filePath)
-	grantCmd.Stderr = os.Stderr
-	grantCmd.Stdout = os.Stdout
-
-	if err := grantCmd.Run(); err != nil {
-		return fmt.Errorf("error running sed command to grant usage on schema cron: %w", err)
+	// add if not exists to partman schema creation
+	replaceCmd := exec.Command("sed", "-i", "s/CREATE SCHEMA partman;/CREATE SCHEMA IF NOT EXISTS partman;/", filePath)
+	replaceCmd.Stderr = os.Stderr
+	replaceCmd.Stdout = os.Stdout
+	if err := replaceCmd.Run(); err != nil {
+		return fmt.Errorf("error running sed command add if not exists to partman schema creation: %w", err)
 	}
 
 	return nil