新增本地持久化存储

2025-11-22 17:48:04 +08:00
parent 62b5a3f9bc
commit ec3fc77403
6 changed files with 470 additions and 0 deletions
--- a/storage/store.go
+++ b/storage/store.go
@@ -0,0 +1,187 @@
+package storage
+
+import (
+	"encoding/json"
+	"fmt"
+	"log"
+	"os"
+	"path/filepath"
+	"time"
+
+	bolt "go.etcd.io/bbolt"
+)
+
+// Bucket 名称常量
+var (
+	BucketSites    = []byte("sites")
+	BucketDownload = []byte("download_records")
+	BucketLogs     = []byte("access_logs")
+	BucketSettings = []byte("settings")
+)
+
+// Store BBolt 存储服务
+type Store struct {
+	db *bolt.DB
+}
+
+// NewStore 创建新的存储服务
+func NewStore(dbPath string) (*Store, error) {
+	// 确保数据库目录存在
+	dir := filepath.Dir(dbPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return nil, fmt.Errorf("创建数据库目录失败: %v", err)
+	}
+
+	// 打开数据库
+	db, err := bolt.Open(dbPath, 0600, &bolt.Options{
+		Timeout: 1 * time.Second,
+	})
+	if err != nil {
+		return nil, fmt.Errorf("打开数据库失败: %v", err)
+	}
+
+	// 创建必要的 buckets
+	err = db.Update(func(tx *bolt.Tx) error {
+		buckets := [][]byte{BucketSites, BucketDownload, BucketLogs, BucketSettings}
+		for _, bucket := range buckets {
+			if _, err := tx.CreateBucketIfNotExists(bucket); err != nil {
+				return fmt.Errorf("创建 bucket %s 失败: %v", string(bucket), err)
+			}
+		}
+		return nil
+	})
+
+	if err != nil {
+		db.Close()
+		return nil, err
+	}
+
+	log.Printf("数据库初始化成功: %s", dbPath)
+
+	return &Store{db: db}, nil
+}
+
+// Close 关闭数据库连接
+func (s *Store) Close() error {
+	if s.db != nil {
+		return s.db.Close()
+	}
+	return nil
+}
+
+// Set 存储键值对到指定 bucket
+func (s *Store) Set(bucket []byte, key string, value interface{}) error {
+	data, err := json.Marshal(value)
+	if err != nil {
+		return fmt.Errorf("序列化数据失败: %v", err)
+	}
+
+	return s.db.Update(func(tx *bolt.Tx) error {
+		b := tx.Bucket(bucket)
+		if b == nil {
+			return fmt.Errorf("bucket %s 不存在", string(bucket))
+		}
+		return b.Put([]byte(key), data)
+	})
+}
+
+// Get 从指定 bucket 获取值
+func (s *Store) Get(bucket []byte, key string, dest interface{}) error {
+	return s.db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket(bucket)
+		if b == nil {
+			return fmt.Errorf("bucket %s 不存在", string(bucket))
+		}
+
+		data := b.Get([]byte(key))
+		if data == nil {
+			return fmt.Errorf("键 %s 不存在", key)
+		}
+
+		return json.Unmarshal(data, dest)
+	})
+}
+
+// Delete 删除指定键
+func (s *Store) Delete(bucket []byte, key string) error {
+	return s.db.Update(func(tx *bolt.Tx) error {
+		b := tx.Bucket(bucket)
+		if b == nil {
+			return fmt.Errorf("bucket %s 不存在", string(bucket))
+		}
+		return b.Delete([]byte(key))
+	})
+}
+
+// GetAll 获取指定 bucket 的所有数据
+func (s *Store) GetAll(bucket []byte, destSlice interface{}) error {
+	var items []json.RawMessage
+
+	err := s.db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket(bucket)
+		if b == nil {
+			return fmt.Errorf("bucket %s 不存在", string(bucket))
+		}
+
+		return b.ForEach(func(k, v []byte) error {
+			items = append(items, json.RawMessage(v))
+			return nil
+		})
+	})
+
+	if err != nil {
+		return err
+	}
+
+	// 将 items 转换为目标切片类型
+	data, err := json.Marshal(items)
+	if err != nil {
+		return err
+	}
+
+	return json.Unmarshal(data, destSlice)
+}
+
+// Exists 检查键是否存在
+func (s *Store) Exists(bucket []byte, key string) bool {
+	exists := false
+	s.db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket(bucket)
+		if b == nil {
+			return nil
+		}
+		exists = b.Get([]byte(key)) != nil
+		return nil
+	})
+	return exists
+}
+
+// Count 统计指定 bucket 的键数量
+func (s *Store) Count(bucket []byte) int {
+	count := 0
+	s.db.View(func(tx *bolt.Tx) error {
+		b := tx.Bucket(bucket)
+		if b == nil {
+			return nil
+		}
+		b.ForEach(func(k, v []byte) error {
+			count++
+			return nil
+		})
+		return nil
+	})
+	return count
+}
+
+// Backup 备份数据库到指定路径
+func (s *Store) Backup(backupPath string) error {
+	// 确保备份目录存在
+	dir := filepath.Dir(backupPath)
+	if err := os.MkdirAll(dir, 0755); err != nil {
+		return fmt.Errorf("创建备份目录失败: %v", err)
+	}
+
+	return s.db.View(func(tx *bolt.Tx) error {
+		return tx.CopyFile(backupPath, 0600)
+	})
+}