feat: add download savepoint route
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user