auto decrypt

internal/wechatdb/datasource/darwinv3/datasource.go

@@ -3,87 +3,127 @@ package darwinv3
 import (
 	"context"
 	"crypto/md5"
-	"database/sql"
 	"encoding/hex"
 	"fmt"
 	"strings"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
 	_ "github.com/mattn/go-sqlite3"
 	"github.com/rs/zerolog/log"
 
 	"github.com/sjzar/chatlog/internal/errors"
 	"github.com/sjzar/chatlog/internal/model"
-	"github.com/sjzar/chatlog/pkg/util"
+	"github.com/sjzar/chatlog/internal/wechatdb/datasource/dbm"
 )
 
 const (
-	MessageFilePattern  = "^msg_([0-9]?[0-9])?\\.db$"
-	ContactFilePattern  = "^wccontact_new2\\.db$"
-	ChatRoomFilePattern = "^group_new\\.db$"
-	SessionFilePattern  = "^session_new\\.db$"
-	MediaFilePattern    = "^hldata\\.db$"
+	Message  = "message"
+	Contact  = "contact"
+	ChatRoom = "chatroom"
+	Session  = "session"
+	Media    = "media"
 )
 
-type DataSource struct {
-	path       string
-	messageDbs []*sql.DB
-	contactDb  *sql.DB
-	chatRoomDb *sql.DB
-	sessionDb  *sql.DB
-	mediaDb    *sql.DB
+var Groups = []dbm.Group{
+	{
+		Name:      Message,
+		Pattern:   `^msg_([0-9]?[0-9])?\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Contact,
+		Pattern:   `^wccontact_new2\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      ChatRoom,
+		Pattern:   `group_new\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Session,
+		Pattern:   `^session_new\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Media,
+		Pattern:   `^hldata\.db$`,
+		BlackList: []string{},
+	},
+}
 
-	talkerDBMap      map[string]*sql.DB
+type DataSource struct {
+	path string
+	dbm  *dbm.DBManager
+
+	talkerDBMap      map[string]string
 	user2DisplayName map[string]string
 }
 
 func New(path string) (*DataSource, error) {
 	ds := &DataSource{
 		path:             path,
-		messageDbs:       make([]*sql.DB, 0),
-		talkerDBMap:      make(map[string]*sql.DB),
+		dbm:              dbm.NewDBManager(path),
+		talkerDBMap:      make(map[string]string),
 		user2DisplayName: make(map[string]string),
 	}
 
-	if err := ds.initMessageDbs(path); err != nil {
+	for _, g := range Groups {
+		ds.dbm.AddGroup(g)
+	}
+
+	if err := ds.dbm.Start(); err != nil {
+		return nil, err
+	}
+
+	if err := ds.initMessageDbs(); err != nil {
 		return nil, errors.DBInitFailed(err)
 	}
-	if err := ds.initContactDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initChatRoomDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initSessionDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initMediaDb(path); err != nil {
+	if err := ds.initChatRoomDb(); err != nil {
 		return nil, errors.DBInitFailed(err)
 	}
 
+	ds.dbm.AddCallback(Message, func(event fsnotify.Event) error {
+		if !event.Op.Has(fsnotify.Create) {
+			return nil
+		}
+		if err := ds.initMessageDbs(); err != nil {
+			log.Err(err).Msgf("Failed to reinitialize message DBs: %s", event.Name)
+		}
+		return nil
+	})
+	ds.dbm.AddCallback(ChatRoom, func(event fsnotify.Event) error {
+		if !event.Op.Has(fsnotify.Create) {
+			return nil
+		}
+		if err := ds.initChatRoomDb(); err != nil {
+			log.Err(err).Msgf("Failed to reinitialize chatroom DB: %s", event.Name)
+		}
+		return nil
+	})
+
 	return ds, nil
 }
 
-func (ds *DataSource) initMessageDbs(path string) error {
+func (ds *DataSource) SetCallback(name string, callback func(event fsnotify.Event) error) error {
+	return ds.dbm.AddCallback(name, callback)
+}
 
-	files, err := util.FindFilesWithPatterns(path, MessageFilePattern, true)
+func (ds *DataSource) initMessageDbs() error {
+
+	dbPaths, err := ds.dbm.GetDBPath(Message)
 	if err != nil {
-		return errors.DBFileNotFound(path, MessageFilePattern, err)
+		return err
 	}
-
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, MessageFilePattern, nil)
-	}
-
 	// 处理每个数据库文件
-	for _, filePath := range files {
-		// 连接数据库
-		db, err := sql.Open("sqlite3", filePath)
+	talkerDBMap := make(map[string]string)
+	for _, filePath := range dbPaths {
+		db, err := ds.dbm.OpenDB(filePath)
 		if err != nil {
-			log.Err(err).Msgf("连接数据库 %s 失败", filePath)
+			log.Err(err).Msgf("获取数据库 %s 失败", filePath)
 			continue
 		}
-		ds.messageDbs = append(ds.messageDbs, db)
 
 		// 获取所有表名
 		rows, err := db.Query("SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'Chat_%'")
@@ -104,96 +144,42 @@ func (ds *DataSource) initMessageDbs(path string) error {
 			if talkerMd5 == "" {
 				continue
 			}
-			ds.talkerDBMap[talkerMd5] = db
+			talkerDBMap[talkerMd5] = filePath
 		}
 		rows.Close()
-
 	}
+	ds.talkerDBMap = talkerDBMap
 	return nil
 }
 
-func (ds *DataSource) initContactDb(path string) error {
-
-	files, err := util.FindFilesWithPatterns(path, ContactFilePattern, true)
+func (ds *DataSource) initChatRoomDb() error {
+	db, err := ds.dbm.GetDB(ChatRoom)
 	if err != nil {
-		return errors.DBFileNotFound(path, ContactFilePattern, err)
+		return err
 	}
 
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, ContactFilePattern, nil)
-	}
-
-	ds.contactDb, err = sql.Open("sqlite3", files[0])
+	rows, err := db.Query("SELECT m_nsUsrName, IFNULL(nickname,\"\") FROM GroupMember")
 	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-
-	return nil
-}
-
-func (ds *DataSource) initChatRoomDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, ChatRoomFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, ChatRoomFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, ChatRoomFilePattern, nil)
-	}
-	ds.chatRoomDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-
-	rows, err := ds.chatRoomDb.Query("SELECT m_nsUsrName, IFNULL(nickname,\"\") FROM GroupMember")
-	if err != nil {
-		log.Err(err).Msgf("数据库 %s 获取群聊成员失败", files[0])
+		log.Err(err).Msg("获取群聊成员失败")
 		return nil
 	}
 
+	user2DisplayName := make(map[string]string)
 	for rows.Next() {
 		var user string
 		var nickName string
 		if err := rows.Scan(&user, &nickName); err != nil {
-			log.Err(err).Msgf("数据库 %s 扫描表名失败", files[0])
+			log.Err(err).Msg("扫描表名失败")
 			continue
 		}
-		ds.user2DisplayName[user] = nickName
+		user2DisplayName[user] = nickName
 	}
 	rows.Close()
+	ds.user2DisplayName = user2DisplayName
 
 	return nil
 }
 
-func (ds *DataSource) initSessionDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, SessionFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, SessionFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, SessionFilePattern, nil)
-	}
-	ds.sessionDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-	return nil
-}
-
-func (ds *DataSource) initMediaDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, MediaFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, MediaFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, MediaFilePattern, nil)
-	}
-	ds.mediaDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-	return nil
-}
-
 // GetMessages 实现获取消息的方法
 func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
 	// 在 darwinv3 中，每个联系人/群聊的消息存储在单独的表中，表名为 Chat_md5(talker)
@@ -204,10 +190,14 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
 
 	_talkerMd5Bytes := md5.Sum([]byte(talker))
 	talkerMd5 := hex.EncodeToString(_talkerMd5Bytes[:])
-	db, ok := ds.talkerDBMap[talkerMd5]
+	dbPath, ok := ds.talkerDBMap[talkerMd5]
 	if !ok {
 		return nil, errors.TalkerNotFound(talker)
 	}
+	db, err := ds.dbm.OpenDB(dbPath)
+	if err != nil {
+		return nil, err
+	}
 	tableName := fmt.Sprintf("Chat_%s", talkerMd5)
 
 	// 构建查询条件
@@ -297,7 +287,11 @@ func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset
 	}
 
 	// 执行查询
-	rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Contact)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -351,7 +345,11 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 	}
 
 	// 执行查询
-	rows, err := ds.chatRoomDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(ChatRoom)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -380,7 +378,7 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 		contacts, err := ds.GetContacts(ctx, key, 1, 0)
 		if err == nil && len(contacts) > 0 && strings.HasSuffix(contacts[0].UserName, "@chatroom") {
 			// 再次尝试通过用户名查找群聊
-			rows, err := ds.chatRoomDb.QueryContext(ctx,
+			rows, err := db.QueryContext(ctx,
 				`SELECT IFNULL(m_nsUsrName,""), IFNULL(nickname,""), IFNULL(m_nsRemark,""), IFNULL(m_nsChatRoomMemList,""), IFNULL(m_nsChatRoomAdminList,"") 
 				FROM GroupContact 
 				WHERE m_nsUsrName = ?`,
@@ -448,7 +446,11 @@ func (ds *DataSource) GetSessions(ctx context.Context, key string, limit, offset
 	}
 
 	// 执行查询
-	rows, err := ds.sessionDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Session)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -506,7 +508,11 @@ WHERE
     r.mediaMd5 = ?`
 	args := []interface{}{key}
 	// 执行查询
-	rows, err := ds.mediaDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Media)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -541,46 +547,5 @@ WHERE
 
 // Close 实现关闭数据库连接的方法
 func (ds *DataSource) Close() error {
-	var errs []error
-
-	// 关闭消息数据库连接
-	for _, db := range ds.messageDbs {
-		if err := db.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭联系人数据库连接
-	if ds.contactDb != nil {
-		if err := ds.contactDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭群聊数据库连接
-	if ds.chatRoomDb != nil {
-		if err := ds.chatRoomDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭会话数据库连接
-	if ds.sessionDb != nil {
-		if err := ds.sessionDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭媒体数据库连接
-	if ds.mediaDb != nil {
-		if err := ds.mediaDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	if len(errs) > 0 {
-		return errors.DBCloseFailed(errs[0])
-	}
-
-	return nil
+	return ds.dbm.Close()
 }

internal/wechatdb/datasource/datasource.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
+
 	"github.com/sjzar/chatlog/internal/errors"
 	"github.com/sjzar/chatlog/internal/model"
 	"github.com/sjzar/chatlog/internal/wechatdb/datasource/darwinv3"
@@ -28,6 +30,9 @@ type DataSource interface {
 	// 媒体
 	GetMedia(ctx context.Context, _type string, key string) (*model.Media, error)
 
+	// 设置回调函数
+	SetCallback(name string, callback func(event fsnotify.Event) error) error
+
 	Close() error
 }
 

internal/wechatdb/datasource/dbm/dbm.go

@@ -0,0 +1,170 @@
+package dbm
+
+import (
+	"database/sql"
+	"runtime"
+	"sync"
+	"time"
+
+	"github.com/fsnotify/fsnotify"
+	_ "github.com/mattn/go-sqlite3"
+	"github.com/rs/zerolog/log"
+
+	"github.com/sjzar/chatlog/internal/errors"
+	"github.com/sjzar/chatlog/pkg/filecopy"
+	"github.com/sjzar/chatlog/pkg/filemonitor"
+)
+
+type DBManager struct {
+	path    string
+	fm      *filemonitor.FileMonitor
+	fgs     map[string]*filemonitor.FileGroup
+	dbs     map[string]*sql.DB
+	dbPaths map[string][]string
+	mutex   sync.RWMutex
+}
+
+func NewDBManager(path string) *DBManager {
+	return &DBManager{
+		path:    path,
+		fm:      filemonitor.NewFileMonitor(),
+		fgs:     make(map[string]*filemonitor.FileGroup),
+		dbs:     make(map[string]*sql.DB),
+		dbPaths: make(map[string][]string),
+	}
+}
+
+func (d *DBManager) AddGroup(g Group) error {
+	fg, err := filemonitor.NewFileGroup(g.Name, d.path, g.Pattern, g.BlackList)
+	if err != nil {
+		return err
+	}
+	fg.AddCallback(d.Callback)
+	d.fm.AddGroup(fg)
+	d.mutex.Lock()
+	d.fgs[g.Name] = fg
+	d.mutex.Unlock()
+	return nil
+}
+
+func (d *DBManager) AddCallback(name string, callback func(event fsnotify.Event) error) error {
+	d.mutex.RLock()
+	fg, ok := d.fgs[name]
+	d.mutex.RUnlock()
+	if !ok {
+		return errors.FileGroupNotFound(name)
+	}
+	fg.AddCallback(callback)
+	return nil
+}
+
+func (d *DBManager) GetDB(name string) (*sql.DB, error) {
+	dbPaths, err := d.GetDBPath(name)
+	if err != nil {
+		return nil, err
+	}
+	return d.OpenDB(dbPaths[0])
+}
+
+func (d *DBManager) GetDBs(name string) ([]*sql.DB, error) {
+	dbPaths, err := d.GetDBPath(name)
+	if err != nil {
+		return nil, err
+	}
+	dbs := make([]*sql.DB, 0)
+	for _, file := range dbPaths {
+		db, err := d.OpenDB(file)
+		if err != nil {
+			return nil, err
+		}
+		dbs = append(dbs, db)
+	}
+	return dbs, nil
+}
+
+func (d *DBManager) GetDBPath(name string) ([]string, error) {
+	d.mutex.RLock()
+	dbPaths, ok := d.dbPaths[name]
+	d.mutex.RUnlock()
+	if !ok {
+		d.mutex.RLock()
+		fg, ok := d.fgs[name]
+		d.mutex.RUnlock()
+		if !ok {
+			return nil, errors.FileGroupNotFound(name)
+		}
+		list, err := fg.List()
+		if err != nil {
+			return nil, errors.DBFileNotFound(d.path, fg.PatternStr, err)
+		}
+		if len(list) == 0 {
+			return nil, errors.DBFileNotFound(d.path, fg.PatternStr, nil)
+		}
+		dbPaths = list
+		d.mutex.Lock()
+		d.dbPaths[name] = dbPaths
+		d.mutex.Unlock()
+	}
+	return dbPaths, nil
+}
+
+func (d *DBManager) OpenDB(path string) (*sql.DB, error) {
+	d.mutex.RLock()
+	db, ok := d.dbs[path]
+	d.mutex.RUnlock()
+	if ok {
+		return db, nil
+	}
+	var err error
+	tempPath := path
+	if runtime.GOOS == "windows" {
+		tempPath, err = filecopy.GetTempCopy(path)
+		if err != nil {
+			log.Err(err).Msgf("获取临时拷贝文件 %s 失败", path)
+			return nil, err
+		}
+	}
+	db, err = sql.Open("sqlite3", tempPath)
+	if err != nil {
+		log.Err(err).Msgf("连接数据库 %s 失败", path)
+		return nil, err
+	}
+	d.mutex.Lock()
+	d.dbs[path] = db
+	d.mutex.Unlock()
+	return db, nil
+}
+
+func (d *DBManager) Callback(event fsnotify.Event) error {
+	if !event.Op.Has(fsnotify.Create) {
+		return nil
+	}
+
+	d.mutex.Lock()
+	db, ok := d.dbs[event.Name]
+	if ok {
+		delete(d.dbs, event.Name)
+		go func(db *sql.DB) {
+			time.Sleep(time.Second * 5)
+			db.Close()
+		}(db)
+	}
+	d.mutex.Unlock()
+
+	return nil
+}
+
+func (d *DBManager) Start() error {
+	return d.fm.Start()
+}
+
+func (d *DBManager) Stop() error {
+	return d.fm.Stop()
+}
+
+func (d *DBManager) Close() error {
+	for _, db := range d.dbs {
+		db.Close()
+	}
+	return d.fm.Stop()
+}

internal/wechatdb/datasource/dbm/dbm_test.go

@@ -0,0 +1,42 @@
+package dbm
+
+import (
+	"fmt"
+	"testing"
+	"time"
+)
+
+func TestXxx(t *testing.T) {
+	path := "/Users/sarv/Documents/chatlog/bigjun_9e7a"
+
+	g := Group{
+		Name:      "session",
+		Pattern:   `session\.db$`,
+		BlackList: []string{},
+	}
+
+	d := NewDBManager(path)
+	d.AddGroup(g)
+	d.Start()
+
+	i := 0
+	for {
+		db, err := d.GetDB("session")
+		if err != nil {
+			fmt.Println(err)
+			break
+		}
+
+		var username string
+		row := db.QueryRow(`SELECT username FROM SessionTable LIMIT 1`)
+		if err := row.Scan(&username); err != nil {
+			fmt.Printf("Error scanning row: %v\n", err)
+			time.Sleep(100 * time.Millisecond)
+			continue
+		}
+		fmt.Printf("%d: Username: %s\n", i, username)
+		i++
+		time.Sleep(1000 * time.Millisecond)
+	}
+
+}

internal/wechatdb/datasource/dbm/group.go

@@ -0,0 +1,7 @@
+package dbm
+
+type Group struct {
+	Name      string
+	Pattern   string
+	BlackList []string
+}

internal/wechatdb/datasource/v4/datasource.go

@@ -10,22 +10,51 @@ import (
 	"strings"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
 	_ "github.com/mattn/go-sqlite3"
 	"github.com/rs/zerolog/log"
 
 	"github.com/sjzar/chatlog/internal/errors"
 	"github.com/sjzar/chatlog/internal/model"
-	"github.com/sjzar/chatlog/pkg/util"
+	"github.com/sjzar/chatlog/internal/wechatdb/datasource/dbm"
 )
 
 const (
-	MessageFilePattern = "^message_([0-9]?[0-9])?\\.db$"
-	ContactFilePattern = "^contact\\.db$"
-	SessionFilePattern = "^session\\.db$"
-	MediaFilePattern   = "^hardlink\\.db$"
-	VoiceFilePattern   = "^media_([0-9]?[0-9])?\\.db$"
+	Message = "message"
+	Contact = "contact"
+	Session = "session"
+	Media   = "media"
+	Voice   = "voice"
 )
 
+var Groups = []dbm.Group{
+	{
+		Name:      Message,
+		Pattern:   `^message_([0-9]?[0-9])?\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Contact,
+		Pattern:   `^contact\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Session,
+		Pattern:   `session\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Media,
+		Pattern:   `^hardlink\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Voice,
+		Pattern:   `^media_([0-9]?[0-9])?\.db$`,
+		BlackList: []string{},
+	},
+}
+
 // MessageDBInfo 存储消息数据库的信息
 type MessageDBInfo struct {
 	FilePath  string
@@ -34,61 +63,65 @@ type MessageDBInfo struct {
 }
 
 type DataSource struct {
-	path       string
-	messageDbs map[string]*sql.DB
-	contactDb  *sql.DB
-	sessionDb  *sql.DB
-	mediaDb    *sql.DB
-	voiceDb    []*sql.DB
+	path string
+	dbm  *dbm.DBManager
 
 	// 消息数据库信息
-	messageFiles []MessageDBInfo
+	messageInfos []MessageDBInfo
 }
 
 func New(path string) (*DataSource, error) {
+
 	ds := &DataSource{
 		path:         path,
-		messageDbs:   make(map[string]*sql.DB),
-		voiceDb:      make([]*sql.DB, 0),
-		messageFiles: make([]MessageDBInfo, 0),
+		dbm:          dbm.NewDBManager(path),
+		messageInfos: make([]MessageDBInfo, 0),
 	}
 
-	if err := ds.initMessageDbs(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initContactDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initSessionDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initMediaDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-	if err := ds.initVoiceDb(path); err != nil {
+	for _, g := range Groups {
+		ds.dbm.AddGroup(g)
+	}
+
+	if err := ds.dbm.Start(); err != nil {
+		return nil, err
+	}
+
+	if err := ds.initMessageDbs(); err != nil {
 		return nil, errors.DBInitFailed(err)
 	}
 
+	ds.dbm.AddCallback(Message, func(event fsnotify.Event) error {
+		if !event.Op.Has(fsnotify.Create) {
+			return nil
+		}
+		if err := ds.initMessageDbs(); err != nil {
+			log.Err(err).Msgf("Failed to reinitialize message DBs: %s", event.Name)
+		}
+		return nil
+	})
+
 	return ds, nil
 }
 
-func (ds *DataSource) initMessageDbs(path string) error {
-	// 查找所有消息数据库文件
-	files, err := util.FindFilesWithPatterns(path, MessageFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, MessageFilePattern, err)
+func (ds *DataSource) SetCallback(name string, callback func(event fsnotify.Event) error) error {
+	if name == "chatroom" {
+		name = Contact
 	}
+	return ds.dbm.AddCallback(name, callback)
+}
 
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, MessageFilePattern, nil)
+func (ds *DataSource) initMessageDbs() error {
+	dbPaths, err := ds.dbm.GetDBPath(Message)
+	if err != nil {
+		return err
 	}
 
 	// 处理每个数据库文件
-	for _, filePath := range files {
-		// 连接数据库
-		db, err := sql.Open("sqlite3", filePath)
+	infos := make([]MessageDBInfo, 0)
+	for _, filePath := range dbPaths {
+		db, err := ds.dbm.OpenDB(filePath)
 		if err != nil {
-			log.Err(err).Msgf("连接数据库 %s 失败", filePath)
+			log.Err(err).Msgf("获取数据库 %s 失败", filePath)
 			continue
 		}
 
@@ -99,108 +132,38 @@ func (ds *DataSource) initMessageDbs(path string) error {
 		row := db.QueryRow("SELECT timestamp FROM Timestamp LIMIT 1")
 		if err := row.Scan(&timestamp); err != nil {
 			log.Err(err).Msgf("获取数据库 %s 的时间戳失败", filePath)
-			db.Close()
 			continue
 		}
 		startTime = time.Unix(timestamp, 0)
 
 		// 保存数据库信息
-		ds.messageFiles = append(ds.messageFiles, MessageDBInfo{
+		infos = append(infos, MessageDBInfo{
 			FilePath:  filePath,
 			StartTime: startTime,
 		})
-
-		// 保存数据库连接
-		ds.messageDbs[filePath] = db
 	}
 
 	// 按照 StartTime 排序数据库文件
-	sort.Slice(ds.messageFiles, func(i, j int) bool {
-		return ds.messageFiles[i].StartTime.Before(ds.messageFiles[j].StartTime)
+	sort.Slice(infos, func(i, j int) bool {
+		return infos[i].StartTime.Before(infos[j].StartTime)
 	})
 
 	// 设置结束时间
-	for i := range ds.messageFiles {
-		if i == len(ds.messageFiles)-1 {
-			ds.messageFiles[i].EndTime = time.Now()
+	for i := range infos {
+		if i == len(infos)-1 {
+			infos[i].EndTime = time.Now()
 		} else {
-			ds.messageFiles[i].EndTime = ds.messageFiles[i+1].StartTime
+			infos[i].EndTime = infos[i+1].StartTime
 		}
 	}
-
-	return nil
-}
-
-func (ds *DataSource) initContactDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, ContactFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, ContactFilePattern, err)
-	}
-
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, ContactFilePattern, nil)
-	}
-
-	ds.contactDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-
-	return nil
-}
-
-func (ds *DataSource) initSessionDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, SessionFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, SessionFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, SessionFilePattern, nil)
-	}
-	ds.sessionDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-	return nil
-}
-
-func (ds *DataSource) initMediaDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, MediaFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, MediaFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, MediaFilePattern, nil)
-	}
-	ds.mediaDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-	return nil
-}
-
-func (ds *DataSource) initVoiceDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, VoiceFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, VoiceFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, VoiceFilePattern, nil)
-	}
-	for _, file := range files {
-		db, err := sql.Open("sqlite3", file)
-		if err != nil {
-			return errors.DBConnectFailed(files[0], err)
-		}
-		ds.voiceDb = append(ds.voiceDb, db)
-	}
+	ds.messageInfos = infos
 	return nil
 }
 
 // getDBInfosForTimeRange 获取时间范围内的数据库信息
 func (ds *DataSource) getDBInfosForTimeRange(startTime, endTime time.Time) []MessageDBInfo {
 	var dbs []MessageDBInfo
-	for _, info := range ds.messageFiles {
+	for _, info := range ds.messageInfos {
 		if info.StartTime.Before(endTime) && info.EndTime.After(startTime) {
 			dbs = append(dbs, info)
 		}
@@ -234,8 +197,8 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
 			return nil, err
 		}
 
-		db, ok := ds.messageDbs[dbInfo.FilePath]
-		if !ok {
+		db, err := ds.dbm.OpenDB(dbInfo.FilePath)
+		if err != nil {
 			log.Error().Msgf("数据库 %s 未打开", dbInfo.FilePath)
 			continue
 		}
@@ -275,8 +238,8 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
 
 // getMessagesSingleFile 从单个数据库文件获取消息
 func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageDBInfo, startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
-	db, ok := ds.messageDbs[dbInfo.FilePath]
-	if !ok {
+	db, err := ds.dbm.OpenDB(dbInfo.FilePath)
+	if err != nil {
 		return nil, errors.DBConnectFailed(dbInfo.FilePath, nil)
 	}
 
@@ -287,7 +250,7 @@ func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageD
 
 	// 检查表是否存在
 	var exists bool
-	err := db.QueryRowContext(ctx,
+	err = db.QueryRowContext(ctx,
 		"SELECT 1 FROM sqlite_master WHERE type='table' AND name=?",
 		tableName).Scan(&exists)
 
@@ -445,7 +408,11 @@ func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset
 	}
 
 	// 执行查询
-	rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Contact)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -477,13 +444,18 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 	var query string
 	var args []interface{}
 
+	// 执行查询
+	db, err := ds.dbm.GetDB(Contact)
+	if err != nil {
+		return nil, err
+	}
+
 	if key != "" {
 		// 按照关键字查询
 		query = `SELECT username, owner, ext_buffer FROM chat_room WHERE username = ?`
 		args = []interface{}{key}
 
-		// 执行查询
-		rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return nil, errors.QueryFailed(query, err)
 		}
@@ -510,7 +482,7 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 			contacts, err := ds.GetContacts(ctx, key, 1, 0)
 			if err == nil && len(contacts) > 0 && strings.HasSuffix(contacts[0].UserName, "@chatroom") {
 				// 再次尝试通过用户名查找群聊
-				rows, err := ds.contactDb.QueryContext(ctx,
+				rows, err := db.QueryContext(ctx,
 					`SELECT username, owner, ext_buffer FROM chat_room WHERE username = ?`,
 					contacts[0].UserName)
 
@@ -560,7 +532,7 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 		}
 
 		// 执行查询
-		rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return nil, errors.QueryFailed(query, err)
 		}
@@ -614,7 +586,11 @@ func (ds *DataSource) GetSessions(ctx context.Context, key string, limit, offset
 	}
 
 	// 执行查询
-	rows, err := ds.sessionDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Session)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -678,7 +654,12 @@ func (ds *DataSource) GetMedia(ctx context.Context, _type string, key string) (*
 	query += " WHERE f.md5 = ? OR f.file_name LIKE ? || '%'"
 	args := []interface{}{key, key}
 
-	rows, err := ds.mediaDb.QueryContext(ctx, query, args...)
+	// 执行查询
+	db, err := ds.dbm.GetDB(Media)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -726,7 +707,12 @@ func (ds *DataSource) GetVoice(ctx context.Context, key string) (*model.Media, e
 	`
 	args := []interface{}{key}
 
-	for _, db := range ds.voiceDb {
+	dbs, err := ds.dbm.GetDBs(Voice)
+	if err != nil {
+		return nil, errors.DBConnectFailed("", err)
+	}
+
+	for _, db := range dbs {
 		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return nil, errors.QueryFailed(query, err)
@@ -755,38 +741,5 @@ func (ds *DataSource) GetVoice(ctx context.Context, key string) (*model.Media, e
 }
 
 func (ds *DataSource) Close() error {
-	var errs []error
-
-	// 关闭消息数据库连接
-	for _, db := range ds.messageDbs {
-		if err := db.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭联系人数据库连接
-	if ds.contactDb != nil {
-		if err := ds.contactDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭会话数据库连接
-	if ds.sessionDb != nil {
-		if err := ds.sessionDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	if ds.mediaDb != nil {
-		if err := ds.mediaDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	if len(errs) > 0 {
-		return errors.DBCloseFailed(errs[0])
-	}
-
-	return nil
+	return ds.dbm.Close()
 }

internal/wechatdb/datasource/windowsv3/datasource.go

@@ -9,23 +9,57 @@ import (
 	"strings"
 	"time"
 
+	"github.com/fsnotify/fsnotify"
 	_ "github.com/mattn/go-sqlite3"
 	"github.com/rs/zerolog/log"
 
 	"github.com/sjzar/chatlog/internal/errors"
 	"github.com/sjzar/chatlog/internal/model"
-	"github.com/sjzar/chatlog/pkg/util"
+	"github.com/sjzar/chatlog/internal/wechatdb/datasource/dbm"
 )
 
 const (
-	MessageFilePattern = "^MSG([0-9]?[0-9])?\\.db$"
-	ContactFilePattern = "^MicroMsg.db$"
-	ImageFilePattern   = "^HardLinkImage\\.db$"
-	VideoFilePattern   = "^HardLinkVideo\\.db$"
-	FileFilePattern    = "^HardLinkFile\\.db$"
-	VoiceFilePattern   = "^MediaMSG([0-9])?\\.db$"
+	Message = "message"
+	Contact = "contact"
+	Image   = "image"
+	Video   = "video"
+	File    = "file"
+	Voice   = "voice"
 )
 
+var Groups = []dbm.Group{
+	{
+		Name:      Message,
+		Pattern:   `^MSG([0-9]?[0-9])?\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Contact,
+		Pattern:   `^MicroMsg.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Image,
+		Pattern:   `^HardLinkImage\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Video,
+		Pattern:   `^HardLinkVideo\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      File,
+		Pattern:   `^HardLinkFile\.db$`,
+		BlackList: []string{},
+	},
+	{
+		Name:      Voice,
+		Pattern:   `^MediaMSG([0-9])?\.db$`,
+		BlackList: []string{},
+	},
+}
+
 // MessageDBInfo 保存消息数据库的信息
 type MessageDBInfo struct {
 	FilePath  string
@@ -36,67 +70,67 @@ type MessageDBInfo struct {
 
 // DataSource 实现了 DataSource 接口
 type DataSource struct {
-	// 消息数据库
-	messageFiles []MessageDBInfo
-	messageDbs   map[string]*sql.DB
+	path string
+	dbm  *dbm.DBManager
 
-	// 联系人数据库
-	contactDbFile string
-	contactDb     *sql.DB
-
-	imageDb *sql.DB
-	videoDb *sql.DB
-	fileDb  *sql.DB
-	voiceDb []*sql.DB
+	// 消息数据库信息
+	messageInfos []MessageDBInfo
 }
 
 // New 创建一个新的 WindowsV3DataSource
 func New(path string) (*DataSource, error) {
 	ds := &DataSource{
-		messageFiles: make([]MessageDBInfo, 0),
-		messageDbs:   make(map[string]*sql.DB),
-		voiceDb:      make([]*sql.DB, 0),
+		path:         path,
+		dbm:          dbm.NewDBManager(path),
+		messageInfos: make([]MessageDBInfo, 0),
 	}
 
-	// 初始化消息数据库
-	if err := ds.initMessageDbs(path); err != nil {
+	for _, g := range Groups {
+		ds.dbm.AddGroup(g)
+	}
+
+	if err := ds.dbm.Start(); err != nil {
+		return nil, err
+	}
+
+	if err := ds.initMessageDbs(); err != nil {
 		return nil, errors.DBInitFailed(err)
 	}
 
-	// 初始化联系人数据库
-	if err := ds.initContactDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-
-	if err := ds.initMediaDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
-
-	if err := ds.initVoiceDb(path); err != nil {
-		return nil, errors.DBInitFailed(err)
-	}
+	ds.dbm.AddCallback(Message, func(event fsnotify.Event) error {
+		if !event.Op.Has(fsnotify.Create) {
+			return nil
+		}
+		if err := ds.initMessageDbs(); err != nil {
+			log.Err(err).Msgf("Failed to reinitialize message DBs: %s", event.Name)
+		}
+		return nil
+	})
 
 	return ds, nil
 }
 
-// initMessageDbs 初始化消息数据库
-func (ds *DataSource) initMessageDbs(path string) error {
-	// 查找所有消息数据库文件
-	files, err := util.FindFilesWithPatterns(path, MessageFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, MessageFilePattern, err)
+func (ds *DataSource) SetCallback(name string, callback func(event fsnotify.Event) error) error {
+	if name == "chatroom" {
+		name = Contact
 	}
+	return ds.dbm.AddCallback(name, callback)
+}
 
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, MessageFilePattern, nil)
+// initMessageDbs 初始化消息数据库
+func (ds *DataSource) initMessageDbs() error {
+	// 获取所有消息数据库文件路径
+	dbPaths, err := ds.dbm.GetDBPath(Message)
+	if err != nil {
+		return err
 	}
 
 	// 处理每个数据库文件
-	for _, filePath := range files {
-		// 连接数据库
-		db, err := sql.Open("sqlite3", filePath)
+	infos := make([]MessageDBInfo, 0)
+	for _, filePath := range dbPaths {
+		db, err := ds.dbm.OpenDB(filePath)
 		if err != nil {
-			log.Err(err).Msgf("连接数据库 %s 失败", filePath)
+			log.Err(err).Msgf("获取数据库 %s 失败", filePath)
 			continue
 		}
 
@@ -106,7 +140,6 @@ func (ds *DataSource) initMessageDbs(path string) error {
 		rows, err := db.Query("SELECT tableIndex, tableVersion, tableDesc FROM DBInfo")
 		if err != nil {
 			log.Err(err).Msgf("查询数据库 %s 的 DBInfo 表失败", filePath)
-			db.Close()
 			continue
 		}
 
@@ -133,7 +166,6 @@ func (ds *DataSource) initMessageDbs(path string) error {
 		rows, err = db.Query("SELECT UsrName FROM Name2ID")
 		if err != nil {
 			log.Err(err).Msgf("查询数据库 %s 的 Name2ID 表失败", filePath)
-			db.Close()
 			continue
 		}
 
@@ -150,123 +182,34 @@ func (ds *DataSource) initMessageDbs(path string) error {
 		rows.Close()
 
 		// 保存数据库信息
-		ds.messageFiles = append(ds.messageFiles, MessageDBInfo{
+		infos = append(infos, MessageDBInfo{
 			FilePath:  filePath,
 			StartTime: startTime,
 			TalkerMap: talkerMap,
 		})
-
-		// 保存数据库连接
-		ds.messageDbs[filePath] = db
 	}
 
 	// 按照 StartTime 排序数据库文件
-	sort.Slice(ds.messageFiles, func(i, j int) bool {
-		return ds.messageFiles[i].StartTime.Before(ds.messageFiles[j].StartTime)
+	sort.Slice(infos, func(i, j int) bool {
+		return infos[i].StartTime.Before(infos[j].StartTime)
 	})
 
 	// 设置结束时间
-	for i := range ds.messageFiles {
-		if i == len(ds.messageFiles)-1 {
-			ds.messageFiles[i].EndTime = time.Now()
+	for i := range infos {
+		if i == len(infos)-1 {
+			infos[i].EndTime = time.Now()
 		} else {
-			ds.messageFiles[i].EndTime = ds.messageFiles[i+1].StartTime
+			infos[i].EndTime = infos[i+1].StartTime
 		}
 	}
-
-	return nil
-}
-
-// initContactDb 初始化联系人数据库
-func (ds *DataSource) initContactDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, ContactFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, ContactFilePattern, err)
-	}
-
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, ContactFilePattern, nil)
-	}
-
-	ds.contactDbFile = files[0]
-
-	ds.contactDb, err = sql.Open("sqlite3", ds.contactDbFile)
-	if err != nil {
-		return errors.DBConnectFailed(ds.contactDbFile, err)
-	}
-
-	return nil
-}
-
-// initContactDb 初始化联系人数据库
-func (ds *DataSource) initMediaDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, ImageFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, ImageFilePattern, err)
-	}
-
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, ImageFilePattern, nil)
-	}
-
-	ds.imageDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-
-	files, err = util.FindFilesWithPatterns(path, VideoFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, VideoFilePattern, err)
-	}
-
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, VideoFilePattern, nil)
-	}
-
-	ds.videoDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-
-	files, err = util.FindFilesWithPatterns(path, FileFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, FileFilePattern, err)
-	}
-
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, FileFilePattern, nil)
-	}
-
-	ds.fileDb, err = sql.Open("sqlite3", files[0])
-	if err != nil {
-		return errors.DBConnectFailed(files[0], err)
-	}
-
-	return nil
-}
-
-func (ds *DataSource) initVoiceDb(path string) error {
-	files, err := util.FindFilesWithPatterns(path, VoiceFilePattern, true)
-	if err != nil {
-		return errors.DBFileNotFound(path, VoiceFilePattern, err)
-	}
-	if len(files) == 0 {
-		return errors.DBFileNotFound(path, VoiceFilePattern, nil)
-	}
-	for _, file := range files {
-		db, err := sql.Open("sqlite3", file)
-		if err != nil {
-			return errors.DBConnectFailed(files[0], err)
-		}
-		ds.voiceDb = append(ds.voiceDb, db)
-	}
+	ds.messageInfos = infos
 	return nil
 }
 
 // getDBInfosForTimeRange 获取时间范围内的数据库信息
 func (ds *DataSource) getDBInfosForTimeRange(startTime, endTime time.Time) []MessageDBInfo {
 	var dbs []MessageDBInfo
-	for _, info := range ds.messageFiles {
+	for _, info := range ds.messageInfos {
 		if info.StartTime.Before(endTime) && info.EndTime.After(startTime) {
 			dbs = append(dbs, info)
 		}
@@ -296,70 +239,19 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
 			return nil, err
 		}
 
-		db, ok := ds.messageDbs[dbInfo.FilePath]
-		if !ok {
+		db, err := ds.dbm.OpenDB(dbInfo.FilePath)
+		if err != nil {
 			log.Error().Msgf("数据库 %s 未打开", dbInfo.FilePath)
 			continue
 		}
 
-		// 构建查询条件
-		conditions := []string{"Sequence >= ? AND Sequence <= ?"}
-		args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
-
-		if len(talker) > 0 {
-			talkerID, ok := dbInfo.TalkerMap[talker]
-			if ok {
-				conditions = append(conditions, "TalkerId = ?")
-				args = append(args, talkerID)
-			} else {
-				conditions = append(conditions, "StrTalker = ?")
-				args = append(args, talker)
-			}
-		}
-
-		query := fmt.Sprintf(`
-            SELECT MsgSvrID, Sequence, CreateTime, StrTalker, IsSender, 
-                Type, SubType, StrContent, CompressContent, BytesExtra
-            FROM MSG 
-            WHERE %s 
-            ORDER BY Sequence ASC
-        `, strings.Join(conditions, " AND "))
-
-		// 执行查询
-		rows, err := db.QueryContext(ctx, query, args...)
+		messages, err := ds.getMessagesFromDB(ctx, db, dbInfo, startTime, endTime, talker)
 		if err != nil {
-			log.Err(err).Msgf("查询数据库 %s 失败", dbInfo.FilePath)
+			log.Err(err).Msgf("从数据库 %s 获取消息失败", dbInfo.FilePath)
 			continue
 		}
 
-		// 处理查询结果
-		for rows.Next() {
-			var msg model.MessageV3
-			var compressContent []byte
-			var bytesExtra []byte
-
-			err := rows.Scan(
-				&msg.MsgSvrID,
-				&msg.Sequence,
-				&msg.CreateTime,
-				&msg.StrTalker,
-				&msg.IsSender,
-				&msg.Type,
-				&msg.SubType,
-				&msg.StrContent,
-				&compressContent,
-				&bytesExtra,
-			)
-			if err != nil {
-				log.Err(err).Msg("扫描消息行失败")
-				continue
-			}
-			msg.CompressContent = compressContent
-			msg.BytesExtra = bytesExtra
-
-			totalMessages = append(totalMessages, msg.Wrap())
-		}
-		rows.Close()
+		totalMessages = append(totalMessages, messages...)
 
 		if limit+offset > 0 && len(totalMessages) >= limit+offset {
 			break
@@ -388,6 +280,11 @@ func (ds *DataSource) GetMessages(ctx context.Context, startTime, endTime time.T
 
 // getMessagesSingleFile 从单个数据库文件获取消息
 func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageDBInfo, startTime, endTime time.Time, talker string, limit, offset int) ([]*model.Message, error) {
+	db, err := ds.dbm.OpenDB(dbInfo.FilePath)
+	if err != nil {
+		return nil, errors.DBConnectFailed(dbInfo.FilePath, nil)
+	}
+
 	// 构建查询条件
 	conditions := []string{"Sequence >= ? AND Sequence <= ?"}
 	args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
@@ -419,7 +316,7 @@ func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageD
 	}
 
 	// 执行查询
-	rows, err := ds.messageDbs[dbInfo.FilePath].QueryContext(ctx, query, args...)
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -453,6 +350,69 @@ func (ds *DataSource) getMessagesSingleFile(ctx context.Context, dbInfo MessageD
 	return totalMessages, nil
 }
 
+// getMessagesFromDB 从数据库获取消息
+func (ds *DataSource) getMessagesFromDB(ctx context.Context, db *sql.DB, dbInfo MessageDBInfo, startTime, endTime time.Time, talker string) ([]*model.Message, error) {
+	// 构建查询条件
+	conditions := []string{"Sequence >= ? AND Sequence <= ?"}
+	args := []interface{}{startTime.Unix() * 1000, endTime.Unix() * 1000}
+
+	if len(talker) > 0 {
+		talkerID, ok := dbInfo.TalkerMap[talker]
+		if ok {
+			conditions = append(conditions, "TalkerId = ?")
+			args = append(args, talkerID)
+		} else {
+			conditions = append(conditions, "StrTalker = ?")
+			args = append(args, talker)
+		}
+	}
+
+	query := fmt.Sprintf(`
+		SELECT MsgSvrID, Sequence, CreateTime, StrTalker, IsSender, 
+			Type, SubType, StrContent, CompressContent, BytesExtra
+		FROM MSG 
+		WHERE %s 
+		ORDER BY Sequence ASC
+	`, strings.Join(conditions, " AND "))
+
+	// 执行查询
+	rows, err := db.QueryContext(ctx, query, args...)
+	if err != nil {
+		return nil, errors.QueryFailed(query, err)
+	}
+	defer rows.Close()
+
+	// 处理查询结果
+	messages := []*model.Message{}
+	for rows.Next() {
+		var msg model.MessageV3
+		var compressContent []byte
+		var bytesExtra []byte
+
+		err := rows.Scan(
+			&msg.MsgSvrID,
+			&msg.Sequence,
+			&msg.CreateTime,
+			&msg.StrTalker,
+			&msg.IsSender,
+			&msg.Type,
+			&msg.SubType,
+			&msg.StrContent,
+			&compressContent,
+			&bytesExtra,
+		)
+		if err != nil {
+			return nil, errors.ScanRowFailed(err)
+		}
+		msg.CompressContent = compressContent
+		msg.BytesExtra = bytesExtra
+
+		messages = append(messages, msg.Wrap())
+	}
+
+	return messages, nil
+}
+
 // GetContacts 实现获取联系人信息的方法
 func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset int) ([]*model.Contact, error) {
 	var query string
@@ -478,7 +438,11 @@ func (ds *DataSource) GetContacts(ctx context.Context, key string, limit, offset
 	}
 
 	// 执行查询
-	rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Contact)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -516,7 +480,11 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 		args = []interface{}{key}
 
 		// 执行查询
-		rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+		db, err := ds.dbm.GetDB(Contact)
+		if err != nil {
+			return nil, err
+		}
+		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return nil, errors.QueryFailed(query, err)
 		}
@@ -543,7 +511,7 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 			contacts, err := ds.GetContacts(ctx, key, 1, 0)
 			if err == nil && len(contacts) > 0 && strings.HasSuffix(contacts[0].UserName, "@chatroom") {
 				// 再次尝试通过用户名查找群聊
-				rows, err := ds.contactDb.QueryContext(ctx,
+				rows, err := db.QueryContext(ctx,
 					`SELECT ChatRoomName, Reserved2, RoomData FROM ChatRoom WHERE ChatRoomName = ?`,
 					contacts[0].UserName)
 
@@ -593,7 +561,11 @@ func (ds *DataSource) GetChatRooms(ctx context.Context, key string, limit, offse
 		}
 
 		// 执行查询
-		rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+		db, err := ds.dbm.GetDB(Contact)
+		if err != nil {
+			return nil, err
+		}
+		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return nil, errors.QueryFailed(query, err)
 		}
@@ -647,7 +619,11 @@ func (ds *DataSource) GetSessions(ctx context.Context, key string, limit, offset
 	}
 
 	// 执行查询
-	rows, err := ds.contactDb.QueryContext(ctx, query, args...)
+	db, err := ds.dbm.GetDB(Contact)
+	if err != nil {
+		return nil, err
+	}
+	rows, err := db.QueryContext(ctx, query, args...)
 	if err != nil {
 		return nil, errors.QueryFailed(query, err)
 	}
@@ -688,25 +664,29 @@ func (ds *DataSource) GetMedia(ctx context.Context, _type string, key string) (*
 		return nil, errors.DecodeKeyFailed(err)
 	}
 
-	var db *sql.DB
+	var dbType string
 	var table1, table2 string
 
 	switch _type {
 	case "image":
-		db = ds.imageDb
+		dbType = Image
 		table1 = "HardLinkImageAttribute"
 		table2 = "HardLinkImageID"
 	case "video":
-		db = ds.videoDb
+		dbType = Video
 		table1 = "HardLinkVideoAttribute"
 		table2 = "HardLinkVideoID"
 	case "file":
-		db = ds.fileDb
+		dbType = File
 		table1 = "HardLinkFileAttribute"
 		table2 = "HardLinkFileID"
 	default:
 		return nil, errors.MediaTypeUnsupported(_type)
+	}
 
+	db, err := ds.dbm.GetDB(dbType)
+	if err != nil {
+		return nil, err
 	}
 
 	query := fmt.Sprintf(`
@@ -768,7 +748,12 @@ func (ds *DataSource) GetVoice(ctx context.Context, key string) (*model.Media, e
 	`
 	args := []interface{}{key}
 
-	for _, db := range ds.voiceDb {
+	dbs, err := ds.dbm.GetDBs(Voice)
+	if err != nil {
+		return nil, errors.DBConnectFailed("", err)
+	}
+
+	for _, db := range dbs {
 		rows, err := db.QueryContext(ctx, query, args...)
 		if err != nil {
 			return nil, errors.QueryFailed(query, err)
@@ -798,41 +783,5 @@ func (ds *DataSource) GetVoice(ctx context.Context, key string) (*model.Media, e
 
 // Close 实现 DataSource 接口的 Close 方法
 func (ds *DataSource) Close() error {
-	var errs []error
-
-	// 关闭消息数据库连接
-	for _, db := range ds.messageDbs {
-		if err := db.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	// 关闭联系人数据库连接
-	if ds.contactDb != nil {
-		if err := ds.contactDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	if ds.imageDb != nil {
-		if err := ds.imageDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-	if ds.videoDb != nil {
-		if err := ds.videoDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-	if ds.fileDb != nil {
-		if err := ds.fileDb.Close(); err != nil {
-			errs = append(errs, err)
-		}
-	}
-
-	if len(errs) > 0 {
-		return errors.DBCloseFailed(errs[0])
-	}
-
-	return nil
+	return ds.dbm.Close()
 }