initialize operator user from environment

Author: michaeldwan

cmd/start/main.go

@@ -2,6 +2,8 @@ package main
 
 import (
 	"bytes"
+	"context"
+	"errors"
 	"fmt"
 	"os"
 	"os/signal"
@@ -12,7 +14,9 @@ import (
 	"time"
 
 	"github.com/fly-examples/postgres-ha/pkg/flypg"
+	"github.com/fly-examples/postgres-ha/pkg/flypg/admin"
 	"github.com/fly-examples/postgres-ha/pkg/supervisor"
+	"github.com/jackc/pgx/v4"
 )
 
 func main() {
@@ -53,16 +57,29 @@ func main() {
 			}
 
 			currentKeeper := cd.Keepers[node.KeeperUID]
-			fmt.Printf("found keeper: %#v\n", currentKeeper)
+			// fmt.Printf("found keeper: %#v\n", currentKeeper)
 			currentDB := cd.FindDB(currentKeeper)
-			fmt.Printf("found db: %#v\n", currentDB)
+			// fmt.Printf("found db: %#v\n", currentDB)
 
 			if currentKeeper == nil || currentDB == nil {
 				continue
 			}
 
-			if currentKeeper.Status.Healthy && currentDB.Status.Healthy {
+			if currentKeeper.Status.Healthy && currentDB.Status.Healthy && currentDB.Spec.Role == "master" {
 				fmt.Println("keeper is healthy, db is healthy, role:", currentDB.Spec.Role)
+
+				pg, err := node.NewLocalConnection(context.TODO())
+				if err != nil {
+					fmt.Println("error connecting to local postgres", err)
+					continue
+				}
+
+				if err = initOperator(context.TODO(), pg, node.OperatorCredentials); err != nil {
+					fmt.Println("error configuring operator:", err)
+					continue
+				}
+				fmt.Println("operator ready!")
+				break
 			}
 		}
 	}()
@@ -153,3 +170,53 @@ func writeStolonctlEnvFile(n *flypg.Node, filename string) {
 
 	os.WriteFile(filename, b.Bytes(), 0644)
 }
+
+func initOperator(ctx context.Context, pg *pgx.Conn, creds flypg.Credentials) error {
+	fmt.Println("configuring operator")
+
+	users, err := admin.ListUsers(ctx, pg)
+	if err != nil {
+		return err
+	}
+
+	var operatorUser *admin.UserInfo
+
+	for _, u := range users {
+		if u.Username == creds.Username {
+			operatorUser = &u
+			break
+		}
+	}
+
+	if operatorUser == nil {
+		fmt.Println("operator user does not exist, creating")
+		err = admin.CreateUser(ctx, pg, creds.Username, creds.Password)
+		if err != nil {
+			return err
+		}
+		operatorUser, err = admin.FindUser(ctx, pg, creds.Username)
+		if err != nil {
+			return err
+		}
+	}
+
+	if operatorUser == nil {
+		return errors.New("error creating operator: user not found")
+	}
+
+	if !operatorUser.SuperUser {
+		fmt.Println("operator is not a superuser, fixing")
+		if err := admin.GrantSuperuser(ctx, pg, creds.Username); err != nil {
+			return err
+		}
+	}
+
+	if !operatorUser.IsPassword(creds.Password) {
+		fmt.Println("operator password does not match config, changing")
+		if err := admin.ChangePassword(ctx, pg, creds.Username, creds.Password); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

docker-compose.yml

@@ -24,6 +24,12 @@ services:
       - consul
     environment:
       FLY_CONSUL_URL: "http://consul:8500/chaos-postgres-wvo4x1opkz9l5ydn/"
+      SU_USERNAME: "flypgadmin"
+      SU_PASSWORD: "supassword"
+      REPL_USERNAME: "repluser"
+      REPL_PASSWORD: "replpassword"
+      OPERATOR_USERNAME: "postgres"
+      OPERATOR_PASSWORD: "operatorpassword"
     ports:
       - "5432:5432"
       - "5433:5433"

pkg/flypg/admin/admin.go

@@ -0,0 +1,142 @@
+package admin
+
+import (
+	"context"
+	"crypto/md5"
+	"fmt"
+	"strings"
+
+	"github.com/jackc/pgx/v4"
+)
+
+func CreateUser(ctx context.Context, pg *pgx.Conn, username string, password string) error {
+	sql := fmt.Sprintf(`CREATE USER %s WITH LOGIN PASSWORD '%s'`, username, password)
+
+	_, err := pg.Exec(ctx, sql)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func GrantSuperuser(ctx context.Context, pg *pgx.Conn, username string) error {
+	sql := fmt.Sprintf("ALTER USER %s WITH SUPERUSER;", username)
+
+	_, err := pg.Exec(ctx, sql)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func ChangePassword(ctx context.Context, pg *pgx.Conn, username, password string) error {
+	sql := fmt.Sprintf("ALTER USER %s WITH LOGIN PASSWORD '%s';", username, password)
+
+	_, err := pg.Exec(ctx, sql)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func ListDatabases(ctx context.Context, pg *pgx.Conn) ([]DbInfo, error) {
+	sql := `
+		SELECT d.datname,
+					(SELECT array_agg(u.usename::text order by u.usename) 
+						from pg_user u 
+						where has_database_privilege(u.usename, d.datname, 'CONNECT')) as allowed_users
+		from pg_database d where d.datistemplate = false
+		order by d.datname;
+		`
+
+	rows, err := pg.Query(ctx, sql)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	values := []DbInfo{}
+
+	for rows.Next() {
+		di := DbInfo{}
+		if err := rows.Scan(&di.Name, &di.Users); err != nil {
+			return nil, err
+		}
+		values = append(values, di)
+	}
+
+	return values, nil
+}
+
+type UserInfo struct {
+	Username     string   `json:"username"`
+	SuperUser    bool     `json:"superuser"`
+	Databases    []string `json:"databases"`
+	PasswordHash string   `json:"-"`
+}
+
+func (ui UserInfo) IsPassword(password string) bool {
+	if !strings.HasPrefix(ui.PasswordHash, "md5") {
+		return false
+	}
+
+	encoded := fmt.Sprintf("md5%x", md5.Sum([]byte(password+ui.Username)))
+	return encoded == ui.PasswordHash
+}
+
+type DbInfo struct {
+	Name  string   `json:"name"`
+	Users []string `json:"users"`
+}
+
+func ListUsers(ctx context.Context, pg *pgx.Conn) ([]UserInfo, error) {
+	sql := `
+		select u.usename,
+			usesuper as superuser,
+			a.rolpassword as passwordhash,
+      (select array_agg(d.datname::text order by d.datname)
+				from pg_database d
+				WHERE datistemplate = false
+				AND has_database_privilege(u.usename, d.datname, 'CONNECT')
+			) as allowed_databases
+			from pg_user u
+			join pg_authid a on u.usesysid = a.oid
+			order by u.usename
+			`
+
+	rows, err := pg.Query(ctx, sql)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+
+	values := []UserInfo{}
+
+	for rows.Next() {
+		ui := UserInfo{}
+		if err := rows.Scan(&ui.Username, &ui.SuperUser, &ui.PasswordHash, &ui.Databases); err != nil {
+			return nil, err
+		}
+		values = append(values, ui)
+	}
+
+	return values, nil
+}
+
+func FindUser(ctx context.Context, pg *pgx.Conn, username string) (*UserInfo, error) {
+	users, err := ListUsers(ctx, pg)
+	if err != nil {
+		return nil, err
+	}
+
+	for _, u := range users {
+		if u.Username == username {
+			return &u, nil
+		}
+	}
+
+	return nil, nil
+}