Merge pull request #136 from Datameer-Inc/turtle/ELT-16035

Allow writer user to create tables in schema
movetokube · Feb 29, 2024 · bd11372 · bd11372
2 parents 9c05343 + 9383b14
commit bd11372
Show file tree

Hide file tree

Showing 5 changed files with 49 additions and 18 deletions.
diff --git a/pkg/controller/postgres/postgres_controller.go b/pkg/controller/postgres/postgres_controller.go
@@ -220,12 +220,20 @@ func (r *ReconcilePostgres) Reconcile(request reconcile.Request) (_ reconcile.Re
 		}
 
 		// Set privileges on schema
-		err = r.pg.SetSchemaPrivileges(database, owner, reader, schema, readerPrivs, reqLogger)
+		schemaPrivilegesReader := postgres.PostgresSchemaPrivileges{database, owner, reader, schema, readerPrivs, false}
+		err = r.pg.SetSchemaPrivileges(schemaPrivilegesReader, reqLogger)
 		if err != nil {
 			reqLogger.Error(err, fmt.Sprintf("Could not give %s permissions \"%s\"", reader, readerPrivs))
 			continue
 		}
-		err = r.pg.SetSchemaPrivileges(database, owner, writer, schema, writerPrivs, reqLogger)
+		schemaPrivilegesWriter := postgres.PostgresSchemaPrivileges{database, owner, writer, schema, readerPrivs, true}
+		err = r.pg.SetSchemaPrivileges(schemaPrivilegesWriter, reqLogger)
+		if err != nil {
+			reqLogger.Error(err, fmt.Sprintf("Could not give %s permissions \"%s\"", writer, writerPrivs))
+			continue
+		}
+		schemaPrivilegesOwner := postgres.PostgresSchemaPrivileges{database, owner, owner, schema, readerPrivs, true}
+		err = r.pg.SetSchemaPrivileges(schemaPrivilegesOwner, reqLogger)
 		if err != nil {
 			reqLogger.Error(err, fmt.Sprintf("Could not give %s permissions \"%s\"", writer, writerPrivs))
 			continue

diff --git a/pkg/controller/postgres/postgres_controller_test.go b/pkg/controller/postgres/postgres_controller_test.go
@@ -682,10 +682,10 @@ var _ = Describe("ReconcilePostgres", func() {
 					// Expected method calls
 					// customers schema
 					pg.EXPECT().CreateSchema(name, name+"-group", "customers", gomock.Any()).Return(nil).Times(1)
-					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "customers", gomock.Any(), gomock.Any()).Return(nil).Times(2)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "customers", gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).Times(3)
 					// stores schema
 					pg.EXPECT().CreateSchema(name, name+"-group", "stores", gomock.Any()).Return(nil).Times(1)
-					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "stores", gomock.Any(), gomock.Any()).Return(nil).Times(2)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "stores", gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).Times(3)
 				})
 
 				It("should update status", func() {
@@ -707,10 +707,12 @@ var _ = Describe("ReconcilePostgres", func() {
 					// Expected method calls
 					// customers schema errors
 					pg.EXPECT().CreateSchema(name, name+"-group", "customers", gomock.Any()).Return(fmt.Errorf("Could not create schema")).Times(1)
-					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "customers", gomock.Any(), gomock.Any()).Return(nil).Times(0)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "customers", gomock.Any(), gomock.Any() ,gomock.Any()).Return(nil).Times(0)
 					// stores schema
 					pg.EXPECT().CreateSchema(name, name+"-group", "stores", gomock.Any()).Return(nil).Times(1)
-					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "stores", gomock.Any(), gomock.Any()).Return(nil).Times(2)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", name+"-reader", "stores", gomock.Any(), false, gomock.Any()).Return(nil).Times(1)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", name+"-writer", "stores", gomock.Any(), true, gomock.Any()).Return(nil).Times(1)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", name+"-group", "stores", gomock.Any(), true, gomock.Any()).Return(nil).Times(1)
 				})
 
 				It("should update status", func() {
@@ -751,7 +753,7 @@ var _ = Describe("ReconcilePostgres", func() {
 					// Expected method calls
 					// customers schema
 					pg.EXPECT().CreateSchema(name, name+"-group", "customers", gomock.Any()).Return(nil).Times(1)
-					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "customers", gomock.Any(), gomock.Any()).Return(nil).Times(2)
+					pg.EXPECT().SetSchemaPrivileges(name, name+"-group", gomock.Any(), "customers", gomock.Any(), gomock.Any(), gomock.Any()).Return(nil).Times(3)
 					// stores schema already exists
 					pg.EXPECT().CreateSchema(name, name+"-group", "stores", gomock.Any()).Times(0)
 					// Call reconcile

diff --git a/pkg/postgres/database.go b/pkg/postgres/database.go
@@ -14,6 +14,7 @@ const (
 	ALTER_DB_OWNER       = `ALTER DATABASE "%s" OWNER 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"`
 	GRANT_ALL_TABLES     = `GRANT %s ON ALL TABLES IN SCHEMA "%s" TO "%s"`
 	DEFAULT_PRIVS_SCHEMA = `ALTER DEFAULT PRIVILEGES FOR ROLE "%s" IN SCHEMA "%s" GRANT %s ON TABLES TO "%s"`
 	REVOKE_CONNECT		 = `REVOKE CONNECT ON DATABASE "%s" FROM public`
@@ -94,29 +95,38 @@ func (c *pg) CreateExtension(db, extension string, logger logr.Logger) error {
 	return nil
 }
 
-func (c *pg) SetSchemaPrivileges(db, creator, role, schema, privs string, logger logr.Logger) error {
-	tmpDb, err := GetConnection(c.user, c.pass, c.host, db, c.args, logger)
+func (c *pg) SetSchemaPrivileges(schemaPrivileges PostgresSchemaPrivileges, logger logr.Logger) error {
+	tmpDb, err := GetConnection(c.user, c.pass, c.host, schemaPrivileges.DB, c.args, logger)
 	if err != nil {
 		return err
 	}
 	defer tmpDb.Close()
 
 	// Grant role usage on schema
-	_, err = tmpDb.Exec(fmt.Sprintf(GRANT_USAGE_SCHEMA, schema, role))
+	_, err = tmpDb.Exec(fmt.Sprintf(GRANT_USAGE_SCHEMA, schemaPrivileges.Schema, schemaPrivileges.Role))
 	if err != nil {
 		return err
 	}
 
 	// Grant role privs on existing tables in schema
-	_, err = tmpDb.Exec(fmt.Sprintf(GRANT_ALL_TABLES, privs, schema, role))
+	_, err = tmpDb.Exec(fmt.Sprintf(GRANT_ALL_TABLES, schemaPrivileges.Privs, schemaPrivileges.Schema, schemaPrivileges.Role))
 	if err != nil {
 		return err
 	}
 
 	// Grant role privs on future tables in schema
-	_, err = tmpDb.Exec(fmt.Sprintf(DEFAULT_PRIVS_SCHEMA, creator, schema, privs, role))
+	_, err = tmpDb.Exec(fmt.Sprintf(DEFAULT_PRIVS_SCHEMA, schemaPrivileges.Creator, schemaPrivileges.Schema, schemaPrivileges.Privs, schemaPrivileges.Role))
 	if err != nil {
 		return err
 	}
+
+	// Grant role usage on schema if createSchema
+	if schemaPrivileges.CreateSchema {
+		_, err = tmpDb.Exec(fmt.Sprintf(GRANT_CREATE_TABLE, schemaPrivileges.Schema, schemaPrivileges.Role))
+		if err != nil {
+			return err
+			}
+	}
+
 	return nil
 }
diff --git a/pkg/postgres/mock/postgres.go b/pkg/postgres/mock/postgres.go
diff --git a/pkg/postgres/postgres.go b/pkg/postgres/postgres.go
@@ -16,7 +16,7 @@ type PG interface {
 	CreateUserRole(role, password string) (string, error)
 	UpdatePassword(role, password string) error
 	GrantRole(role, grantee string) error
-	SetSchemaPrivileges(db, creator, role, schema, privs string, logger logr.Logger) error
+	SetSchemaPrivileges(schemaPrivileges PostgresSchemaPrivileges, logger logr.Logger) error
 	RevokeRole(role, revoked string) error
 	AlterDefaultLoginRole(role, setRole string) error
 	DropDatabase(db string, logger logr.Logger) error
@@ -35,6 +35,15 @@ type pg struct {
 	default_database string
 }
 
+type PostgresSchemaPrivileges struct {
+	DB           string
+	Creator      string
+	Role         string
+	Schema       string
+	Privs        string
+	CreateSchema bool
+}
+
 func NewPG(host, user, password, uri_args, default_database, cloud_type string, logger logr.Logger) (PG, error) {
 	db, err := GetConnection(user, password, host, default_database, uri_args, logger)
 	if err != nil {