feat: add download savepoint route

This commit is contained in:
2025-01-17 18:34:06 +03:30
parent ef7b16af68
commit c977c8a15d
7 changed files with 159 additions and 71 deletions

View File

@@ -1,12 +1,23 @@
package controller
import (
"bufio"
"context"
"errors"
"flink-kube-operator/internal/crd"
"flink-kube-operator/internal/crd/v1alpha1"
"flink-kube-operator/internal/managed_job"
"flink-kube-operator/internal/rest/types"
"flink-kube-operator/pkg"
"fmt"
"io"
"os"
"os/exec"
"path/filepath"
"k8s.io/apimachinery/pkg/types"
"github.com/danielgtaylor/huma/v2"
"go.uber.org/zap"
k8sTypes "k8s.io/apimachinery/pkg/types"
)
type GetJobsReq struct {
@@ -40,7 +51,7 @@ type StopJobResp struct {
func StopJob(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
mgr := managed_job.GetManager()
job := mgr.GetJob(types.UID(req.JobUId))
job := mgr.GetJob(k8sTypes.UID(req.JobUId))
err := job.Stop()
if err != nil {
return nil, err
@@ -52,7 +63,7 @@ func StopJob(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
func StartJob(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
mgr := managed_job.GetManager()
job := mgr.GetJob(types.UID(req.JobUId))
job := mgr.GetJob(k8sTypes.UID(req.JobUId))
err := job.Run(true)
if err != nil {
return nil, err
@@ -64,7 +75,7 @@ func StartJob(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
func RemoveJobJar(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
mgr := managed_job.GetManager()
job := mgr.GetJob(types.UID(req.JobUId))
job := mgr.GetJob(k8sTypes.UID(req.JobUId))
job.RemoveJar()
return &StopJobResp{Body: StopJobRespBody{
Success: true,
@@ -73,9 +84,75 @@ func RemoveJobJar(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
func PauseJob(ctx context.Context, req *StopJobReq) (*StopJobResp, error) {
mgr := managed_job.GetManager()
job := mgr.GetJob(types.UID(req.JobUId))
job := mgr.GetJob(k8sTypes.UID(req.JobUId))
job.Pause()
return &StopJobResp{Body: StopJobRespBody{
Success: true,
}}, nil
}
func DownloadSavepoint(ctx context.Context, req *types.SavepointDownloadReq) (*huma.StreamResponse, error) {
folderPath := req.SavepointPath // Change this to your folder path
// Create a temporary zip file
zipFilePath, err := filepath.Abs("./savepoint.zip")
pkg.Logger.Debug("[controller] [savepoint]",
zap.String("zipFileName", zipFilePath),
zap.String("folderPath", folderPath),
)
// Run the zip command
cmd := exec.Command("zip", "-r", zipFilePath, folderPath)
// Capture any output or errors
output, err := cmd.CombinedOutput()
if err != nil {
return nil, fmt.Errorf("failed to run zip command: %v\nOutput: %s", err, string(output))
}
pkg.Logger.Debug("[controller] [savepoint]", zap.Any("zip command output", string(output)))
// Open the zip file for reading
zipFileReader, err := os.Open(zipFilePath)
if err != nil {
return nil, fmt.Errorf("failed to open zip file: %w", err)
}
//defer zipFileReader.Close()
fileInfo, err := zipFileReader.Stat()
if err != nil {
return nil, fmt.Errorf("failed to get info of zipped file: %w", err)
}
resp := &huma.StreamResponse{
Body: func(ctx huma.Context) {
ctx.SetHeader("Content-Type", "application/zip")
ctx.SetHeader("Content-Length", fmt.Sprintf("%d", fileInfo.Size()))
ctx.SetHeader("Content-Disposition", fmt.Sprintf("attachment; filename=%s", zipFilePath))
writer := ctx.BodyWriter()
br := bufio.NewReader(zipFileReader)
for {
b, err := br.ReadByte()
if err != nil && !errors.Is(err, io.EOF) {
fmt.Println(err)
break
}
// process the one byte b
if err != nil {
// end of file
break
}
writer.Write([]byte{b})
}
os.Remove(zipFilePath)
},
}
return resp, nil
}