Files
dolthub--doltgresql/testing/go/replication_test.go
T
2026-07-13 12:32:25 +08:00

786 lines
27 KiB
Go

// Copyright 2024 Dolthub, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package _go
import (
"context"
"fmt"
"log"
"net"
"os"
"strings"
"testing"
"time"
"github.com/cockroachdb/errors"
"github.com/dolthub/go-mysql-server/sql"
"github.com/jackc/pgx/v5"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/dolthub/doltgresql/server/logrepl"
)
// special pseudo-queries for orchestrating replication tests
const (
createReplicationSlot = "createReplicationSlot"
dropReplicationSlot = "dropReplicationSlot"
stopReplication = "stopReplication"
startReplication = "startReplication"
waitForCatchup = "waitForCatchup"
sleep = "sleep"
)
type ReplicationTest struct {
// Name of the script.
Name string
// The database to create and use. If not provided, then it defaults to "postgres".
Database string
// The SQL statements to execute as setup, in order. Results are not checked, but statements must not error.
// An initial comment can be used to Setup is always run on the primary.
SetUpScript []string
// The set of assertions to make after setup, in order. These should not use ExpectedRaw, as it is not handled.
Assertions []ScriptTestAssertion
// When using RunScripts, setting this on one (or more) tests causes RunScripts to ignore all tests that have this
// set to false (which is the default value). This allows a developer to easily "focus" on a specific test without
// having to comment out other tests, pull it into a different function, etc. In addition, CI ensures that this is
// false before passing, meaning this prevents the commented-out situation where the developer forgets to uncomment
// their code.
Focus bool
// Skip is used to completely skip a test including setup
Skip bool
}
var replicationTests = []ReplicationTest{
{
Name: "simple replication, strings and integers",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"INSERT INTO public.test VALUES (1, 'one')",
"INSERT INTO public.test VALUES (2, 'two')",
"UPDATE public.test SET name = 'three' WHERE id = 2",
"DELETE FROM public.test WHERE id = 1",
"INSERT INTO public.test VALUES (3, 'one')",
"INSERT INTO public.test VALUES (4, 'two')",
"UPDATE public.test SET name = 'five' WHERE id = 4",
"DELETE FROM public.test WHERE id = 3",
"INSERT INTO public.test VALUES (5, 'one')",
"INSERT INTO public.test VALUES (6, 'two')",
"UPDATE public.test SET name = 'six' WHERE id = 6",
"DELETE FROM public.test WHERE id = 5",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "five"},
{int32(6), "six"},
},
},
},
},
{
Name: "stale start",
SetUpScript: []string{
// Postgres will not start tracking which WAL locations to send until the replication slot is created, so we have
// to do that first. Customers have the same constraint: they must import any table data that existed before
// they create the replication slot.
dropReplicationSlot,
createReplicationSlot,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"INSERT INTO public.test VALUES (1, 'one')",
"INSERT INTO public.test VALUES (2, 'two')",
"UPDATE public.test SET name = 'three' WHERE id = 2",
"DELETE FROM public.test WHERE id = 1",
"INSERT INTO public.test VALUES (3, 'one')",
"INSERT INTO public.test VALUES (4, 'two')",
"UPDATE public.test SET name = 'five' WHERE id = 4",
"DELETE FROM public.test WHERE id = 3",
"INSERT INTO public.test VALUES (5, 'one')",
"INSERT INTO public.test VALUES (6, 'two')",
"UPDATE public.test SET name = 'six' WHERE id = 6",
"DELETE FROM public.test WHERE id = 5",
startReplication,
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "five"},
{int32(6), "six"},
},
},
},
},
{
Name: "stopping and resuming replication",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"INSERT INTO public.test VALUES (1, 'one')",
"INSERT INTO public.test VALUES (2, 'two')",
waitForCatchup,
stopReplication,
"UPDATE public.test SET name = 'three' WHERE id = 2",
"DELETE FROM public.test WHERE id = 1",
"INSERT INTO public.test VALUES (3, 'one')",
"INSERT INTO public.test VALUES (4, 'two')",
"UPDATE public.test SET name = 'five' WHERE id = 4",
"DELETE FROM public.test WHERE id = 3",
startReplication,
"INSERT INTO public.test VALUES (5, 'one')",
"INSERT INTO public.test VALUES (6, 'two')",
"UPDATE public.test SET name = 'six' WHERE id = 6",
"DELETE FROM public.test WHERE id = 5",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "five"},
{int32(6), "six"},
},
},
},
},
{
Name: "extended stop/start",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"INSERT INTO public.test VALUES (1, 'one')",
"INSERT INTO public.test VALUES (2, 'two')",
"UPDATE public.test SET name = 'three' WHERE id = 2",
"DELETE FROM public.test WHERE id = 1",
"INSERT INTO public.test VALUES (3, 'one')",
"INSERT INTO public.test VALUES (4, 'two')",
"UPDATE public.test SET name = 'five' WHERE id = 4",
"DELETE FROM public.test WHERE id = 3",
"INSERT INTO public.test VALUES (5, 'one')",
startReplication,
"INSERT INTO public.test VALUES (6, 'two')",
"UPDATE public.test SET name = 'six' WHERE id = 6",
stopReplication,
"DELETE FROM public.test WHERE id = 5",
"INSERT INTO public.test VALUES (7, 'one')",
"INSERT INTO public.test VALUES (8, 'two')",
startReplication,
"UPDATE public.test SET name = 'nine' WHERE id = 8",
"DELETE FROM public.test WHERE id = 7",
"INSERT INTO public.test VALUES (9, 'one')",
stopReplication,
startReplication,
"INSERT INTO public.test VALUES (10, 'two')",
"UPDATE public.test SET name = 'eleven' WHERE id = 10",
stopReplication,
"DELETE FROM public.test WHERE id = 9",
"INSERT INTO public.test VALUES (11, 'one')",
"INSERT INTO public.test VALUES (12, 'two')",
"UPDATE public.test SET name = 'thirteen' WHERE id = 12",
"DELETE FROM public.test WHERE id = 11",
startReplication,
"INSERT INTO public.test VALUES (13, 'one')",
"INSERT INTO public.test VALUES (14, 'two')",
"UPDATE public.test SET name = 'fifteen' WHERE id = 14",
"DELETE FROM public.test WHERE id = 13",
waitForCatchup,
// since replication lag is a heuristic, this sleep is necessary to ensure that the replica has caught
// up in all cases before we shut it off
sleep,
stopReplication,
// below this point we don't expect to find any values replicated because replication was stopped
"INSERT INTO public.test VALUES (15, 'one')",
"INSERT INTO public.test VALUES (16, 'two')",
"UPDATE public.test SET name = 'seventeen' WHERE id = 16",
"DELETE FROM public.test WHERE id = 15",
sleep, // final sleep to make sure that any replication events that will arrive have
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "five"},
{int32(6), "six"},
{int32(8), "nine"},
{int32(10), "eleven"},
{int32(12), "thirteen"},
{int32(14), "fifteen"},
},
},
},
},
{
Name: "all supported types",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100), u_id uuid, age INT, height FLOAT)",
"drop table if exists public.test",
"create table public.test (id INT primary key, name varchar(100), u_id uuid, age INT, height FLOAT)",
"INSERT INTO public.test VALUES (1, 'one', '5ef34887-e635-4c9c-a994-97b1cb810786', 1, 1.1)",
"INSERT INTO public.test VALUES (2, 'two', '2de55648-76ec-4f66-9fae-bd3d853fb0da', 2, 2.2)",
"UPDATE public.test SET name = 'three' WHERE id = 2",
"update public.test set u_id = '3232abe7-560b-4714-a020-2b1a11a1ec65' where id = 2",
"DELETE FROM public.test WHERE id = 1",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three", "3232abe7-560b-4714-a020-2b1a11a1ec65", int32(2), 2.2},
},
},
},
},
{
Name: "concurrent writes",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"/* primary a */ START TRANSACTION",
"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
"/* primary b */ START TRANSACTION",
"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
"/* primary a */ DELETE FROM public.test WHERE id = 1",
"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
"/* primary b */ DELETE FROM public.test WHERE id = 3",
"/* primary b */ COMMIT",
"/* primary a */ COMMIT",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "five"},
},
},
},
},
{
Name: "concurrent writes with restarts",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"/* primary a */ START TRANSACTION",
"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
stopReplication,
"/* primary b */ START TRANSACTION",
"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
"/* primary c */ START TRANSACTION",
"/* primary c */ INSERT INTO public.test VALUES (5, 'one')",
"/* primary c */ INSERT INTO public.test VALUES (6, 'two')",
"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
startReplication,
"/* primary a */ DELETE FROM public.test WHERE id = 1",
"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
"/* primary b */ DELETE FROM public.test WHERE id = 3",
"/* primary b */ COMMIT",
stopReplication,
"/* primary c */ UPDATE public.test SET name = 'seven' WHERE id > 0",
"/* primary c */ DELETE FROM public.test WHERE id = 5",
"/* primary a */ COMMIT",
startReplication,
"/* primary c */ COMMIT",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "seven"},
{int32(6), "seven"},
},
},
},
},
{
Name: "concurrent writes with rollbacks",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"/* primary a */ START TRANSACTION",
"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
stopReplication,
"/* primary b */ START TRANSACTION",
"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
"/* primary c */ START TRANSACTION",
"/* primary c */ INSERT INTO public.test VALUES (5, 'one')",
"/* primary c */ INSERT INTO public.test VALUES (6, 'two')",
"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
startReplication,
"/* primary a */ DELETE FROM public.test WHERE id = 1",
"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
"/* primary b */ DELETE FROM public.test WHERE id = 3",
"/* primary b */ COMMIT",
stopReplication,
"/* primary c */ UPDATE public.test SET name = 'seven' WHERE id > 0",
"/* primary c */ DELETE FROM public.test WHERE id = 5",
"/* primary a */ ROLLBACK",
startReplication,
"/* primary c */ COMMIT",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(4), "seven"},
{int32(6), "seven"},
},
},
},
},
{
Name: "concurrent writes, stale commits",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"/* primary a */ START TRANSACTION",
"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
"/* primary b */ START TRANSACTION",
"/* primary b */ INSERT INTO public.test VALUES (2, 'two')",
"/* primary b */ COMMIT",
waitForCatchup,
stopReplication,
startReplication,
// this tx includes several WAL locations before our last flush, but it must still be replicated
"/* primary a */ COMMIT",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(1), "one"},
{int32(2), "two"},
},
},
},
},
{
Name: "concurrent writes, very stale commits",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100))",
"drop table if exists public.test",
"CREATE TABLE public.test (id INT primary key, name varchar(100))",
"/* primary a */ START TRANSACTION",
"/* primary a */ INSERT INTO public.test VALUES (1, 'one')",
"/* primary a */ INSERT INTO public.test VALUES (2, 'two')",
"/* primary a */ UPDATE public.test SET name = 'three' WHERE id > 0",
"/* primary a */ DELETE FROM public.test WHERE id = 1",
"/* primary b */ START TRANSACTION",
"/* primary b */ INSERT INTO public.test VALUES (3, 'one')",
"/* primary b */ INSERT INTO public.test VALUES (4, 'two')",
"/* primary c */ START TRANSACTION",
"/* primary c */ INSERT INTO public.test VALUES (5, 'one')",
"/* primary c */ INSERT INTO public.test VALUES (6, 'two')",
"/* primary c */ UPDATE public.test SET name = 'seven' WHERE id > 0",
"/* primary c */ DELETE FROM public.test WHERE id = 5",
"/* primary c */ COMMIT",
"/* primary b */ UPDATE public.test SET name = 'five' WHERE id > 0",
"/* primary b */ DELETE FROM public.test WHERE id = 3",
"/* primary b */ COMMIT",
waitForCatchup,
stopReplication,
startReplication,
// this tx includes several WAL locations before our last flush, but it must still be replicated
"/* primary a */ COMMIT",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three"},
{int32(4), "five"},
{int32(6), "five"},
},
},
},
},
{
Name: "all types",
SetUpScript: []string{
dropReplicationSlot,
createReplicationSlot,
startReplication,
"/* replica */ drop table if exists public.test",
"/* replica */ create table public.test (id INT primary key, name varchar(100), age INT, is_cool BOOLEAN, height FLOAT, birth_date DATE, birth_timestamp TIMESTAMP)",
"drop table if exists public.test",
"create table public.test (id INT primary key, name varchar(100), age INT, is_cool BOOLEAN, height FLOAT, birth_date DATE, birth_timestamp TIMESTAMP)",
"INSERT INTO public.test VALUES (1, 'one', 1, true, 1.1, '2021-01-01', '2021-01-01 12:00:00')",
"INSERT INTO public.test VALUES (2, 'two', 2, false, 2.2, '2021-02-02', '2021-02-02 13:00:00')",
"UPDATE public.test SET name = 'three' WHERE id = 2",
"DELETE FROM public.test WHERE id = 1",
waitForCatchup,
},
Assertions: []ScriptTestAssertion{
{
Query: "/* replica */ SELECT * FROM public.test order by id",
Expected: []sql.Row{
{int32(2), "three", int32(2), "f", 2.2, "2021-02-02", "2021-02-02 13:00:00"},
},
},
},
},
}
func TestReplication(t *testing.T) {
if _, ok := os.LookupEnv("GITHUB_ACTION"); !ok {
t.Skip("Replication tests are only run from within GitHub Actions, so we're skipping for local testing")
}
RunReplicationScripts(t, replicationTests)
}
// RunScripts runs the given collection of scripts.
func RunReplicationScripts(t *testing.T, scripts []ReplicationTest) {
// First, we'll run through the scripts to check for the Focus variable. If it's true, then append it to the new slice.
focusScripts := make([]ReplicationTest, 0, len(scripts))
for _, script := range scripts {
if script.Focus {
// If this is running in GitHub Actions, then we'll panic, because someone forgot to disable it before committing
if _, ok := os.LookupEnv("GITHUB_ACTION"); ok {
panic(fmt.Sprintf("The script `%s` has Focus set to `true`. GitHub Actions requires that "+
"all tests are run, which Focus circumvents, leading to this error. Please disable Focus on "+
"all tests.", script.Name))
}
focusScripts = append(focusScripts, script)
}
}
// If we have scripts with Focus set, then we replace the normal script slice with the new slice.
if len(focusScripts) > 0 {
scripts = focusScripts
}
primaryDns := fmt.Sprintf("postgres://postgres:password@127.0.0.1:%d/%s?sslmode=disable&replication=database", localPostgresPort, "postgres")
// We drop and recreate the replication slot once at the beginning of the test suite. Postgres seems to do a little
// work in the background with a publication, so we need to wait a little bit before running any test scripts.
require.NoError(t, logrepl.DropPublication(primaryDns, slotName))
require.NoError(t, logrepl.CreatePublication(primaryDns, slotName))
time.Sleep(500 * time.Millisecond)
for _, script := range scripts {
RunReplicationScript(t, script)
}
}
const slotName = "doltgres_slot"
const localPostgresPort = 5432
// RunReplicationScript runs the given ReplicationTest.
func RunReplicationScript(t *testing.T, script ReplicationTest) {
scriptDatabase := script.Database
if len(scriptDatabase) == 0 {
scriptDatabase = "postgres"
}
database := "postgres"
// primaryDns is the connection to the actual postgres (not doltgres) database, which is why we use port 5342.
// If you have postgres running on a different port, you'll need to change this.
primaryDns := fmt.Sprintf("postgres://postgres:password@127.0.0.1:%d/%s?sslmode=disable", localPostgresPort, database)
ctx, replicaConn, controller := CreateServer(t, scriptDatabase)
defer func() {
replicaConn.Close(ctx)
controller.Stop()
err := controller.WaitForStop()
require.NoError(t, err)
}()
ctx = context.Background()
t.Run(script.Name, func(t *testing.T) {
runReplicationScript(ctx, t, script, replicaConn.Default, primaryDns)
})
}
func newReplicator(t *testing.T, walFilePath string, replicaConn *pgx.Conn, primaryDns string) *logrepl.LogicalReplicator {
connString := replicaConn.PgConn().Conn().RemoteAddr().String()
_, port, err := net.SplitHostPort(connString)
require.NoError(t, err)
replicaDns := fmt.Sprintf("postgres://postgres:password@127.0.0.1:%s/", port)
r, err := logrepl.NewLogicalReplicator(walFilePath, primaryDns, replicaDns)
require.NoError(t, err)
return r
}
// runReplicationScript runs the script given on the postgres connection provided. This does not handle assertions that
// use ExpectedRaw.
func runReplicationScript(
ctx context.Context,
t *testing.T,
script ReplicationTest,
replicaConn *pgx.Conn,
primaryDns string,
) {
walFile := fmt.Sprintf("%s/%s", t.TempDir(), "wal")
r := newReplicator(t, walFile, replicaConn, primaryDns)
defer r.Stop()
if script.Skip {
t.Skip("Skip has been set in the script")
}
connections := map[string]*pgx.Conn{
"replica": replicaConn,
}
defer func() {
for _, conn := range connections {
if conn != nil {
conn.Close(ctx)
}
}
}()
// Run the setup
for _, query := range script.SetUpScript {
// handle logic for special pseudo-queries
if handlePseudoQuery(t, query, r) {
continue
}
conn := connectionForQuery(t, query, connections, primaryDns)
log.Println("Running setup query:", query)
_, err := conn.Exec(ctx, query)
require.NoError(t, err)
}
// Run the assertions
for _, assertion := range script.Assertions {
t.Run(assertion.Query, func(t *testing.T) {
if assertion.Skip {
t.Skip("Skip has been set in the assertion")
}
// handle logic for special pseudo-queries
if handlePseudoQuery(t, assertion.Query, r) {
return
}
target, _ := clientSpecFromQueryComment(assertion.Query)
enableRetries := target == "replica"
conn := connectionForQuery(t, assertion.Query, connections, primaryDns)
numRetries := 3
for retries := 0; retries < numRetries; retries++ {
// If we're skipping the results check, then we call Execute, as it uses a simplified message model.
if assertion.SkipResultsCheck || assertion.ExpectedErr != "" {
_, err := conn.Exec(ctx, assertion.Query, assertion.BindVars...)
if assertion.ExpectedErr != "" {
require.Error(t, err)
assert.Contains(t, err.Error(), assertion.ExpectedErr)
} else {
require.NoError(t, err)
}
} else {
rows, err := conn.Query(ctx, assertion.Query, assertion.BindVars...)
require.NoError(t, err)
readRows, _, err := ReadRows(rows, true)
require.NoError(t, err)
normalizedRows := NormalizeExpectedRow(rows.FieldDescriptions(), assertion.Expected)
// For queries against the replica, whether or not replication is caught up is a heuristic that can be
// incorrect. So we retry queries with sleeps in between to give replication a chance to catch up when this
// happens.
if !assert.ObjectsAreEqual(normalizedRows, readRows) {
if enableRetries && retries < numRetries-1 {
log.Println("Assertion failed, retrying")
time.Sleep(500 * time.Millisecond)
continue
} else {
assert.Equal(t, normalizedRows, readRows)
break
}
}
}
}
})
}
}
// connectionForQuery returns the connection to use for the given query
func connectionForQuery(t *testing.T, query string, connections map[string]*pgx.Conn, primaryDns string) *pgx.Conn {
target, client := clientSpecFromQueryComment(query)
var conn *pgx.Conn
switch target {
case "primary":
conn = connections[client]
if conn == nil {
var err error
conn, err = pgx.Connect(context.Background(), primaryDns)
require.NoError(t, err)
connections[client] = conn
}
case "replica":
conn = connections["replica"]
default:
require.Fail(t, "Invalid target in setup script: ", target)
}
return conn
}
// handlePseudoQuery handles special pseudo-queries that are used to orchestrate replication tests and returns whether
// one was handled.
func handlePseudoQuery(t *testing.T, query string, r *logrepl.LogicalReplicator) bool {
switch query {
case createReplicationSlot:
require.NoError(t, r.CreateReplicationSlotIfNecessary(slotName))
return true
case dropReplicationSlot:
require.NoError(t, r.DropReplicationSlot(slotName))
return true
case startReplication:
go func() {
require.NoError(t, r.StartReplication(slotName))
}()
require.NoError(t, waitForRunning(r))
return true
case stopReplication:
r.Stop()
return true
case waitForCatchup:
require.NoError(t, waitForCaughtUp(r))
return true
case sleep:
time.Sleep(200 * time.Millisecond)
return true
}
return false
}
// clientSpecFromQueryComment returns "replica" if the query is meant to be run on the replica, and "primary" if it's meant
// to be run on the primary, based on the comment in the query. If not comment, the query runs on the primary
func clientSpecFromQueryComment(query string) (string, string) {
startCommentIdx := strings.Index(query, "/*")
endCommentIdx := strings.Index(query, "*/")
if startCommentIdx < 0 || endCommentIdx < 0 {
return "primary", "a"
}
query = query[startCommentIdx+2 : endCommentIdx]
if strings.Contains(query, "replica") {
return "replica", "a"
}
if i := strings.Index(query, "primary "); i > 0 && i+len("primary ") < len(query) {
return "primary", query[i+len("primary "):]
}
return "primary", "a"
}
func waitForRunning(r *logrepl.LogicalReplicator) error {
start := time.Now()
for {
if r.Running() {
break
}
if time.Since(start) > 500*time.Millisecond {
return errors.New("Replication did not start")
}
time.Sleep(5 * time.Millisecond)
}
return nil
}
func waitForCaughtUp(r *logrepl.LogicalReplicator) error {
log.Println("Waiting for replication to catch up")
start := time.Now()
for {
if caughtUp, err := r.CaughtUp(150); caughtUp {
log.Println("replication caught up")
break
} else if err != nil {
return err
}
log.Println("replication not caught up, waiting")
if time.Since(start) >= 2*time.Second {
return errors.New("Replication did not catch up")
}
time.Sleep(20 * time.Millisecond)
}
return nil
}