WIP

Signed-off-by: Tim Vaillancourt <[email protected]>

go/vt/vtorc/logic/tablet_discovery.go

@@ -39,7 +39,6 @@ import (
 	"vitess.io/vitess/go/vt/vtorc/db"
 	"vitess.io/vitess/go/vt/vtorc/inst"
 	"vitess.io/vitess/go/vt/vtorc/process"
-	"vitess.io/vitess/go/vt/vttablet/tmclient"
 )
 
 var (
@@ -60,18 +59,20 @@ func RegisterFlags(fs *pflag.FlagSet) {
 type TabletDiscovery struct {
 	keyspaceShards []*topo.KeyspaceShard
 	refreshTicker  *time.Ticker
+	tm             *TabletManager
 	vtorc          *VTOrc
 }
 
 // OpenTabletDiscovery opens the vitess topo if enables and returns a ticker
 // channel for polling.
-func NewTabletDiscovery(v *VTOrc) (*TabletDiscovery, error) {
+func NewTabletDiscovery(v *VTOrc, tm *TabletManager) (*TabletDiscovery, error) {
 	// Clear existing cache and perform a new refresh.
 	if _, err := db.ExecVTOrc("DELETE FROM vitess_tablet"); err != nil {
 		return nil, err
 	}
 
 	td := &TabletDiscovery{
+		tm:    tm,
 		vtorc: v,
 	}
 
@@ -98,8 +99,7 @@ func (td *TabletDiscovery) Close() {
 	}
 }
 
-func (td *TabletDiscovery) tmClient() tmclient.TabletManagerClient { return td.vtorc.tmc }
-func (td *TabletDiscovery) topoServer() *topo.Server               { return td.vtorc.ts }
+func (td *TabletDiscovery) topoServer() *topo.Server { return td.vtorc.ts }
 
 // Parse input and build list of keyspaces / shards
 func (td *TabletDiscovery) parseKeyspaceShardsToWatch() {
@@ -338,41 +338,6 @@ func (td *TabletDiscovery) LockShard(ctx context.Context, tabletAlias string, lo
 	}, nil
 }
 
-// tabletUndoDemotePrimary calls the said RPC for the given tablet.
-func (td *TabletDiscovery) tabletUndoDemotePrimary(ctx context.Context, tablet *topodatapb.Tablet, semiSync bool) error {
-	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer tmcCancel()
-	return td.tmClient().UndoDemotePrimary(tmcCtx, tablet, semiSync)
-}
-
-// setReadOnly calls the said RPC for the given tablet
-func (td *TabletDiscovery) setReadOnly(ctx context.Context, tablet *topodatapb.Tablet) error {
-	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer tmcCancel()
-	return td.tmClient().SetReadOnly(tmcCtx, tablet)
-}
-
-// changeTabletType calls the said RPC for the given tablet with the given parameters.
-func (td *TabletDiscovery) changeTabletType(ctx context.Context, tablet *topodatapb.Tablet, tabletType topodatapb.TabletType, semiSync bool) error {
-	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer tmcCancel()
-	return td.tmClient().ChangeType(tmcCtx, tablet, tabletType, semiSync)
-}
-
-// resetReplicationParameters resets the replication parameters on the given tablet.
-func (td *TabletDiscovery) resetReplicationParameters(ctx context.Context, tablet *topodatapb.Tablet) error {
-	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer tmcCancel()
-	return td.tmClient().ResetReplicationParameters(tmcCtx, tablet)
-}
-
-// setReplicationSource calls the said RPC with the parameters provided
-func (td *TabletDiscovery) setReplicationSource(ctx context.Context, replica *topodatapb.Tablet, primary *topodatapb.Tablet, semiSync bool, heartbeatInterval float64) error {
-	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
-	defer tmcCancel()
-	return td.tmClient().SetReplicationSource(tmcCtx, replica, primary.Alias, 0, "", true, semiSync, heartbeatInterval)
-}
-
 // shardPrimary finds the primary of the given keyspace-shard by reading the vtorc backend
 func shardPrimary(keyspace string, shard string) (primary *topodatapb.Tablet, err error) {
 	query := `SELECT

go/vt/vtorc/logic/tablet_manager.go

@@ -0,0 +1,68 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+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 logic
+
+import (
+	"context"
+
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/topo"
+	"vitess.io/vitess/go/vt/vttablet/tmclient"
+)
+
+type TabletManager struct {
+	tmclient.TabletManagerClient
+}
+
+func NewTabletManager(tmc tmclient.TabletManagerClient) *TabletManager {
+	return &TabletManager{tmc}
+}
+
+// tabletUndoDemotePrimary calls the said RPC for the given tablet.
+func (tm *TabletManager) tabletUndoDemotePrimary(ctx context.Context, tablet *topodatapb.Tablet, semiSync bool) error {
+	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
+	defer tmcCancel()
+	return tm.UndoDemotePrimary(tmcCtx, tablet, semiSync)
+}
+
+// setReadOnly calls the said RPC for the given tablet
+func (tm *TabletManager) setReadOnly(ctx context.Context, tablet *topodatapb.Tablet) error {
+	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
+	defer tmcCancel()
+	return tm.SetReadOnly(tmcCtx, tablet)
+}
+
+// changeTabletType calls the said RPC for the given tablet with the given parameters.
+func (tm *TabletManager) changeTabletType(ctx context.Context, tablet *topodatapb.Tablet, tabletType topodatapb.TabletType, semiSync bool) error {
+	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
+	defer tmcCancel()
+	return tm.ChangeType(tmcCtx, tablet, tabletType, semiSync)
+}
+
+// resetReplicationParameters resets the replication parameters on the given tablet.
+func (tm *TabletManager) resetReplicationParameters(ctx context.Context, tablet *topodatapb.Tablet) error {
+	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
+	defer tmcCancel()
+	return tm.ResetReplicationParameters(tmcCtx, tablet)
+}
+
+// setReplicationSource calls the said RPC with the parameters provided
+func (tm *TabletManager) setReplicationSource(ctx context.Context, replica *topodatapb.Tablet, primary *topodatapb.Tablet, semiSync bool, heartbeatInterval float64) error {
+	tmcCtx, tmcCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout)
+	defer tmcCancel()
+	return tm.SetReplicationSource(tmcCtx, replica, primary.Alias, 0, "", true, semiSync, heartbeatInterval)
+}