feat: Support for customizing account refresh time (#143)

nianhua99 · Nov 13, 2024 · a1fb099 · a1fb099
1 parent f102494
commit a1fb099
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 6 deletions.
diff --git a/cmd/server/wire/wire_gen.go b/cmd/server/wire/wire_gen.go
diff --git a/data/config.json b/data/config.json
@@ -46,11 +46,12 @@
       "claude": "http://localhost:9002",
       "index": "http://localhost:9001",
       "token": "https://token.oaifree.com"
-    }
+    },
+    "account_refresh_cron": ""
   },
   "security": {
     "2fa_secret": "",
-    "admin_password": ""
+    "admin_password": "12345678"
   },
   "share": {
     "custom": true,

diff --git a/internal/server/task.go b/internal/server/task.go
@@ -5,19 +5,22 @@ import (
 	"PandoraHelper/pkg/log"
 	"context"
 	"github.com/go-co-op/gocron"
+	"github.com/spf13/viper"
 	"go.uber.org/zap"
 	"time"
 )
 
 type Task struct {
+	conf           *viper.Viper
 	log            *log.Logger
 	scheduler      *gocron.Scheduler
 	accountService service.AccountService
 	shareService   service.ShareService
 }
 
-func NewTask(log *log.Logger, accountService service.AccountService, shareService service.ShareService) *Task {
+func NewTask(conf *viper.Viper, log *log.Logger, accountService service.AccountService, shareService service.ShareService) *Task {
 	return &Task{
+		conf:           conf,
 		log:            log,
 		accountService: accountService,
 		shareService:   shareService,
@@ -68,11 +71,17 @@ func (t *Task) Start(ctx context.Context) error {
 
 	t.scheduler = gocron.NewScheduler(time.UTC)
 
-	_, err := t.scheduler.Every(1).Day().At("00:00").Do(t.RefreshAllAccountEveryday, ctx)
+	var refreshCron = t.conf.GetString("pandora.account_refresh_cron")
+	t.log.Info("refresh cron with second is:" + refreshCron)
+	var err error
+	if refreshCron != "" {
+		_, err = t.scheduler.CronWithSeconds(refreshCron).Do(t.RefreshAllAccountEveryday, ctx)
+	} else {
+		_, err = t.scheduler.Every(1).Day().At("00:00").Do(t.RefreshAllAccountEveryday, ctx)
+	}
 	if err != nil {
 		return err
 	}
-
 	_, err = t.scheduler.Every(1).Day().At("00:05").Do(t.RefreshShareLimitEveryday, ctx)
 	if err != nil {
 		return err

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -53,6 +53,7 @@ func setDefaults(conf *viper.Viper) {
 	conf.SetDefault("pandora.domain.token", "https://token.oaifree.com")
 	conf.SetDefault("pandora.domain.index", "https://new.oaifree.com")
 	conf.SetDefault("pandora.domain.claude", "https://demo.fuclaude.com")
+	conf.SetDefault("pandora.account_refresh_cron", "")
 
 	// Share settings
 	conf.SetDefault("share.random", true)