From a1fb0997aa4beea60fd3bd57338a73bad8a2ef96 Mon Sep 17 00:00:00 2001
From: Jazzylol <zhitao.zzt@gmail.com>
Date: Wed, 13 Nov 2024 22:12:10 +0800
Subject: [PATCH] feat: Support for customizing account refresh time (#143)

---
 cmd/server/wire/wire_gen.go |  2 +-
 data/config.json            |  5 +++--
 internal/server/task.go     | 15 ++++++++++++---
 pkg/config/config.go        |  1 +
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/cmd/server/wire/wire_gen.go b/cmd/server/wire/wire_gen.go
index bd193aa..4a30d06 100644
--- a/cmd/server/wire/wire_gen.go
+++ b/cmd/server/wire/wire_gen.go
@@ -46,7 +46,7 @@ func NewWire(viperViper *viper.Viper, logger *log.Logger) (*app.App, func(), err
 	healthCheckHandler := handler.NewHealthCheckHandler()
 	httpServer := server.NewHTTPServer(logger, viperViper, jwtJWT, userHandler, shareHandler, accountHandler, healthCheckHandler)
 	job := server.NewJob(logger)
-	task := server.NewTask(logger, accountService, shareService)
+	task := server.NewTask(viperViper, logger, accountService, shareService)
 	migrate := server.NewMigrate(db, logger)
 	conversationRepository := repository.NewConversationRepository(repositoryRepository)
 	conversationLoggerMiddleware := middleware.NewConversationLoggerMiddleware(logger, conversationRepository)
diff --git a/data/config.json b/data/config.json
index a5efb0e..1ba243e 100644
--- 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
index 205a3a5..850eac1 100644
--- 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
index 134658a..bb00758 100644
--- 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)