feat: AWS RDS pg_repack unit tests

Author: tkcontiant

go.mod

@@ -16,6 +16,7 @@ require (
 
 require (
 	cel.dev/expr v0.24.0 // indirect
+	github.com/DATA-DOG/go-sqlmock v1.5.2 // indirect
 	github.com/Masterminds/semver/v3 v3.4.0 // indirect
 	github.com/antlr4-go/antlr/v4 v4.13.0 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect

go.sum

@@ -1,5 +1,7 @@
 cel.dev/expr v0.24.0 h1:56OvJKSH3hDGL0ml5uSxZmz3/3Pq4tJ+fb1unVLAFcY=
 cel.dev/expr v0.24.0/go.mod h1:hLPLo1W4QUmuYdA72RBX06QTs6MXw941piREPl3Yfiw=
+github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU=
+github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU=
 github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0=
 github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM=
 github.com/antlr4-go/antlr/v4 v4.13.0 h1:lxCg3LAv+EUK6t1i0y1V6/SLeUi0eKEKdhQAlS8TVTI=
@@ -82,6 +84,7 @@ github.com/joshdk/go-junit v1.0.0 h1:S86cUKIdwBHWwA6xCmFlf3RTLfVXYQfvanM5Uh+K6GE
 github.com/joshdk/go-junit v1.0.0/go.mod h1:TiiV0PqkaNfFXjEiyjWM3XXrhVyCa1K4Zfga6W52ung=
 github.com/json-iterator/go v1.1.12 h1:PV8peI4a0ysnczrg+LtxykD8LfKY9ML6u2jnxaEnrnM=
 github.com/json-iterator/go v1.1.12/go.mod h1:e30LSqwooZae/UwlEbR2852Gd8hjQvJoHmT4TnhNGBo=
+github.com/kisielk/sqlstruct v0.0.0-20201105191214-5f3e10d3ab46/go.mod h1:yyMNCyc/Ib3bDTKd379tNMpB/7/H5TjM2Y9QJ5THLbE=
 github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
 github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
 github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=

pkg/postgres/aws.go

@@ -16,6 +16,9 @@ const (
 	AWS_ALTER_REPACK_DEFAULT_PRIVS_SEQUENCES = `ALTER DEFAULT PRIVILEGES FOR ROLE "%s" IN SCHEMA "repack" GRANT USAGE, SELECT ON SEQUENCES TO PUBLIC`
 )
 
+// Test seam: defaults to GetConnection in production, but can be overridden in unit tests.
+var awsGetConnection = GetConnection
+
 func newAWSPG(postgres *pg) PG {
 	return &awspg{
 		*postgres,
@@ -56,15 +59,19 @@ func (c *awspg) CreateExtension(dbname, extension string) error {
 		return nil
 	}
 
+	return c.applyPgRepackPrivileges(dbname)
+}
+
+func (c *awspg) applyPgRepackPrivileges(dbname string) error {
 	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)
+	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)
+	tmpDb, err := awsGetConnection(c.user, c.pass, c.host, dbname, c.args)
 	if err != nil {
 		return err
 	}

pkg/postgres/aws_test.go

@@ -0,0 +1,68 @@
+package postgres
+
+import (
+	"database/sql"
+	"fmt"
+	"regexp"
+	"testing"
+
+	sqlmock "github.com/DATA-DOG/go-sqlmock"
+)
+
+func TestApplyPgRepackPrivileges(t *testing.T) {
+	originalGetConnection := awsGetConnection
+	defer func() {
+		awsGetConnection = originalGetConnection
+	}()
+
+	mainDB, mainMock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("failed to create main sqlmock: %v", err)
+	}
+	defer mainDB.Close()
+
+	tmpDB, tmpMock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("failed to create tmp sqlmock: %v", err)
+	}
+	defer tmpDB.Close()
+
+	dbname := "test-db-dev"
+	owner := "test-db-dev-owner"
+
+	mainMock.ExpectQuery(regexp.QuoteMeta(fmt.Sprintf(GET_DB_OWNER, dbname))).
+		WillReturnRows(sqlmock.NewRows([]string{"pg_get_userbyid"}).AddRow(owner))
+
+	awsGetConnection = func(user, password, host, database, uriArgs string) (*sql.DB, error) {
+		if database != dbname {
+			t.Fatalf("expected database %s, got %s", dbname, database)
+		}
+		return tmpDB, nil
+	}
+
+	tmpMock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(AWS_ALTER_REPACK_DEFAULT_PRIVS_TABLES, owner))).
+		WillReturnResult(sqlmock.NewResult(0, 0))
+	tmpMock.ExpectExec(regexp.QuoteMeta(fmt.Sprintf(AWS_ALTER_REPACK_DEFAULT_PRIVS_SEQUENCES, owner))).
+		WillReturnResult(sqlmock.NewResult(0, 0))
+
+	c := &awspg{
+		pg: pg{
+			db:   mainDB,
+			host: "localhost:5432",
+			user: "postgres",
+			pass: "postgres",
+			args: "sslmode=disable",
+		},
+	}
+
+	if err := c.applyPgRepackPrivileges(dbname); err != nil {
+		t.Fatalf("expected nil error, got %v", err)
+	}
+
+	if err := mainMock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("main DB expectations were not met: %v", err)
+	}
+	if err := tmpMock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("tmp DB expectations were not met: %v", err)
+	}
+}