diff --git a/cmd/static/cmd.go b/cmd/static/cmd.go index 95c0a5c9..7643a1c9 100644 --- a/cmd/static/cmd.go +++ b/cmd/static/cmd.go @@ -22,6 +22,7 @@ import ( "github.com/cloudwego/cwgo/pkg/api_list" "github.com/cloudwego/cwgo/pkg/client" "github.com/cloudwego/cwgo/pkg/consts" + "github.com/cloudwego/cwgo/pkg/cronjob" "github.com/cloudwego/cwgo/pkg/curd/doc" "github.com/cloudwego/cwgo/pkg/fallback" "github.com/cloudwego/cwgo/pkg/job" @@ -107,6 +108,17 @@ func Init() *cli.App { return job.Job(globalArgs.JobArgument) }, }, + { + Name: CronJobName, + Usage: CronJobUsage, + Flags: cronjobFlags(), + Action: func(c *cli.Context) error { + if err := globalArgs.CronJobArgument.ParseCli(c); err != nil { + return err + } + return cronjob.Cronjob(globalArgs.CronJobArgument) + }, + }, { Name: ApiListName, Usage: ApiUsage, @@ -214,6 +226,12 @@ Examples: Examples: cwgo job --job_name jobOne --job_name jobTwo --module my_job +` + CronJobName = "cronjob" + CronJobUsage = `generate cronjob code + +Examples: + cwgo cronjob --job_name jobOne --job_name jobTwo --module my_cronjob ` FallbackName = "fallback" FallbackUsage = "fallback to hz or kitex" diff --git a/cmd/static/cronjob_flags.go b/cmd/static/cronjob_flags.go new file mode 100644 index 00000000..966fa7ae --- /dev/null +++ b/cmd/static/cronjob_flags.go @@ -0,0 +1,31 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package static + +import ( + "github.com/cloudwego/cwgo/pkg/consts" + "github.com/urfave/cli/v2" +) + +func cronjobFlags() []cli.Flag { + return []cli.Flag{ + &cli.StringSliceFlag{Name: consts.JobName, Usage: "Specify the cronjob name."}, + &cli.StringFlag{Name: consts.Module, Aliases: []string{"mod"}, Usage: "Specify the Go module name to generate go.mod."}, + &cli.StringFlag{Name: consts.OutDir, Usage: "Specify output directory, default is current dir."}, + &cli.StringFlag{Name: consts.JobFile, Usage: "Specify cronjob output file, default is job.go.", Value: "job.go"}, + } +} diff --git a/config/argument.go b/config/argument.go index 68ef14da..1daee0d6 100644 --- a/config/argument.go +++ b/config/argument.go @@ -39,6 +39,7 @@ type Argument struct { *ModelArgument *DocArgument *JobArgument + *CronJobArgument *ApiArgument *FallbackArgument } @@ -50,6 +51,7 @@ func NewArgument() *Argument { ModelArgument: NewModelArgument(), DocArgument: NewDocArgument(), JobArgument: NewJobArgument(), + CronJobArgument: NewCronJobArgument(), ApiArgument: NewApiArgument(), FallbackArgument: NewFallbackArgument(), } diff --git a/config/cronjob.go b/config/cronjob.go new file mode 100644 index 00000000..306fcbf0 --- /dev/null +++ b/config/cronjob.go @@ -0,0 +1,42 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package config + +import ( + "github.com/cloudwego/cwgo/pkg/consts" + "github.com/urfave/cli/v2" +) + +type CronJobArgument struct { + GoMod string + PackagePrefix string + JobName []string + OutDir string + JobFile string +} + +func NewCronJobArgument() *CronJobArgument { + return &CronJobArgument{} +} + +func (cj *CronJobArgument) ParseCli(ctx *cli.Context) error { + cj.JobName = ctx.StringSlice(consts.JobName) + cj.GoMod = ctx.String(consts.Module) + cj.OutDir = ctx.String(consts.OutDir) + cj.JobFile = ctx.String(consts.JobFile) + return nil +} diff --git a/pkg/common/utils/string.go b/pkg/common/utils/string.go new file mode 100644 index 00000000..2fb9bf12 --- /dev/null +++ b/pkg/common/utils/string.go @@ -0,0 +1,30 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package utils + +import "unicode" + +func CapitalizeFirstLetter(s string) string { + if len(s) == 0 { + return s + } + runes := []rune(s) + if !unicode.IsUpper(runes[0]) { + runes[0] = unicode.ToUpper(runes[0]) + } + return string(runes) +} diff --git a/pkg/consts/const.go b/pkg/consts/const.go index f3ba8505..04250859 100644 --- a/pkg/consts/const.go +++ b/pkg/consts/const.go @@ -167,6 +167,7 @@ const ( const ( JobName = "job_name" + JobFile = "job_file" ) const ( diff --git a/pkg/cronjob/cronjob.go b/pkg/cronjob/cronjob.go new file mode 100644 index 00000000..0ac7bda2 --- /dev/null +++ b/pkg/cronjob/cronjob.go @@ -0,0 +1,560 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package cronjob + +import ( + "bytes" + "errors" + "fmt" + "go/ast" + "go/parser" + "go/printer" + "go/token" + "os" + "path/filepath" + "strings" + "text/template" + + "github.com/cloudwego/cwgo/meta" + + "github.com/cloudwego/cwgo/config" + "github.com/cloudwego/cwgo/pkg/common/utils" + "github.com/cloudwego/cwgo/pkg/consts" + "github.com/cloudwego/kitex/tool/internal_pkg/log" +) + +func Cronjob(c *config.CronJobArgument) error { + if err := check(c); err != nil { + return err + } + + err := generateCronjobFile(c) + if err != nil { + return err + } + + return nil +} + +func check(c *config.CronJobArgument) (err error) { + if len(c.JobName) == 0 { + return errors.New("job name is empty") + } + + c.OutDir, err = filepath.Abs(c.OutDir) + if err != nil { + return err + } + + gopath, err := utils.GetGOPATH() + if err != nil { + return fmt.Errorf("GET gopath failed: %s", err) + } + if gopath == "" { + return fmt.Errorf("GOPATH is not set") + } + + gosrc := filepath.Join(gopath, consts.Src) + gosrc, err = filepath.Abs(gosrc) + if err != nil { + log.Warn("Get GOPATH/src path failed:", err.Error()) + os.Exit(1) + } + curpath, err := filepath.Abs(consts.CurrentDir) + if err != nil { + log.Warn("Get current path failed:", err.Error()) + os.Exit(1) + } + + if !strings.HasSuffix(c.JobFile, ".go") { + log.Warn("--job_file must be a go file") + os.Exit(1) + } + + for k := range c.JobName { + c.JobName[k] = utils.CapitalizeFirstLetter(c.JobName[k]) + } + + if strings.HasPrefix(curpath, gosrc) { + goPkg := "" + if goPkg, err = filepath.Rel(gosrc, curpath); err != nil { + log.Warn("Get GOPATH/src relpath failed:", err.Error()) + os.Exit(1) + } + + if c.GoMod == "" { + if utils.IsWindows() { + c.GoMod = strings.ReplaceAll(goPkg, consts.BackSlash, consts.Slash) + } else { + c.GoMod = goPkg + } + } + + if c.GoMod != "" { + if utils.IsWindows() { + goPkgSlash := strings.ReplaceAll(goPkg, consts.BackSlash, consts.Slash) + if goPkgSlash != c.GoMod { + return fmt.Errorf("module name: %s is not the same with GoPkg under GoPath: %s", c.GoMod, goPkgSlash) + } + } else { + if c.GoMod != goPkg { + return fmt.Errorf("module name: %s is not the same with GoPkg under GoPath: %s", c.GoMod, goPkg) + } + } + } + } + + if !strings.HasPrefix(curpath, gosrc) && c.GoMod == "" { + log.Warn("Outside of $GOPATH. Please specify a module name with the '-module' flag.") + os.Exit(1) + } + + if c.GoMod != "" { + module, path, ok := utils.SearchGoMod(curpath, true) + + if ok { + // go.mod exists + if module != c.GoMod { + log.Warnf("The module name given by the '-module' option ('%s') is not consist with the name defined in go.mod ('%s' from %s)\n", + c.GoMod, module, path) + os.Exit(1) + } + if c.PackagePrefix, err = filepath.Rel(path, c.OutDir); err != nil { + log.Warn("Get package prefix failed:", err.Error()) + os.Exit(1) + } + c.PackagePrefix = filepath.Join(c.GoMod, c.PackagePrefix) + } else { + if err = utils.InitGoMod(c.GoMod); err != nil { + log.Warn("Init go mod failed:", err.Error()) + os.Exit(1) + } + if c.PackagePrefix, err = filepath.Rel(curpath, c.OutDir); err != nil { + log.Warn("Get package prefix failed:", err.Error()) + os.Exit(1) + } + c.PackagePrefix = filepath.Join(c.GoMod, c.PackagePrefix) + } + } + + c.PackagePrefix = strings.ReplaceAll(c.PackagePrefix, consts.BackSlash, consts.Slash) + + return nil +} + +type JobsData struct { + JobInfos []JobInfo +} +type JobInfo struct { + JobName string + GoModule string + PackagePrefix string +} + +func addScheduleNewJobs(data string, jobs []JobInfo) (string, error) { + fSet := token.NewFileSet() + file, err := parser.ParseFile(fSet, "", data, parser.ParseComments) + if err != nil { + return "", err + } + + // Extract existing cronjob calls + existingJobs := make(map[string]bool) + ast.Inspect(file, func(n ast.Node) bool { + switch x := n.(type) { + case *ast.CallExpr: + if sel, ok := x.Fun.(*ast.SelectorExpr); ok { + if ident, ok := sel.X.(*ast.Ident); ok && sel.Sel.Name == "AddFunc" && ident.Name == "c" { + for _, arg := range x.Args { + if funcLit, ok := arg.(*ast.FuncLit); ok { + ast.Inspect(funcLit.Body, func(n ast.Node) bool { + switch stmt := n.(type) { + case *ast.ExprStmt: + if callExpr, ok := stmt.X.(*ast.CallExpr); ok { + if sel, ok := callExpr.Fun.(*ast.SelectorExpr); ok { + if id, ok := sel.X.(*ast.Ident); ok && id.Name == "job" { + existingJobs[sel.Sel.Name] = true + } + } + } + } + return true + }) + } + } + } + } + } + return true + }) + + buf := new(bytes.Buffer) + + for _, decl := range file.Decls { + if fn, ok := decl.(*ast.FuncDecl); ok && fn.Name.Name == "Init" { + for k := range jobs { + if _, ok := existingJobs[jobs[k].JobName]; !ok { + jobName := jobs[k].JobName + addCronJobCall(fn, jobName) + } + } + } + } + + // Generate the modified code + if err := printer.Fprint(buf, fSet, file); err != nil { + return "", err + } + return buf.String(), nil +} + +func addNewJobs(data string, jobs []JobInfo) (string, error) { + fSet := token.NewFileSet() + file, err := parser.ParseFile(fSet, "", data, parser.ParseComments) + if err != nil { + return "", err + } + + // Extract existing cronjob + existingJobs := make(map[string]bool) + ast.Inspect(file, func(n ast.Node) bool { + if fn, ok := n.(*ast.FuncDecl); ok { + existingJobs[fn.Name.Name] = true + } + return true + }) + + for _, job := range jobs { + if _, ok := existingJobs[job.JobName]; !ok { + addCronJobFunction(file, job.JobName) + } + } + + // Generate the modified code + buf := new(bytes.Buffer) + if err = printer.Fprint(buf, fSet, file); err != nil { + return "", err + } + + return buf.String(), nil +} + +func addCronJobFunction(file *ast.File, jobName string) { + funcDecl := &ast.FuncDecl{ + Name: ast.NewIdent(jobName), + Type: &ast.FuncType{ + Params: &ast.FieldList{ + List: []*ast.Field{ + { + Names: []*ast.Ident{ast.NewIdent("ctx")}, + Type: &ast.SelectorExpr{ + X: ast.NewIdent("context"), + Sel: ast.NewIdent("Context"), + }, + }, + }, + }, + }, + Body: &ast.BlockStmt{ + List: []ast.Stmt{ + &ast.ExprStmt{ + X: &ast.BasicLit{ + Kind: token.STRING, + Value: `// TODO: fill with your own logic`, + }, + }, + }, + }, + Doc: &ast.CommentGroup{}, + } + + // Append the new function declaration to the AST file + file.Decls = append(file.Decls, funcDecl) +} + +func addCronJobCall(fn *ast.FuncDecl, jobName string) { + // Construct the AST nodes for the new cron job call + addFuncCall := &ast.AssignStmt{ + Lhs: []ast.Expr{ + ast.NewIdent("_, err"), + }, + Tok: token.ASSIGN, + Rhs: []ast.Expr{ + &ast.CallExpr{ + Fun: &ast.SelectorExpr{ + X: ast.NewIdent("c"), + Sel: ast.NewIdent("AddFunc"), + }, + Args: []ast.Expr{ + &ast.BasicLit{ + Kind: token.STRING, + Value: `"* * * * *"`, + }, + &ast.FuncLit{ + Type: &ast.FuncType{ + Params: &ast.FieldList{ + List: []*ast.Field{}, + }, + }, + Body: &ast.BlockStmt{ + List: []ast.Stmt{ + &ast.SelectStmt{ + Body: &ast.BlockStmt{ + List: []ast.Stmt{ + &ast.CommClause{ + Comm: &ast.ExprStmt{ + X: &ast.UnaryExpr{ + Op: token.ARROW, + X: ast.NewIdent("ctx.Done()"), + }, + }, + Body: []ast.Stmt{ + &ast.ExprStmt{ + X: &ast.CallExpr{ + Fun: &ast.SelectorExpr{ + X: ast.NewIdent("log"), + Sel: ast.NewIdent("Println"), + }, + Args: []ast.Expr{ + &ast.BasicLit{ + Kind: token.STRING, + Value: fmt.Sprintf(`"%s terminated."`, jobName), + }, + }, + }, + }, + &ast.ReturnStmt{}, + }, + }, + &ast.CommClause{ + Body: []ast.Stmt{ + &ast.ExprStmt{ + X: &ast.CallExpr{ + Fun: &ast.SelectorExpr{ + X: ast.NewIdent("job"), + Sel: ast.NewIdent(jobName), + }, + Args: []ast.Expr{ + ast.NewIdent("ctx"), + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + }, + } + + checkErrCall := &ast.IfStmt{ + Cond: &ast.BinaryExpr{ + X: ast.NewIdent("err"), + Op: token.NEQ, + Y: ast.NewIdent("nil"), + }, + Body: &ast.BlockStmt{ + List: []ast.Stmt{ + &ast.ExprStmt{ + X: &ast.CallExpr{ + Fun: &ast.SelectorExpr{ + X: ast.NewIdent("log"), + Sel: ast.NewIdent("Fatalf"), + }, + Args: []ast.Expr{ + &ast.BasicLit{ + Kind: token.STRING, + Value: `"Error adding new cron job: %v"`, + }, + ast.NewIdent("err"), + }, + }, + }, + }, + }, + } + + dummyCall := &ast.ExprStmt{ + X: &ast.BasicLit{ + Kind: token.STRING, + Value: "", + }, + } + + fn.Body.List = append(fn.Body.List, dummyCall, addFuncCall, checkErrCall) +} + +func generateCronjobFile(c *config.CronJobArgument) error { + // Ensure the base output directory exists + err := os.MkdirAll(c.OutDir, 0o755) + if err != nil { + return fmt.Errorf("failed to create output directory: %w", err) + } + + // Create cmd/main.go and overwrite each time + cmdDir := filepath.Join(c.OutDir, "cmd") + err = os.MkdirAll(cmdDir, 0o755) + if err != nil { + return fmt.Errorf("failed to create cmd directory: %w", err) + } + + mainGoPath := filepath.Join(cmdDir, "main.go") + tmpl, err := template.New("job_main").Parse(jobMainTemplate) + if err != nil { + return err + } + + jobsInfo := &JobsData{ + JobInfos: make([]JobInfo, 0), + } + for _, v := range c.JobName { + jobsInfo.JobInfos = append(jobsInfo.JobInfos, JobInfo{ + JobName: v, + GoModule: c.GoMod, + PackagePrefix: c.PackagePrefix, + }) + } + + var jobFileContent bytes.Buffer + data := struct { + PackagePrefix string + Version string + }{ + PackagePrefix: c.PackagePrefix, + Version: meta.Version, + } + err = tmpl.Execute(&jobFileContent, data) + if err != nil { + return err + } + err = utils.CreateFile(mainGoPath, jobFileContent.String()) + if err != nil { + return err + } + jobFileContent.Reset() + + // Create or append to schedule.go + internalDir := filepath.Join(c.OutDir, "internal") + err = os.MkdirAll(internalDir, 0o755) + if err != nil { + return fmt.Errorf("failed to create internal directory: %w", err) + } + scheduleGoPath := filepath.Join(internalDir, "schedule.go") + scheduleTmpl, err := template.New("job_schedule").Parse(jobScheduleTemplate) + if err != nil { + return err + } + + if exist, _ := utils.PathExist(scheduleGoPath); !exist { + err = scheduleTmpl.Execute(&jobFileContent, jobsInfo) + if err != nil { + return err + } + err = utils.CreateFile(scheduleGoPath, jobFileContent.String()) + if err != nil { + return fmt.Errorf("failed to write schedule.go: %w", err) + } + jobFileContent.Reset() + } else { + src, err := utils.ReadFileContent(scheduleGoPath) + if err != nil { + return fmt.Errorf("failed to read schedule.go: %w", err) + } + + res, err := addScheduleNewJobs(string(src), jobsInfo.JobInfos) + if err != nil { + return err + } + + err = utils.CreateFile(scheduleGoPath, res) + if err != nil { + return err + } + } + + // Create or append to job_file + jobDir := filepath.Join(internalDir, "job") + err = os.MkdirAll(jobDir, 0o755) + if err != nil { + return fmt.Errorf("failed to create job directory: %w", err) + } + jobGoPath := filepath.Join(jobDir, c.JobFile) + jobTmpl, err := template.New("job_file").Parse(jobTemplate) + if err != nil { + return err + } + + if exist, _ := utils.PathExist(jobGoPath); !exist { + err = jobTmpl.Execute(&jobFileContent, jobsInfo) + if err != nil { + return err + } + err = utils.CreateFile(jobGoPath, jobFileContent.String()) + if err != nil { + return fmt.Errorf("failed to write job_file: %w", err) + } + jobFileContent.Reset() + } else { + src, err := utils.ReadFileContent(jobGoPath) + if err != nil { + return fmt.Errorf("failed to read job_file: %w", err) + } + + res, err := addNewJobs(string(src), jobsInfo.JobInfos) + if err != nil { + return err + } + + err = utils.CreateFile(jobGoPath, res) + if err != nil { + return err + } + } + + // Create or append to run.sh + scriptsDir := filepath.Join(c.OutDir, "scripts") + err = os.MkdirAll(scriptsDir, 0o755) + if err != nil { + return fmt.Errorf("failed to create scripts directory: %w", err) + } + + // Create run.sh + runShPath := filepath.Join(scriptsDir, "run.sh") + scriptTmpl, err := template.New("job_script").Parse(scriptTemplate) + if err != nil { + return err + } + + if exist, _ := utils.PathExist(runShPath); !exist { + err = scriptTmpl.Execute(&jobFileContent, nil) + if err != nil { + return err + } + err = utils.CreateFile(runShPath, jobFileContent.String()) + if err != nil { + return err + } + jobFileContent.Reset() + } + + return nil +} diff --git a/pkg/cronjob/cronjob_test.go b/pkg/cronjob/cronjob_test.go new file mode 100644 index 00000000..911960ee --- /dev/null +++ b/pkg/cronjob/cronjob_test.go @@ -0,0 +1,368 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package cronjob + +import ( + "os" + "os/exec" + "path/filepath" + "testing" + + "github.com/cloudwego/cwgo/config" + "github.com/stretchr/testify/assert" +) + +func TestCronJobValidCreation(t *testing.T) { + args := &config.CronJobArgument{ + JobName: []string{"cronjob1"}, + GoMod: "github.com/cloudwego/cwgo", + PackagePrefix: "github.com/cloudwego/cwgo", + OutDir: "./test_out", + JobFile: "job.go", + } + + err := Cronjob(args) + assert.NoError(t, err) + + checkFiles := []string{ + "test_out/cmd/main.go", + "test_out/internal/schedule.go", + "test_out/scripts/run.sh", + "test_out/internal/job/job.go", + } + + for _, file := range checkFiles { + _, err := os.Stat(file) + assert.NoError(t, err) + } + + contentChecks := map[string]string{ + "test_out/internal/schedule.go": "job.Cronjob1(ctx)", + } + + for file, content := range contentChecks { + data, err := os.ReadFile(file) + assert.NoError(t, err) + assert.Contains(t, string(data), content) + } + + err = os.RemoveAll(args.OutDir) + assert.NoError(t, err) +} + +func TestAddScheduleNewJobs(t *testing.T) { + original := ` + +package schedule + +import ( + "context" + "log" + + "github.com/robfig/cron/v3" + + "test/internal/job" +) + +func Init(ctx context.Context, c *cron.Cron) { + var err error + + _, err = c.AddFunc("* * * * *", func() { + select { + case <-ctx.Done(): + log.Println("JobOne terminated.") + return + default: + job.JobOne(ctx) + } + }) + if err != nil { + log.Fatalf("Error adding cron job: %v", err) + } + + _, err = c.AddFunc("* * * * *", func() { + select { + case <-ctx.Done(): + log.Println("JobTwo terminated.") + return + default: + job.JobTwo(ctx) + } + }) + if err != nil { + log.Fatalf("Error adding new cron job: %v", err) + } + +} + + +` + + jobs := []JobInfo{ + {JobName: "JobOne", PackagePrefix: "github.com/cloudwego/cwgo", GoModule: "github.com/cloudwego/cwgo"}, + {JobName: "JobThree", PackagePrefix: "github.com/cloudwego/cwgo", GoModule: "github.com/cloudwego/cwgo"}, + } + + expected := `package schedule + +import ( + "context" + "log" + + "github.com/robfig/cron/v3" + + "test/internal/job" +) + +func Init(ctx context.Context, c *cron.Cron) { + var err error + + _, err = c.AddFunc("* * * * *", func() { + select { + case <-ctx.Done(): + log.Println("JobOne terminated.") + return + default: + job.JobOne(ctx) + } + }) + if err != nil { + log.Fatalf("Error adding cron job: %v", err) + } + + _, err = c.AddFunc("* * * * *", func() { + select { + case <-ctx.Done(): + log.Println("JobTwo terminated.") + return + default: + job.JobTwo(ctx) + } + }) + if err != nil { + log.Fatalf("Error adding new cron job: %v", err) + } + + _, err = c.AddFunc("* * * * *", func() { + select { + case <-ctx.Done(): + log.Println("JobThree terminated.") + return + default: + job.JobThree(ctx) + } + }) + if err != nil { + log.Fatalf("Error adding new cron job: %v", err) + } + +} +` + + result, err := addScheduleNewJobs(original, jobs) + assert.NoError(t, err) + assert.Equal(t, expected, result) +} + +func TestAddNewJobs(t *testing.T) { + original := ` + +package job + +import ( + "context" +) + +func JobOne(ctx context.Context) { + // TODO: fill with your own logic +} + +func JobTwo(ctx context.Context) { + // TODO: fill with your own logic +} + +` + + jobs := []JobInfo{ + {JobName: "JobOne", PackagePrefix: "github.com/cloudwego/cwgo", GoModule: "github.com/cloudwego/cwgo"}, + {JobName: "JobThree", PackagePrefix: "github.com/cloudwego/cwgo", GoModule: "github.com/cloudwego/cwgo"}, + } + + expected := `package job + +import ( + "context" +) + +func JobOne(ctx context.Context) { + // TODO: fill with your own logic +} + +func JobTwo(ctx context.Context) { + // TODO: fill with your own logic +} + +func JobThree(ctx context.Context) { + // TODO: fill with your own logic +} +` + + result, err := addNewJobs(original, jobs) + assert.NoError(t, err) + assert.Equal(t, expected, result) +} + +//func TestCreateCronRunCall(t *testing.T) { +// jobName := "cronjob1" +// expected := `wg.Add(1) +//go func() { +// defer wg.Done() +// cronjob1.Run() +//}()` +// +// result, err := addNewJobs(jobName) +// assert.NoError(t, err) +// var buf bytes.Buffer +// err := printer.Fprint(&buf, token.NewFileSet(), result) +// assert.NoError(t, err) +// assert.Equal(t, expected, buf.String()) +//} + +func TestCronJobMissingJobName(t *testing.T) { + args := &config.CronJobArgument{ + JobName: []string{""}, + GoMod: "github.com/cloudwego/cwgo", + PackagePrefix: "github.com/cloudwego/cwgo", + OutDir: "./test_out", + JobFile: "job.go", + } + + if os.Getenv("BE_CRASHER") == "1" { + err := Cronjob(args) + assert.Error(t, err) + return + } + cmd := exec.Command(os.Args[0], "-test.run=TestCronJobMissingJobName") + cmd.Env = append(os.Environ(), "BE_CRASHER=1") + err := cmd.Run() + defer func() { + _ = os.RemoveAll(args.OutDir) + }() + if e, ok := err.(*exec.ExitError); ok && !e.Success() { + return + } + t.Logf("process ran with err %v, want exit status 1", err) +} + +func TestCronJobWithDifferentModule(t *testing.T) { + args := &config.CronJobArgument{ + JobName: []string{"cronjob2"}, + GoMod: "github.com/cloudwego/another_test", + PackagePrefix: "github.com/cloudwego/another_test", + OutDir: "./another_test_out", + JobFile: "job.go", + } + + if os.Getenv("BE_CRASHER") == "1" { + err := Cronjob(args) + assert.Error(t, err) + return + } + cmd := exec.Command(os.Args[0], "-test.run=TestCronJobWithDifferentModule") + cmd.Env = append(os.Environ(), "BE_CRASHER=1") + err := cmd.Run() + if e, ok := err.(*exec.ExitError); ok && !e.Success() { + return + } + t.Logf("process ran with err %v, want exit status 1", err) +} + +func TestCronJobWithEmptyModule(t *testing.T) { + args := &config.CronJobArgument{ + JobName: []string{"cronjob3"}, + GoMod: "", + PackagePrefix: "github.com/cloudwego/cwgo", + OutDir: "./test_out", + JobFile: "job.go", + } + + if os.Getenv("BE_CRASHER") == "1" { + err := Cronjob(args) + assert.Error(t, err) + return + } + cmd := exec.Command(os.Args[0], "-test.run=TestCronJobWithEmptyModule") + cmd.Env = append(os.Environ(), "BE_CRASHER=1") + err := cmd.Run() + if e, ok := err.(*exec.ExitError); ok && !e.Success() { + return + } + t.Logf("process ran with err %v, want exit status 1", err) +} + +func TestCronJobWithoutGoModInGOPATH(t *testing.T) { + args := &config.CronJobArgument{ + JobName: []string{"cronjob3"}, + GoMod: "", + PackagePrefix: "github.com/cloudwego/cwgo", + OutDir: "./test_out", + JobFile: "job.go", + } + + // Setup a temporary GOPATH + tmpGopath, err := os.MkdirTemp("", "gopath") + if err != nil { + t.Fatalf("Failed to create temp GOPATH: %v", err) + } + defer os.RemoveAll(tmpGopath) + + // Create a directory structure under GOPATH/src + projectPath := filepath.Join(tmpGopath, "src", "github.com", "cloudwego", "cwgo") + if err := os.MkdirAll(projectPath, 0o755); err != nil { + t.Fatalf("Failed to create project directory: %v", err) + } + + // Move test file to the project directory + currentTestFile := filepath.Join(projectPath, "cronjob_test.go") + if err := os.Rename(os.Args[0], currentTestFile); err != nil { + t.Fatalf("Failed to move test file to project directory: %v", err) + } + + // Set the GOPATH environment variable + oldGopath := os.Getenv("GOPATH") + os.Setenv("GOPATH", tmpGopath) + defer os.Setenv("GOPATH", oldGopath) + + if os.Getenv("BE_CRASHER") == "1" { + err := Cronjob(args) + assert.Error(t, err) + return + } + + // Run the test in a subprocess + cmd := exec.Command(os.Args[0], "-test.run=TestCronJobWithoutGoModInGOPATH") + cmd.Env = append(os.Environ(), "BE_CRASHER=1") + output, err := cmd.CombinedOutput() + + if err == nil { + t.Fatalf("Expected process to exit with an error, got nil. Output: %s", output) + } + + if exitErr, ok := err.(*exec.ExitError); ok && exitErr.ExitCode() != 0 { + t.Logf("Expected process to exit with status 0, got %d. Output: %s", exitErr.ExitCode(), output) + } +} diff --git a/pkg/cronjob/template.go b/pkg/cronjob/template.go new file mode 100644 index 00000000..d00e0aa2 --- /dev/null +++ b/pkg/cronjob/template.go @@ -0,0 +1,128 @@ +/* + * Copyright 2024 CloudWeGo Authors + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package cronjob + +const jobTemplate = `package job + +import( + "context" +) + +{{- range .JobInfos }} + +func {{.JobName}}(ctx context.Context) { + // TODO: fill with your own logic +} +{{- end }} +` + +const jobMainTemplate = `// Code generated by cwgo ({{.Version}}). DO NOT EDIT. +package main + +import ( + "context" + "log" + "os" + "os/signal" + "syscall" + + "{{.PackagePrefix}}/internal" + + "time" + + "github.com/robfig/cron/v3" +) + +func main() { + ctx, cancel := context.WithCancel(context.Background()) + + c := cron.New() + + schedule.Init(ctx,c) + + c.Start() + + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + <-sigs + + cancel() + + time.Sleep(2 * time.Second) + + c.Stop() + log.Println("Cron job stopped gracefully.") +} + +` + +const jobScheduleTemplate = `package schedule + +import ( + "context" + "github.com/robfig/cron/v3" + "log" + "{{(index .JobInfos 0).PackagePrefix}}/internal/job" +) + +func Init(ctx context.Context, c *cron.Cron){ + var err error + {{range .JobInfos}} + _, err = c.AddFunc("* * * * *", func(){ + select { + case <-ctx.Done(): + log.Println("{{.JobName}} terminated.") + return + default: + job.{{.JobName}}(ctx) + } + }) + if err != nil { + log.Fatalf("Error adding cron job: %v", err) + } + {{end}} +} +` + +const scriptTemplate = `#!/bin/bash + +echo "Building cronjob binary..." +go build -o cronjob ../cmd/main.go + +if [ $? -ne 0 ]; then + echo "Error: Failed to build cronjob." + exit 1 +fi + +echo "Running cronjob..." +./cronjob + +if [ $? -ne 0 ]; then + echo "Error: cronjob execution failed." + exit 1 +fi + +echo "cronjob Done." + +rm cronjob + +if [ $? -ne 0 ]; then + echo "Error: Failed to remove cronjob binary." + exit 1 +fi + +`