Dragonfly2

1
/*
2
 *     Copyright 2020 The Dragonfly Authors
3
 *
4
 * Licensed under the Apache License, Version 2.0 (the "License");
5
 * you may not use this file except in compliance with the License.
6
 * You may obtain a copy of the License at
7
 *
8
 *      http://www.apache.org/licenses/LICENSE-2.0
9
 *
10
 * Unless required by applicable law or agreed to in writing, software
11
 * distributed under the License is distributed on an "AS IS" BASIS,
12
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
 * See the License for the specific language governing permissions and
14
 * limitations under the License.
15
 */
16

17
package service
18

19
import (
20
	"context"
21
	"strconv"
22

23
	machineryv1tasks "github.com/RichardKnop/machinery/v1/tasks"
24
	"google.golang.org/grpc/codes"
25
	"google.golang.org/grpc/status"
26

27
	logger "d7y.io/dragonfly/v2/internal/dflog"
28
	internaljob "d7y.io/dragonfly/v2/internal/job"
29
	"d7y.io/dragonfly/v2/manager/models"
30
	"d7y.io/dragonfly/v2/manager/types"
31
)
32

33
const (
34
	// V1PreheatingStatePending is the preheating is waiting for starting
35
	V1PreheatingStatePending = "WAITING"
36

37
	// V1PreheatingStateRunning is the preheating is running
38
	V1PreheatingStateRunning = "RUNNING"
39

40
	// V1PreheatingStateSuccess is the preheating is success
41
	V1PreheatingStateSuccess = "SUCCESS"
42

43
	// V1PreheatingStateFail is the preheating is failed
44
	V1PreheatingStateFail = "FAIL"
45
)
46

47
func (s *service) CreateV1Preheat(ctx context.Context, json types.CreateV1PreheatRequest) (*types.CreateV1PreheatResponse, error) {
48
	job, err := s.CreatePreheatJob(ctx, types.CreatePreheatJobRequest{
49
		Type: internaljob.PreheatJob,
50
		Args: types.PreheatArgs{
51
			Type:                json.Type,
52
			URL:                 json.URL,
53
			FilteredQueryParams: json.FilteredQueryParams,
54
			Headers:             json.Headers,
55
		},
56
	})
57
	if err != nil {
58
		return nil, err
59
	}
60

61
	return &types.CreateV1PreheatResponse{
62
		ID: strconv.FormatUint(uint64(job.ID), 10),
63
	}, nil
64
}
65

66
func (s *service) GetV1Preheat(ctx context.Context, rawID string) (*types.GetV1PreheatResponse, error) {
67
	id, err := strconv.ParseUint(rawID, 10, 32)
68
	if err != nil {
69
		logger.Errorf("preheat convert error", err)
70
		return nil, status.Error(codes.Unknown, err.Error())
71
	}
72

73
	job := models.Job{}
74
	if err := s.db.WithContext(ctx).First(&job, uint(id)).Error; err != nil {
75
		return nil, status.Error(codes.Unknown, err.Error())
76
	}
77

78
	return &types.GetV1PreheatResponse{
79
		ID:         strconv.FormatUint(uint64(job.ID), 10),
80
		Status:     convertState(job.State),
81
		StartTime:  job.CreatedAt.String(),
82
		FinishTime: job.UpdatedAt.String(),
83
	}, nil
84
}
85

86
func convertState(state string) string {
87
	switch state {
88
	case machineryv1tasks.StatePending, machineryv1tasks.StateReceived, machineryv1tasks.StateRetry:
89
		return V1PreheatingStatePending
90
	case machineryv1tasks.StateStarted:
91
		return V1PreheatingStateRunning
92
	case machineryv1tasks.StateSuccess:
93
		return V1PreheatingStateSuccess
94
	case machineryv1tasks.StateFailure:
95
		return V1PreheatingStateFail
96
	}
97

98
	return V1PreheatingStateFail
99
}
100