kenaito-media-picker/internal/store/store.go

// Package store 提供媒体信息的 SQLite 持久化（xorm + 纯 Go 驱动，免 CGO）。
package store

import (
	"database/sql"
	"path/filepath"
	"time"

	sqlite "github.com/glebarez/go-sqlite"
	"xorm.io/xorm"
)

// DBFileName 数据库文件名（存放于被整理目录的根目录，扫描时跳过）
const DBFileName = "media-info.db"

// init 将纯Go的sqlite驱动注册为 xorm 认识的 sqlite3 驱动名（避免CGO依赖）
func init() {
	sql.Register("sqlite3", &sqlite.Driver{})
}

// MediaInfo 媒体文件信息记录，按字段拆分存放，便于查询
type MediaInfo struct {
	Id        int64     `xorm:"pk autoincr 'id'"`
	FileMd5   string    `xorm:"char(32) index 'file_md5'"`           // 文件MD5，唯一标识
	FilePath  string    `xorm:"text 'file_path'"`                    // 扫描时的源路径
	FileName  string    `xorm:"varchar(255) 'file_name'"`            // 文件名(含扩展名)
	FileExt   string    `xorm:"varchar(32) 'file_ext'"`              // 扩展名
	FileSize  int64     `xorm:"bigint 'file_size'"`                  // 文件大小(字节)
	MimeType  string    `xorm:"varchar(128) 'mime_type'"`            // MIME类型
	MediaType string    `xorm:"varchar(16) index 'media_type'"`      // image/video/audio
	Width     int       `xorm:"int index 'width'"`                   // 宽(像素)
	Height    int       `xorm:"int index 'height'"`                  // 高(像素)
	Duration  float64   `xorm:"real 'duration'"`                     // 时长(秒)
	Codec     string    `xorm:"varchar(64) 'codec'"`                 // 视频编码，如h264/hevc
	Format    string    `xorm:"varchar(64) 'format'"`                // 容器格式，如mp4/mkv
	RawInfo   string    `xorm:"text 'raw_info'"`                     // ffprobe原始JSON（完整信息兜底）
	HeadMd5   string    `xorm:"char(32) index('idx_fp') 'head_md5'"` // 文件头256KB的MD5（指纹加速）
	TailMd5   string    `xorm:"char(32) index('idx_fp') 'tail_md5'"` // 文件尾256KB的MD5（指纹加速）
	Category  string    `xorm:"varchar(64) index 'category'"`        // 归类的目录名（去向）
	DestDir   string    `xorm:"text 'dest_dir'"`                     // 移动后的目录路径（去向）
	CreatedAt time.Time `xorm:"created 'created_at'"`                // 首次入库时间
	UpdatedAt time.Time `xorm:"updated 'updated_at'"`                // 最近更新时间
}

// Store 封装 SQLite 存取，按实例持有连接，不再依赖全局状态
type Store struct {
	engine *xorm.Engine
}

// Open 打开（或创建）数据库并同步表结构
func Open(dbPath string) (*Store, error) {
	e, err := xorm.NewEngine("sqlite3", dbPath)
	if err != nil {
		return nil, err
	}
	e.ShowSQL(false)
	e.SetMaxOpenConns(1) // sqlite同一时刻只允许一个写连接
	if err := e.Sync(new(MediaInfo)); err != nil {
		_ = e.Close()
		return nil, err
	}
	return &Store{engine: e}, nil
}

// Close 关闭数据库连接
func (s *Store) Close() error {
	if s == nil || s.engine == nil {
		return nil
	}
	err := s.engine.Close()
	s.engine = nil
	return err
}

// SaveMedia 按MD5去重保存/更新媒体信息
func (s *Store) SaveMedia(info *MediaInfo) error {
	if s == nil || s.engine == nil {
		return nil
	}
	existing := new(MediaInfo)
	has, err := s.engine.Where("file_md5 = ?", info.FileMd5).Get(existing)
	if err != nil {
		return err
	}
	if has {
		info.Id = existing.Id
		_, err = s.engine.ID(existing.Id).Update(info)
		return err
	}
	_, err = s.engine.Insert(info)
	return err
}

// FindMediaByFingerprint 按 (文件大小 + 头256KB + 尾256KB) 查找已入库记录，避免对大文件重复计算全量MD5
func (s *Store) FindMediaByFingerprint(size int64, headMd5, tailMd5 string) (*MediaInfo, bool) {
	if s == nil || s.engine == nil || headMd5 == "" || tailMd5 == "" {
		return nil, false
	}
	info := new(MediaInfo)
	has, err := s.engine.Where("file_size = ? AND head_md5 = ? AND tail_md5 = ?", size, headMd5, tailMd5).Get(info)
	if err != nil || !has {
		return nil, false
	}
	return info, true
}

// FindMediaByMd5 按文件MD5查找已入库记录（用于识别重复文件）
func (s *Store) FindMediaByMd5(md5 string) (*MediaInfo, bool) {
	if s == nil || s.engine == nil || md5 == "" {
		return nil, false
	}
	info := new(MediaInfo)
	has, err := s.engine.Where("file_md5 = ?", md5).Get(info)
	if err != nil || !has {
		return nil, false
	}
	return info, true
}

// UpdateMediaDest 按源路径回填文件的归类目录与去向
func (s *Store) UpdateMediaDest(filePath, destDir string) error {
	if s == nil || s.engine == nil {
		return nil
	}
	_, err := s.engine.Where("file_path = ?", filePath).Update(&MediaInfo{
		Category: filepath.Base(destDir),
		DestDir:  destDir,
	})
	return err
}