Extract control replica connection config in throttler

Author: danieljoos

go/logic/throttler.go

@@ -176,6 +176,17 @@ func (thlr *Throttler) collectReplicationLag(firstThrottlingCollected chan<- boo
 	}
 }
 
+// controlReplicaConnectionConfig returns the connection config used to read
+// replication lag from a control replica. In move-tables mode the lag is read
+// from the target cluster's replicas, otherwise from the source (inspector)
+// cluster's replicas.
+func (thlr *Throttler) controlReplicaConnectionConfig(replicaKey mysql.InstanceKey) *mysql.ConnectionConfig {
+	if thlr.migrationContext.IsMoveTablesMode() {
+		return thlr.migrationContext.MoveTables.ConnectionConfig.DuplicateCredentials(replicaKey)
+	}
+	return thlr.migrationContext.InspectorConnectionConfig.DuplicateCredentials(replicaKey)
+}
+
 // collectControlReplicasLag polls all the control replicas to get maximum lag value
 func (thlr *Throttler) collectControlReplicasLag() {
 	if atomic.LoadInt64(&thlr.migrationContext.HibernateUntil) > 0 {
@@ -213,12 +224,7 @@ func (thlr *Throttler) collectControlReplicasLag() {
 		}
 		lagResults := make(chan *mysql.ReplicationLagResult, instanceKeyMap.Len())
 		for replicaKey := range *instanceKeyMap {
-			var connectionConfig *mysql.ConnectionConfig
-			if thlr.migrationContext.IsMoveTablesMode() {
-				connectionConfig = thlr.migrationContext.MoveTables.ConnectionConfig.DuplicateCredentials(replicaKey)
-			} else {
-				connectionConfig = thlr.migrationContext.InspectorConnectionConfig.DuplicateCredentials(replicaKey)
-			}
+			connectionConfig := thlr.controlReplicaConnectionConfig(replicaKey)
 
 			if err := connectionConfig.RegisterTLSConfig(); err != nil {
 				return &mysql.ReplicationLagResult{Err: err}

go/logic/throttler_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/stretchr/testify/assert"
 
 	"github.com/github/gh-ost/go/base"
+	"github.com/github/gh-ost/go/mysql"
 )
 
 func newTestThrottler() *Throttler {
@@ -86,6 +87,40 @@ func TestThrottleReturnsOnContextCancellation(t *testing.T) {
 	}
 }
 
+func TestControlReplicaConnectionConfig(t *testing.T) {
+	replicaKey := mysql.InstanceKey{Hostname: "replica-host", Port: 3307}
+
+	t.Run("uses inspector connection config when not in move-tables mode", func(t *testing.T) {
+		thlr := newTestThrottler()
+		thlr.migrationContext.InspectorConnectionConfig.Key = mysql.InstanceKey{Hostname: "source-host", Port: 3306}
+		thlr.migrationContext.InspectorConnectionConfig.User = "source-user"
+		thlr.migrationContext.InspectorConnectionConfig.Password = "source-pass"
+
+		connectionConfig := thlr.controlReplicaConnectionConfig(replicaKey)
+
+		assert.False(t, thlr.migrationContext.IsMoveTablesMode())
+		assert.Equal(t, replicaKey, connectionConfig.Key)
+		assert.Equal(t, "source-user", connectionConfig.User)
+		assert.Equal(t, "source-pass", connectionConfig.Password)
+	})
+
+	t.Run("uses move-tables connection config when in move-tables mode", func(t *testing.T) {
+		thlr := newTestThrottler()
+		thlr.migrationContext.MoveTables.TableNames = []string{"my_table"}
+		thlr.migrationContext.MoveTables.ConnectionConfig = mysql.NewConnectionConfig()
+		thlr.migrationContext.MoveTables.ConnectionConfig.Key = mysql.InstanceKey{Hostname: "target-host", Port: 3306}
+		thlr.migrationContext.MoveTables.ConnectionConfig.User = "target-user"
+		thlr.migrationContext.MoveTables.ConnectionConfig.Password = "target-pass"
+
+		connectionConfig := thlr.controlReplicaConnectionConfig(replicaKey)
+
+		assert.True(t, thlr.migrationContext.IsMoveTablesMode())
+		assert.Equal(t, replicaKey, connectionConfig.Key)
+		assert.Equal(t, "target-user", connectionConfig.User)
+		assert.Equal(t, "target-pass", connectionConfig.Password)
+	})
+}
+
 func TestThrottleCallsOnThrottledCallback(t *testing.T) {
 	thlr := newTestThrottler()
 	thlr.migrationContext.SetThrottled(true, "test", base.NoThrottleReasonHint)