♻️ Migrate from gorm/sqlite to boltdb

This commit is contained in:
2024-04-24 12:51:37 +08:00
parent 3d3b4e0bea
commit 566965bb6a
12 changed files with 183 additions and 264 deletions

View File

@ -1,56 +1,71 @@
package database
import (
"encoding/json"
"errors"
"path/filepath"
"sub2clash/common"
"sub2clash/logger"
"sub2clash/model"
"github.com/glebarez/sqlite"
"go.uber.org/zap"
"gorm.io/gorm"
"go.etcd.io/bbolt"
)
var DB *gorm.DB
var DB *bbolt.DB
func ConnectDB() error {
// 用上面的数据库连接初始化 gorm
err := common.MKDir("data")
if err != nil {
return err
}
db, err := gorm.Open(
sqlite.Open(filepath.Join("data", "sub2clash.db")), &gorm.Config{
Logger: nil,
},
)
// 确保数据目录存在
path := filepath.Join("data", "sub2clash.db")
// 打开或创建数据库文件
db, err := bbolt.Open(path, 0600, nil)
if err != nil {
return err
}
DB = db
err = db.AutoMigrate(&model.ShortLink{})
if err != nil {
// 确保存储桶存在
return db.Update(func(tx *bbolt.Tx) error {
_, err := tx.CreateBucketIfNotExists([]byte("ShortLinks"))
return err
})
}
func FindShortLinkByHash(hash string) (*model.ShortLink, error) {
var shortLink model.ShortLink
err := DB.View(func(tx *bbolt.Tx) error {
b := tx.Bucket([]byte("ShortLinks"))
v := b.Get([]byte(hash))
if v == nil {
return errors.New("ShortLink not found")
}
return json.Unmarshal(v, &shortLink)
})
if err != nil {
return nil, err
}
return nil
return &shortLink, nil
}
func FindShortLinkByUrl(url string, shortLink *model.ShortLink) *gorm.DB {
logger.Logger.Debug("find short link by url", zap.String("url", url))
return DB.Where("url = ?", url).First(&shortLink)
func SaveShortLink(shortLink *model.ShortLink) error {
return DB.Update(func(tx *bbolt.Tx) error {
b := tx.Bucket([]byte("ShortLinks"))
encoded, err := json.Marshal(shortLink)
if err != nil {
return err
}
return b.Put([]byte(shortLink.Hash), encoded)
})
}
func FindShortLinkByHash(hash string, shortLink *model.ShortLink) *gorm.DB {
logger.Logger.Debug("find short link by hash", zap.String("hash", hash))
return DB.Where("hash = ?", hash).First(&shortLink)
}
func SaveShortLink(shortLink *model.ShortLink) {
logger.Logger.Debug("save short link", zap.String("hash", shortLink.Hash))
DB.Save(shortLink)
}
func FirstOrCreateShortLink(shortLink *model.ShortLink) {
logger.Logger.Debug("first or create short link", zap.String("hash", shortLink.Hash))
DB.FirstOrCreate(shortLink)
func CheckShortLinkHashExists(hash string) (bool, error) {
exists := false
err := DB.View(func(tx *bbolt.Tx) error {
b := tx.Bucket([]byte("ShortLinks"))
v := b.Get([]byte(hash))
exists = v != nil
return nil
})
if err != nil {
return false, err
}
return exists, nil
}

View File

@ -26,5 +26,8 @@ func MkEssentialDir() error {
if err := MKDir("logs"); err != nil {
return errors.New("create logs dir failed" + err.Error())
}
if err := MKDir("data"); err != nil {
return errors.New("create data dir failed" + err.Error())
}
return nil
}

View File

@ -5,6 +5,7 @@ import (
"encoding/hex"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"sub2clash/config"
@ -57,11 +58,11 @@ func FetchSubscriptionFromAPI(url string) ([]byte, error) {
if err != nil {
return nil, err
}
defer func(Body io.ReadCloser) {
if Body != nil {
_ = Body.Close()
defer func(resp *http.Response) {
if resp != nil && resp.Body != nil {
_ = resp.Body.Close()
}
}(resp.Body)
}(resp)
data, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)