Replice terminate_backend with WITH (FORCE)

Author: pcallewaert

pkg/postgres/database.go

@@ -12,7 +12,7 @@ const (
 	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"`
+	DROP_DATABASE           = `DROP DATABASE "%s" WITH (FORCE)`
 	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"`
@@ -22,9 +22,9 @@ const (
 	GRANT_ALL_SEQUENCES     = `GRANT %s ON ALL SEQUENCES IN SCHEMA "%s" TO "%s"`
 	DEFAULT_PRIVS_SEQUENCES = `ALTER DEFAULT PRIVILEGES IN SCHEMA "%s" GRANT %s ON SEQUENCES TO "%s"`
 	REVOKE_CONNECT          = `REVOKE CONNECT ON DATABASE "%s" FROM public`
-	TERMINATE_BACKEND       = `SELECT pg_terminate_backend(pg_stat_activity.pid) FROM pg_stat_activity	WHERE pg_stat_activity.datname = '%s' AND pid <> pg_backend_pid()`
 	GET_DB_OWNER            = `SELECT pg_catalog.pg_get_userbyid(d.datdba) FROM pg_catalog.pg_database d WHERE d.datname = '%s'`
 	GRANT_CREATE_SCHEMA     = `GRANT CREATE ON DATABASE "%s" TO "%s"`
+	GRANT_CONNECT           = `GRANT CONNECT ON DATABASE "%s" TO "%s"`
 )
 
 func (c *pg) CreateDB(dbname, role string) error {
@@ -42,9 +42,20 @@ func (c *pg) CreateDB(dbname, role string) error {
 		return err
 	}
 
-	err = c.execute(fmt.Sprintf(GRANT_CREATE_SCHEMA, dbname, role))
-	if err != nil {
-		return err
+	// Grant CREATE on database to owner and operator user
+	usersToGrant := []string{c.user, role}
+	for _, u := range usersToGrant {
+		err = c.execute(fmt.Sprintf(GRANT_CREATE_SCHEMA, dbname, u))
+		if err != nil {
+			return fmt.Errorf("failed to grant create schema on %s to %s: %w", dbname, u, err)
+		}
+	}
+	// Grant CONNECT on database to owner and operator user
+	for _, u := range usersToGrant {
+		err = c.execute(fmt.Sprintf(GRANT_CONNECT, dbname, u))
+		if err != nil {
+			return fmt.Errorf("failed to grant connect on %s to %s: %w", dbname, u, err)
+		}
 	}
 	return nil
 }
@@ -99,11 +110,6 @@ func (c *pg) DropDatabase(database string, logger logr.Logger) error {
 		return err
 	}
 
-	err = c.execute(fmt.Sprintf(TERMINATE_BACKEND, database))
-	// Error code 3D000 is returned if database doesn't exist
-	if err != nil && !isPgError(err, "3D000") {
-		return err
-	}
 	err = c.execute(fmt.Sprintf(DROP_DATABASE, database))
 	// Error code 3D000 is returned if database doesn't exist
 	if err != nil && !isPgError(err, "3D000") {

pkg/postgres/database_grants_test.go

@@ -14,11 +14,14 @@ func TestCreateDB_GrantsCreateOnDatabase(t *testing.T) {
 	}
 	defer db.Close()
 
-	p := &pg{db: db}
+	p := &pg{db: db, user: "operator"}
 
 	mock.ExpectExec(regexp.QuoteMeta(`CREATE DATABASE "mydb"`)).WillReturnResult(sqlmock.NewResult(0, 1))
 	mock.ExpectExec(regexp.QuoteMeta(`ALTER DATABASE "mydb" OWNER TO "owner"`)).WillReturnResult(sqlmock.NewResult(0, 1))
+	mock.ExpectExec(regexp.QuoteMeta(`GRANT CREATE ON DATABASE "mydb" TO "operator"`)).WillReturnResult(sqlmock.NewResult(0, 1))
 	mock.ExpectExec(regexp.QuoteMeta(`GRANT CREATE ON DATABASE "mydb" TO "owner"`)).WillReturnResult(sqlmock.NewResult(0, 1))
+	mock.ExpectExec(regexp.QuoteMeta(`GRANT CONNECT ON DATABASE "mydb" TO "operator"`)).WillReturnResult(sqlmock.NewResult(0, 1))
+	mock.ExpectExec(regexp.QuoteMeta(`GRANT CONNECT ON DATABASE "mydb" TO "owner"`)).WillReturnResult(sqlmock.NewResult(0, 1))
 
 	if err := p.CreateDB("mydb", "owner"); err != nil {
 		t.Fatalf("CreateDB: %v", err)