add download artifacts endpoint to enclave manager server

kurtosis-tech · Nov 9, 2023 · b6cbaae · b6cbaae
1 parent b16894c
commit b6cbaae
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 17 deletions.
diff --git a/enclave-manager/server/server.go b/enclave-manager/server/server.go
@@ -210,11 +210,11 @@ func (c *WebServer) RunStarlarkPackage(ctx context.Context, req *connect.Request
 	if err != nil {
 		return stacktrace.Propagate(err, "Failed to create the APIC client")
 	}
-	serviceRequest := &connect.Request[kurtosis_core_rpc_api_bindings.RunStarlarkPackageArgs]{
+	runStarlarkRequest := &connect.Request[kurtosis_core_rpc_api_bindings.RunStarlarkPackageArgs]{
 		Msg: req.Msg.RunStarlarkPackageArgs,
 	}
 
-	starlarkLogsStream, err := (*apiContainerServiceClient).RunStarlarkPackage(ctx, serviceRequest)
+	starlarkLogsStream, err := (*apiContainerServiceClient).RunStarlarkPackage(ctx, runStarlarkRequest)
 
 	for starlarkLogsStream.Receive() {
 		resp := starlarkLogsStream.Msg()
@@ -296,34 +296,42 @@ func (c *WebServer) InspectFilesArtifactContents(ctx context.Context, req *conne
 	return resp, nil
 }
 
-func (c *WebServer) DownloadFilesArtifact(ctx context.Context, req *connect.Request[kurtosis_enclave_manager_api_bindings.DownloadFilesArtifactArgs]) (*connect.Response[kurtosis_core_rpc_api_bindings.InspectFilesArtifactContentsResponse], error) {
+func (c *WebServer) DownloadFilesArtifact(
+	ctx context.Context,
+	req *connect.Request[kurtosis_enclave_manager_api_bindings.DownloadFilesArtifactRequest],
+	str *connect.ServerStream[kurtosis_core_rpc_api_bindings.StreamedDataChunk],
+) error {
 	auth, err := c.ValidateRequestAuthorization(ctx, c.enforceAuth, req.Header())
 	if err != nil {
-		return nil, stacktrace.Propagate(err, "Authentication attempt failed")
+		return stacktrace.Propagate(err, "Authentication attempt failed")
 	}
 	if !auth {
-		return nil, stacktrace.Propagate(err, "User not authorized")
+		return stacktrace.Propagate(err, "User not authorized")
 	}
 	apiContainerServiceClient, err := c.createAPICClient(req.Msg.ApicIpAddress, req.Msg.ApicPort)
 	if err != nil {
-		return nil, stacktrace.Propagate(err, "Failed to create the APIC client")
+		return stacktrace.Propagate(err, "Failed to create the APIC client")
 	}
 
-	request := &connect.Request[kurtosis_core_rpc_api_bindings.InspectFilesArtifactContentsRequest]{
-		Msg: &kurtosis_core_rpc_api_bindings.InspectFilesArtifactContentsRequest{
-			FileNamesAndUuid: req.Msg.FileNamesAndUuid,
+	filesArtifactIdentifier := req.Msg.FileUuid.Identifier
+	downloadFilesArtifactRequest := &connect.Request[kurtosis_core_rpc_api_bindings.DownloadFilesArtifactArgs]{
+		Msg: &kurtosis_core_rpc_api_bindings.DownloadFilesArtifactArgs{
+			Identifier: filesArtifactIdentifier,
 		},
 	}
-	result, err := (*apiContainerServiceClient).InspectFilesArtifactContents(ctx, request)
-	if err != nil {
-		return nil, err
+
+	filesArtifactStream, err := (*apiContainerServiceClient).DownloadFilesArtifact(ctx, downloadFilesArtifactRequest)
+	for filesArtifactStream.Receive() {
+		resp := filesArtifactStream.Msg()
+		err = str.Send(resp)
+		if err != nil {
+			return stacktrace.Propagate(err, "An error occurred in the enclave manager server attempting to send streamed data chunks for files artifact with identifier: %v.", filesArtifactIdentifier)
+		}
 	}
-	resp := &connect.Response[kurtosis_core_rpc_api_bindings.InspectFilesArtifactContentsResponse]{
-		Msg: &kurtosis_core_rpc_api_bindings.InspectFilesArtifactContentsResponse{
-			FileDescriptions: result.Msg.FileDescriptions,
-		},
+	if err = filesArtifactStream.Err(); err != nil {
+		return stacktrace.Propagate(err, "An error occurred in the enclave manager server attempting to receive streamed data chunks for files artifact with identifier %v.", filesArtifactIdentifier)
 	}
-	return resp, nil
+	return nil
 }
 
 func (c *WebServer) GetStarlarkRun(

diff --git a/go.work.sum b/go.work.sum
@@ -677,6 +677,7 @@ github.com/hashicorp/consul/api v1.3.0 h1:HXNYlRkkM/t+Y/Yhxtwcy02dlYwIaoxzvxPnS+
 github.com/hashicorp/consul/sdk v0.3.0 h1:UOxjlb4xVNF93jak1mzzoBatyFju9nrkxpVwIp/QqxQ=
 github.com/hashicorp/errwrap v1.0.0 h1:hLrqtEDnRye3+sgx6z4qVLNuviH3MR5aQ0ykNJa/UYA=
 github.com/hashicorp/go-cleanhttp v0.5.1 h1:dH3aiDG9Jvb5r5+bYHsikaOUIpcM0xvgMXVoDkXMzJM=
+github.com/hashicorp/go-immutable-radix v1.0.0 h1:AKDB1HM5PWEA7i4nhcpwOrO2byshxBjXVn/J/3+z5/0=
 github.com/hashicorp/go-immutable-radix v1.1.0 h1:vN9wG1D6KG6YHRTWr8512cxGOVgTMEfgEdSj/hr8MPc=
 github.com/hashicorp/go-memdb v1.0.4 h1:sIdJHAEtV3//iXcUb4LumSQeorYos5V0ptvqvQvFgDA=
 github.com/hashicorp/go-msgpack v0.5.3 h1:zKjpN5BK/P5lMYrLmBHdBULWbJ0XpYR+7NGzqkZzoD4=