1 files changed, 5 insertions, 13 deletions
diff --git a/go/internal/handler/upload_archive.go b/go/internal/handler/upload_archive.go
index 2175300..eb2fa90 100644
--- a/go/internal/handler/upload_archive.go
+++ b/go/internal/handler/upload_archive.go
@@ -2,25 +2,17 @@ package handler
 
 import (
 	"context"
-	"fmt"
 	"os"
 
 	pb "gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/client"
+	"google.golang.org/grpc"
 )
 
-func UploadArchive(gitalyAddress string, request *pb.SSHUploadArchiveRequest) (int32, error) {
-	if gitalyAddress == "" {
-		return 0, fmt.Errorf("no gitaly_address given")
-	}
-
-	conn, err := client.Dial(gitalyAddress, dialOpts())
-	if err != nil {
-		return 0, err
-	}
-	defer conn.Close()
-
-	ctx, cancel := context.WithCancel(context.Background())
+// UploadArchive issues a Gitaly upload-archive rpc to the provided address
+func UploadArchive(ctx context.Context, conn *grpc.ClientConn, request *pb.SSHUploadArchiveRequest) (int32, error) {
+	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
+
 	return client.UploadArchive(ctx, conn, os.Stdin, os.Stdout, os.Stderr, request)
 }