feat: add backup and restore
This commit is contained in:
engigu
@@ -0,0 +1,304 @@
|
||||
package services
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
"baihu/internal/constant"
|
||||
"baihu/internal/database"
|
||||
"baihu/internal/models"
|
||||
)
|
||||
|
||||
type BackupService struct {
|
||||
settingsService *SettingsService
|
||||
}
|
||||
|
||||
func NewBackupService() *BackupService {
|
||||
return &BackupService{
|
||||
settingsService: NewSettingsService(),
|
||||
}
|
||||
}
|
||||
|
||||
const (
|
||||
BackupSection = "backup"
|
||||
BackupFileKey = "backup_file"
|
||||
BackupDir = "./data/backups"
|
||||
)
|
||||
|
||||
// tableConfig 表备份配置
|
||||
type tableConfig struct {
|
||||
filename string
|
||||
export func() (any, error)
|
||||
restore func([]byte) error
|
||||
}
|
||||
|
||||
func (s *BackupService) getTableConfigs() []tableConfig {
|
||||
return []tableConfig{
|
||||
{"tasks.json", s.exportTable(&[]models.Task{}, true), s.restoreTable(&[]models.Task{}, true)},
|
||||
{"task_logs.json", s.exportTable(&[]models.TaskLog{}, false), s.restoreTable(&[]models.TaskLog{}, false)},
|
||||
{"envs.json", s.exportTable(&[]models.EnvironmentVariable{}, true), s.restoreTable(&[]models.EnvironmentVariable{}, true)},
|
||||
{"scripts.json", s.exportTable(&[]models.Script{}, true), s.restoreTable(&[]models.Script{}, true)},
|
||||
{"settings.json", s.exportSettings, s.restoreSettings},
|
||||
{"send_stats.json", s.exportTable(&[]models.SendStats{}, false), s.restoreTable(&[]models.SendStats{}, false)},
|
||||
{"login_logs.json", s.exportTable(&[]models.LoginLog{}, false), s.restoreTable(&[]models.LoginLog{}, false)},
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) exportTable(dest any, unscoped bool) func() (any, error) {
|
||||
return func() (any, error) {
|
||||
db := database.DB
|
||||
if unscoped {
|
||||
db = db.Unscoped()
|
||||
}
|
||||
db.Find(dest)
|
||||
return dest, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) restoreTable(dest any, unscoped bool) func([]byte) error {
|
||||
return func(data []byte) error {
|
||||
if err := json.Unmarshal(data, dest); err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) exportSettings() (any, error) {
|
||||
var data []models.Setting
|
||||
database.DB.Where("section != ?", BackupSection).Find(&data)
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func (s *BackupService) restoreSettings(data []byte) error {
|
||||
var settings []models.Setting
|
||||
return json.Unmarshal(data, &settings)
|
||||
}
|
||||
|
||||
// CreateBackup 创建备份
|
||||
func (s *BackupService) CreateBackup() (string, error) {
|
||||
if err := os.MkdirAll(BackupDir, 0755); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
timestamp := time.Now().Format("20060102_150405")
|
||||
zipPath := filepath.Join(BackupDir, fmt.Sprintf("backup_%s.zip", timestamp))
|
||||
|
||||
zipFile, err := os.Create(zipPath)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer zipFile.Close()
|
||||
|
||||
zipWriter := zip.NewWriter(zipFile)
|
||||
defer zipWriter.Close()
|
||||
|
||||
// 导出各表
|
||||
for _, cfg := range s.getTableConfigs() {
|
||||
data, err := cfg.export()
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
jsonData, err := json.MarshalIndent(data, "", " ")
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
w, err := zipWriter.Create(cfg.filename)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, err := w.Write(jsonData); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
// 打包 scripts 文件夹
|
||||
scriptsDir := constant.ScriptsWorkDir
|
||||
if _, err := os.Stat(scriptsDir); err == nil {
|
||||
if err := s.addDirToZip(zipWriter, scriptsDir, "scripts"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
|
||||
s.settingsService.Set(BackupSection, BackupFileKey, zipPath)
|
||||
return zipPath, nil
|
||||
}
|
||||
|
||||
// Restore 恢复备份
|
||||
func (s *BackupService) Restore(zipPath string) error {
|
||||
r, err := zip.OpenReader(zipPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer r.Close()
|
||||
|
||||
// 构建文件名到配置的映射
|
||||
configs := s.getTableConfigs()
|
||||
fileMap := make(map[string]*zip.File)
|
||||
for _, f := range r.File {
|
||||
fileMap[f.Name] = f
|
||||
}
|
||||
|
||||
// 读取所有表数据
|
||||
tableData := make(map[string][]byte)
|
||||
for _, cfg := range configs {
|
||||
if f, ok := fileMap[cfg.filename]; ok {
|
||||
data, err := s.readZipFile(f)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
tableData[cfg.filename] = data
|
||||
}
|
||||
}
|
||||
|
||||
// 清空现有数据（物理删除）
|
||||
database.DB.Unscoped().Where("1=1").Delete(&models.Task{})
|
||||
database.DB.Unscoped().Where("1=1").Delete(&models.TaskLog{})
|
||||
database.DB.Unscoped().Where("1=1").Delete(&models.EnvironmentVariable{})
|
||||
database.DB.Unscoped().Where("1=1").Delete(&models.Script{})
|
||||
database.DB.Unscoped().Where("section != ?", BackupSection).Delete(&models.Setting{})
|
||||
database.DB.Unscoped().Where("1=1").Delete(&models.SendStats{})
|
||||
database.DB.Unscoped().Where("1=1").Delete(&models.LoginLog{})
|
||||
|
||||
// 恢复数据
|
||||
s.restoreFromData(tableData, "tasks.json", &[]models.Task{})
|
||||
s.restoreFromData(tableData, "task_logs.json", &[]models.TaskLog{})
|
||||
s.restoreFromData(tableData, "envs.json", &[]models.EnvironmentVariable{})
|
||||
s.restoreFromData(tableData, "scripts.json", &[]models.Script{})
|
||||
s.restoreFromData(tableData, "settings.json", &[]models.Setting{})
|
||||
s.restoreFromData(tableData, "send_stats.json", &[]models.SendStats{})
|
||||
s.restoreFromData(tableData, "login_logs.json", &[]models.LoginLog{})
|
||||
|
||||
// 恢复 scripts 文件夹
|
||||
s.restoreScriptsDir(r)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *BackupService) restoreFromData(tableData map[string][]byte, filename string, dest any) {
|
||||
if data, ok := tableData[filename]; ok {
|
||||
if err := json.Unmarshal(data, dest); err == nil {
|
||||
s.insertRecords(dest)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) insertRecords(records any) {
|
||||
switch v := records.(type) {
|
||||
case *[]models.Task:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
case *[]models.TaskLog:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
case *[]models.EnvironmentVariable:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
case *[]models.Script:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
case *[]models.Setting:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
case *[]models.SendStats:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
case *[]models.LoginLog:
|
||||
for _, r := range *v {
|
||||
database.DB.Create(&r)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) restoreScriptsDir(r *zip.ReadCloser) {
|
||||
scriptsDir := constant.ScriptsWorkDir
|
||||
for _, f := range r.File {
|
||||
if len(f.Name) > 8 && f.Name[:8] == "scripts/" {
|
||||
relPath := f.Name[8:]
|
||||
if relPath == "" {
|
||||
continue
|
||||
}
|
||||
fpath := filepath.Join(scriptsDir, relPath)
|
||||
if f.FileInfo().IsDir() {
|
||||
os.MkdirAll(fpath, 0755)
|
||||
continue
|
||||
}
|
||||
os.MkdirAll(filepath.Dir(fpath), 0755)
|
||||
if outFile, err := os.Create(fpath); err == nil {
|
||||
if rc, err := f.Open(); err == nil {
|
||||
io.Copy(outFile, rc)
|
||||
rc.Close()
|
||||
}
|
||||
outFile.Close()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BackupService) readZipFile(f *zip.File) ([]byte, error) {
|
||||
rc, err := f.Open()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rc.Close()
|
||||
return io.ReadAll(rc)
|
||||
}
|
||||
|
||||
func (s *BackupService) addDirToZip(zipWriter *zip.Writer, srcDir, prefix string) error {
|
||||
return filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
relPath, err := filepath.Rel(srcDir, path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
zipPath := filepath.ToSlash(filepath.Join(prefix, relPath))
|
||||
if info.IsDir() {
|
||||
if relPath != "." {
|
||||
_, err := zipWriter.Create(zipPath + "/")
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
w, err := zipWriter.Create(zipPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer file.Close()
|
||||
_, err = io.Copy(w, file)
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
func (s *BackupService) GetBackupFile() string {
|
||||
var setting models.Setting
|
||||
if err := database.DB.Where("section = ? AND `key` = ?", BackupSection, BackupFileKey).First(&setting).Error; err != nil {
|
||||
return ""
|
||||
}
|
||||
return setting.Value
|
||||
}
|
||||
|
||||
func (s *BackupService) ClearBackup() error {
|
||||
filePath := s.GetBackupFile()
|
||||
if filePath != "" {
|
||||
os.Remove(filePath)
|
||||
database.DB.Where("section = ? AND `key` = ?", BackupSection, BackupFileKey).Delete(&models.Setting{})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user