go-site-clone/storage/store.go

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)
	})
}