Enable additional repository to have its storage re-written

2 files changed, 51 insertions, 7 deletions

diff --git a/internal/praefect/protoregistry/protoregistry.go b/internal/praefect/protoregistry/protoregistry.go
index 931da1461..94fd464be 100644
--- a/internal/praefect/protoregistry/protoregistry.go
+++ b/internal/praefect/protoregistry/protoregistry.go
@@ -82,12 +82,29 @@ type MethodInfo struct {
 	Operation      OpType
 	Scope          Scope
 	targetRepo     []int
+	additionalRepo []int
 	requestName    string // protobuf message name for input type
 	requestFactory protoFactory
 }
 
 // TargetRepo returns the target repository for a protobuf message if it exists
 func (mi MethodInfo) TargetRepo(msg proto.Message) (*gitalypb.Repository, error) {
+	return mi.getRepo(msg, mi.targetRepo)
+}
+
+// AdditionalRepo returns the additional repository for a protobuf message that needs a storage rewritten
+// if it exists
+func (mi MethodInfo) AdditionalRepo(msg proto.Message) (*gitalypb.Repository, bool, error) {
+	if mi.additionalRepo == nil {
+		return nil, false, nil
+	}
+
+	repo, err := mi.getRepo(msg, mi.additionalRepo)
+
+	return repo, true, err
+}
+
+func (mi MethodInfo) getRepo(msg proto.Message, targetOid []int) (*gitalypb.Repository, error) {
 	if mi.requestName != proto.MessageName(msg) {
 		return nil, fmt.Errorf(
 			"proto message %s does not match expected RPC request message %s",
@@ -95,7 +112,7 @@ func (mi MethodInfo) TargetRepo(msg proto.Message) (*gitalypb.Repository, error)
 		)
 	}
 
-	return reflectFindRepoTarget(msg, mi.targetRepo)
+	return reflectFindRepoTarget(msg, targetOid)
 }
 
 // UnmarshalRequestProto will unmarshal the bytes into the method's request
@@ -235,6 +252,13 @@ func parseMethodInfo(methodDesc *descriptor.MethodDescriptorProto) (MethodInfo,
 			return MethodInfo{}, err
 		}
 		mi.targetRepo = targetRepo
+
+		if opMsg.GetAdditionalRepositoryField() != "" {
+			mi.additionalRepo, err = parseOID(opMsg.GetAdditionalRepositoryField())
+			if err != nil {
+				return MethodInfo{}, err
+			}
+		}
 	}
 
 	return mi, nil
diff --git a/internal/praefect/protoregistry/targetrepo_test.go b/internal/praefect/protoregistry/targetrepo_test.go
index 8d6629524..23708b53b 100644
--- a/internal/praefect/protoregistry/targetrepo_test.go
+++ b/internal/praefect/protoregistry/targetrepo_test.go
@@ -35,12 +35,13 @@ func TestProtoRegistryTargetRepo(t *testing.T) {
 	}
 
 	testcases := []struct {
-		desc       string
-		svc        string
-		method     string
-		pbMsg      proto.Message
-		expectRepo *gitalypb.Repository
-		expectErr  error
+		desc                 string
+		svc                  string
+		method               string
+		pbMsg                proto.Message
+		expectRepo           *gitalypb.Repository
+		expectAdditionalRepo *gitalypb.Repository
+		expectErr            error
 	}{
 		{
 			desc:   "valid request type single depth",
@@ -71,6 +72,18 @@ func TestProtoRegistryTargetRepo(t *testing.T) {
 			},
 			expectRepo: testRepos[1],
 		},
+		{
+			desc:   "target nested, includes additional repository",
+			svc:    "ObjectPoolService",
+			method: "FetchIntoObjectPool",
+			pbMsg: &gitalypb.FetchIntoObjectPoolRequest{
+				Origin:     testRepos[0],
+				ObjectPool: &gitalypb.ObjectPool{Repository: testRepos[1]},
+				Repack:     false,
+			},
+			expectRepo:           testRepos[1],
+			expectAdditionalRepo: testRepos[0],
+		},
 	}
 
 	for _, tc := range testcases {
@@ -87,6 +100,13 @@ func TestProtoRegistryTargetRepo(t *testing.T) {
 			if tc.expectRepo != actualTarget {
 				t.Fatal("pointers do not match")
 			}
+
+			if tc.expectAdditionalRepo != nil {
+				additionalRepo, ok, err := info.AdditionalRepo(tc.pbMsg)
+				require.True(t, ok)
+				require.NoError(t, err)
+				require.Equal(t, tc.expectAdditionalRepo, additionalRepo)
+			}
 		})
 	}
 }