Add Support for Querying Against Read Replicas

main_test.go

@@ -20,7 +20,13 @@ const (
 	postgresImage = "docker.io/postgres:16.0-alpine"
 )
 
-var postgresDataSource string
+// Connection strings for the three databases created in [TestMain]
+// These are used in other tests
+var (
+	postgresDataSource    string
+	postgresDataSourceRR1 string
+	postgresDataSourceRR2 string
+)
 
 func withSchemaSQL() testcontainers.CustomizeRequestOption {
 	return func(req *testcontainers.GenericContainerRequest) error {
@@ -34,24 +40,69 @@ func withSchemaSQL() testcontainers.CustomizeRequestOption {
 }
 
 func TestMain(m *testing.M) {
+	ctx := context.Background()
+	cleanups, err := createPostgresContainers(ctx)
+	defer func() {
+		for _, cleanup := range cleanups {
+			if cleanup != nil {
+				cleanup()
+			}
+		}
+	}()
+	if err != nil {
+		fmt.Printf("did not create postgres containers: %v\n", err)
+	}
+
+	// nolint:gocritic // The docs for Run specify that the returned int is to be passed to os.Exit
+	os.Exit(m.Run())
+}
+
+// createPostgresContainers creates 3 databases as containers. The first is intended to mimic a master database and
+// the last 2 are intended to mimic read replicas.
+// A []func() is returned to cleanups.
+// Package-level connection strings are set.
+func createPostgresContainers(ctx context.Context) ([]func(), error) {
+	// Database connection strings are the same, except the port
+	connString := func(port string) string {
+		return fmt.Sprintf("postgres://%s:%s@localhost:%s/%s?sslmode=disable&application_name=test", dbUser, dbPassword, port, dbName)
+	}
+
 	var cleanups []func()
-	cleanup := func(fn func()) {
-		cleanups = append(cleanups, fn)
+
+	// Create the master database
+	cM, mpM, err := createPostgresContainer(ctx, "init-db.sh")
+	cleanups = append(cleanups, cM)
+	if err != nil {
+		return cleanups, err
 	}
-	fatal := func(args ...any) {
-		fmt.Fprintln(os.Stderr, args...)
-		for _, fn := range cleanups {
-			fn()
-		}
-		os.Exit(1)
+	postgresDataSource = connString(mpM)
+
+	// Create first read replica
+	cRR1, mpRR1, err := createPostgresContainer(ctx, "init-db-rr.sh")
+	cleanups = append(cleanups, cRR1)
+	if err != nil {
+		return cleanups, err
 	}
+	postgresDataSourceRR1 = connString(mpRR1)
 
-	ctx := context.Background()
+	// Create second read replica
+	cRR2, mpRR2, err := createPostgresContainer(ctx, "init-db-rr.sh")
+	cleanups = append(cleanups, cRR2)
+	if err != nil {
+		return cleanups, err
+	}
+	postgresDataSourceRR2 = connString(mpRR2)
+
+	return cleanups, nil
+}
+
+// createPostgresContainer creates a single postgres container on the specified port
+func createPostgresContainer(ctx context.Context, initFilename string) (func(), string, error) {
 	postgresContainer, err := postgres.Run(ctx, postgresImage,
 		postgres.WithDatabase(dbName),
 		postgres.WithUsername(dbUser),
 		postgres.WithPassword(dbPassword),
-		postgres.WithInitScripts(filepath.Join("testdata", "init-db.sh")),
+		postgres.WithInitScripts(filepath.Join("testdata", initFilename)),
 		withSchemaSQL(),
 		testcontainers.WithWaitStrategy(
 			wait.ForLog("database system is ready to accept connections").
@@ -60,28 +111,27 @@ func TestMain(m *testing.M) {
 		),
 	)
 	if err != nil {
-		fatal("error creating postgres container:", err)
+		return nil, "", fmt.Errorf("error creating postgres container: %w", err)
 	}
-	cleanup(func() {
+
+	cleanup := func() {
 		if err := postgresContainer.Terminate(ctx); err != nil {
 			fmt.Fprintln(os.Stderr, "error terminating postgres:", err)
 		}
-	})
+	}
 
 	postgresState, err := postgresContainer.State(ctx)
 	if err != nil {
-		fatal(err)
+		return cleanup, "", fmt.Errorf("checking container state: %w", err)
 	}
 	if !postgresState.Running {
-		fatal("Postgres status:", postgresState.Status)
+		return cleanup, "", fmt.Errorf("Postgres status %q is not \"running\"", postgresState.Status)
 	}
 
-	postgresPort, err := postgresContainer.MappedPort(ctx, "5432/tcp")
+	mp, err := postgresContainer.MappedPort(ctx, "5432/tcp")
 	if err != nil {
-		fatal(err)
+		return cleanup, "", fmt.Errorf("mapped port 5432/tcp does not seem to be available: %w", err)
 	}
 
-	postgresDataSource = fmt.Sprintf("postgres://%s:%s@localhost:%s/%s?sslmode=disable&application_name=test", dbUser, dbPassword, postgresPort.Port(), dbName)
-
-	os.Exit(m.Run())
+	return cleanup, mp.Port(), nil
 }

read_replicas.go

@@ -0,0 +1,87 @@
+package sequel
+
+import (
+	"errors"
+	"sync"
+
+	"github.com/go-sqlx/sqlx"
+)
+
+var ErrNoReadReplicaConnection = errors.New("no read replica connections")
+
+type readReplica struct {
+	db *sqlx.DB
+
+	next *readReplica
+}
+
+// readReplicas contains a set of DB connections. It is intended to give fair round robin access
+// through a circular singularly linked list.
+//
+// Replicas are appended after the current one. The intended use is to build the replica set before
+// querying, but all operations are concurrent-safe.
+type readReplicas struct {
+	m sync.Mutex
+
+	current *readReplica
+}
+
+// add adds a DB to the collection of read replicas
+func (rr *readReplicas) add(db *sqlx.DB) {
+	rr.m.Lock()
+	defer rr.m.Unlock()
+
+	r := readReplica{
+		db: db,
+	}
+
+	// Empty ring, add new DB
+	if rr.current == nil {
+		r.next = &r
+
+		rr.current = &r
+
+		return
+	}
+
+	// Insert new db after current
+	n := rr.current.next
+	r.next = n
+	rr.current.next = &r
+}
+
+// next returns the current DB. The current pointer is advanced.
+func (rr *readReplicas) next() (*sqlx.DB, error) {
+	rr.m.Lock()
+	defer rr.m.Unlock()
+
+	if rr.current == nil {
+		return nil, ErrNoReadReplicaConnection
+	}
+
+	c := rr.current
+
+	rr.current = rr.current.next
+
+	return c.db, nil
+}
+
+// Close closes all read replica connections
+func (rr *readReplicas) Close() {
+	rr.m.Lock()
+	defer rr.m.Unlock()
+
+	// If this instance has no replicas, current would be nil
+	if rr.current == nil {
+		return
+	}
+
+	first := rr.current
+	for c := first; ; c = c.next {
+		c.db.Close()
+
+		if c.next == first {
+			break
+		}
+	}
+}

read_replicas_test.go

@@ -0,0 +1,145 @@
+package sequel
+
+import (
+	"errors"
+	"slices"
+	"testing"
+
+	"github.com/go-sqlx/sqlx"
+)
+
+func Test_readReplicas_add(t *testing.T) {
+	tests := []struct {
+		name string
+		rrs  []*sqlx.DB
+	}{
+		{
+			name: "add single",
+			rrs:  []*sqlx.DB{sqlx.NewDb(nil, "fakeDriver")},
+		},
+		{
+			name: "add multiple",
+			rrs:  []*sqlx.DB{sqlx.NewDb(nil, "fakeDriver"), sqlx.NewDb(nil, "fakeDriver"), sqlx.NewDb(nil, "fakeDriver")},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// Create readReplicas
+			rr := &readReplicas{}
+
+			// Add connections to read replicas
+			for _, c := range tt.rrs {
+				rr.add(c)
+			}
+
+			// Confirm they are all added
+			if !rrContains(t, rr, tt.rrs) {
+				t.Fatal("readReplicas does not contain all added conns")
+			}
+		})
+	}
+}
+
+func Test_readReplicas_next(t *testing.T) {
+	newRRs := func(rrs ...*sqlx.DB) *readReplicas {
+		rr := &readReplicas{}
+		for _, c := range rrs {
+			rr.add(c)
+		}
+
+		return rr
+	}
+
+	c1, c2, c3 := sqlx.NewDb(nil, "fakeDriver"), sqlx.NewDb(nil, "fakeDriver"), sqlx.NewDb(nil, "fakeDriver")
+
+	emptyRR := newRRs()
+	singleRR := newRRs(c1)
+	threeRR := newRRs(c1, c2, c3)
+
+	tests := []struct {
+		name    string
+		rr      *readReplicas
+		want    *sqlx.DB
+		wantErr error
+	}{
+		{
+			name:    "empty readReplicas",
+			rr:      emptyRR,
+			want:    nil,
+			wantErr: ErrNoReadReplicaConnection,
+		},
+		{
+			name:    "one read replica",
+			rr:      singleRR,
+			want:    c1,
+			wantErr: nil,
+		},
+		{
+			name:    "one read replica (2 calls)",
+			rr:      singleRR,
+			want:    c1,
+			wantErr: nil,
+		},
+		{
+			name:    "one read replica (3 calls)",
+			rr:      singleRR,
+			want:    c1,
+			wantErr: nil,
+		},
+		{
+			name:    "three read replicas",
+			rr:      threeRR,
+			want:    c1,
+			wantErr: nil,
+		},
+		{
+			name:    "three read replicas (2 calls)",
+			rr:      threeRR,
+			want:    c3,
+			wantErr: nil,
+		},
+		{
+			name:    "three read replicas (3 calls)",
+			rr:      threeRR,
+			want:    c2,
+			wantErr: nil,
+		},
+		{
+			name:    "three read replicas (4 calls)",
+			rr:      threeRR,
+			want:    c1,
+			wantErr: nil,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := tt.rr.next()
+			if !errors.Is(err, tt.wantErr) {
+				t.Fatalf("got err %v, expecting %v", err, tt.wantErr)
+			}
+
+			if got != tt.want {
+				t.Fatalf("got rr %p, want %p", got, tt.want)
+			}
+		})
+	}
+}
+
+// rrContains is a test helper to see if a [readReplicas] contains all the connections passed in
+func rrContains(t *testing.T, rr *readReplicas, conns []*sqlx.DB) bool {
+	t.Helper()
+
+	var findCnt int
+	head := rr.current
+	for c := head.next; ; c = c.next {
+		if slices.Contains(conns, c.db) {
+			findCnt++
+		}
+
+		if c == head {
+			break
+		}
+	}
+
+	return findCnt == len(conns)
+}