Add GetPrimary to NodeManager

This commit implements the praefect.PrimaryGetter interface on NodeManager.
This allows for NodeManager to be used as a PrimaryGetter in code which has
been written to support repository specific primaries.

2 files changed, 18 insertions, 0 deletions

diff --git a/internal/praefect/nodes/manager.go b/internal/praefect/nodes/manager.go
index bc1b382ff..393159668 100644
--- a/internal/praefect/nodes/manager.go
+++ b/internal/praefect/nodes/manager.go
@@ -275,6 +275,17 @@ func (n *Mgr) GetShard(ctx context.Context, virtualStorageName string) (Shard, e
 	return strategy.GetShard(ctx)
 }
 
+// GetPrimary returns the current primary of a repository. This is an adapter so NodeManager can be used
+// as a praefect.PrimaryGetter in newer code which written to support repository specific primaries.
+func (n *Mgr) GetPrimary(ctx context.Context, virtualStorage, _ string) (string, error) {
+	shard, err := n.GetShard(ctx, virtualStorage)
+	if err != nil {
+		return "", err
+	}
+
+	return shard.Primary.GetStorage(), nil
+}
+
 func (n *Mgr) GetSyncedNode(ctx context.Context, virtualStorageName, repoPath string) (Node, error) {
 	upToDateStorages, err := n.csg.GetConsistentStorages(ctx, virtualStorageName, repoPath)
 	if err != nil {
diff --git a/internal/praefect/nodes/manager_test.go b/internal/praefect/nodes/manager_test.go
index 8661adfe6..e715078ef 100644
--- a/internal/praefect/nodes/manager_test.go
+++ b/internal/praefect/nodes/manager_test.go
@@ -294,12 +294,19 @@ func TestNodeManager(t *testing.T) {
 		Secondaries: []nodeAssertion{{node2.Storage, node2.Address}},
 	}, shard)
 
+	primary, err := nm.GetPrimary(ctx, "virtual-storage-0", "")
+	require.NoError(t, err)
+	require.Equal(t, shard.Primary.GetStorage(), primary)
+
 	healthSrv0.SetServingStatus("", grpc_health_v1.HealthCheckResponse_UNKNOWN)
 	healthSrv1.SetServingStatus("", grpc_health_v1.HealthCheckResponse_UNKNOWN)
 	checkShards(unhealthyCheckCount)
 
 	_, err = nm.GetShard(ctx, "virtual-storage-0")
 	require.Error(t, err, "should return error since no nodes are healthy")
+
+	_, err = nm.GetPrimary(ctx, "virtual-storage-0", "")
+	require.Equal(t, ErrPrimaryNotHealthy, err)
 }
 
 func TestMgr_GetSyncedNode(t *testing.T) {