feat: 抓取任务队列 + 媒体机直传 + 向量回填
crawl（品牌抓取任务队列） - /admin/internal/crawl/tasks 领取 / :id/result 回报，spider 可常驻拉任务 - 后台品牌列表支持「加入抓取队列」，队列页可看状态与重试 - 租约超时的 running 自动回收为 queued，避免 worker 崩溃后任务卡死 media（媒体机直传） - ingest 载荷新增 main_key/detail_keys/main_size/detail_sizes： 有 key 就直接落库，跳过下载与上传；与 URL 形态可逐图混用 - 细节图按 Details/DetailKeys 两者较长者遍历（只给 key 时不再丢图） - 媒体机上传的 key 不进失败回滚清单（内容寻址、可能与他人共享） - /admin/internal/upload-url 签发 S4 PUT 预签名（spider 不持有 AK/SK） embeddings（向量回填） - GET /admin/internal/embeddings/pending 派发缺向量的图（含可直接下载的 url） - POST /admin/internal/embeddings 回写向量并重算近邻；失败项计入 image_embedding_failures，连续失败到上限后退出队列（坏图不占位） - 近邻重算改为异步串行：同步会让 100 张的批次 POST 超时（实测 >60s） - 新增 key 分支单测：混用 key/url、细节图只给 key、字节数、回滚清单、向量占位
This commit is contained in:
toom1996
@ -52,6 +52,16 @@ type BrandRepository interface {
|
||||
// 规则：先按品牌总图片数取前 limit 个热门品牌，再为每个品牌取最新一篇走秀。
|
||||
// 返回的是仓储层裸数据（数字主键），对外编码由 service 层负责。
|
||||
PopularWithCover(ctx context.Context, limit int) ([]PopularBrandRow, error)
|
||||
|
||||
// StatsByIDs 批量统计若干品牌的「图集数 / 图片数」，供后台列表展示数据量。
|
||||
// 只查当前页的 id，避免为了两列统计把整张品牌表都聚合一遍。
|
||||
StatsByIDs(ctx context.Context, ids []uint32) (map[uint32]BrandStats, error)
|
||||
}
|
||||
|
||||
// BrandStats 品牌数据量统计（后台列表用）。
|
||||
type BrandStats struct {
|
||||
Runways int64 // 未删除的图集数
|
||||
Images int64 // 未删除的图片数
|
||||
}
|
||||
|
||||
// PopularBrandRow 热门品牌裸数据：品牌数字主键 + 中英文名称 + 代表走秀封面/中英文标题/文章 ID。
|
||||
@ -174,11 +184,30 @@ func (r *brandRepository) list(ctx context.Context, q dto.BrandQuery, restrictID
|
||||
return []model.Brand{}, 0, nil
|
||||
}
|
||||
|
||||
var items []model.Brand
|
||||
if err := r.db.WithContext(ctx).
|
||||
// 后台可按「最近有动静」排序：走秀图集的 updated_at 与抓取任务的 updated_at 取大。
|
||||
//
|
||||
// 性能注意：不要写成逐行的相关子查询（2.3 万品牌 × 2 个子查询实测 2.7 秒）。
|
||||
// 这里先把两张小表各自按品牌聚合（brand_runways / crawl_tasks 都只有几百行），
|
||||
// 再 LEFT JOIN 到品牌表排序 —— 一次哈希连接 + 一次排序，量级降到百毫秒。
|
||||
// 公开侧不传 Sort，保持原有 name_en ASC，前台顺序不变。
|
||||
itemsQ := r.db.WithContext(ctx).
|
||||
Model(&model.Brand{}).
|
||||
Scopes(scope).
|
||||
Order("name_en ASC").
|
||||
Scopes(scope)
|
||||
if admin && q.Sort == "updated" {
|
||||
itemsQ = itemsQ.
|
||||
Joins(`LEFT JOIN (SELECT brand_id, max(updated_at) AS u
|
||||
FROM brand_runways WHERE is_deleted = 0 GROUP BY brand_id) br
|
||||
ON br.brand_id = brands.id`).
|
||||
Joins(`LEFT JOIN (SELECT brand_id, max(updated_at) AS u
|
||||
FROM crawl_tasks GROUP BY brand_id) ct
|
||||
ON ct.brand_id = brands.id`).
|
||||
Order("GREATEST(coalesce(br.u, 0), coalesce(ct.u, 0)) DESC, name_en ASC")
|
||||
} else {
|
||||
itemsQ = itemsQ.Order("name_en ASC")
|
||||
}
|
||||
|
||||
var items []model.Brand
|
||||
if err := itemsQ.
|
||||
Offset(q.Offset()).
|
||||
Limit(q.Size).
|
||||
Find(&items).Error; err != nil {
|
||||
@ -319,6 +348,36 @@ func (r *brandRepository) CrawlBrands(ctx context.Context, brandID uint32) ([]mo
|
||||
// 3. 结果按热门品牌排名顺序返回
|
||||
//
|
||||
// 注意：created_at/updated_at 是数据导入时间，不能代表走秀新旧，因此必须用 year + season_code。
|
||||
// StatsByIDs 批量统计图集数 / 图片数（只针对传入的 id，通常是一页 30 个）。
|
||||
func (r *brandRepository) StatsByIDs(ctx context.Context, ids []uint32) (map[uint32]BrandStats, error) {
|
||||
out := make(map[uint32]BrandStats, len(ids))
|
||||
if len(ids) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
var rows []struct {
|
||||
BrandID uint32
|
||||
Runways int64
|
||||
Images int64
|
||||
}
|
||||
// 一次查询同时拿两个计数：按品牌聚合图集，再左连图片聚合，避免 N+1。
|
||||
err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT r.brand_id,
|
||||
count(DISTINCT r.id) AS runways,
|
||||
count(i.id) AS images
|
||||
FROM brand_runways r
|
||||
LEFT JOIN brand_runway_images i
|
||||
ON i.runway_id = r.id AND i.is_deleted = 0
|
||||
WHERE r.is_deleted = 0 AND r.brand_id IN ?
|
||||
GROUP BY r.brand_id`, ids).Scan(&rows).Error
|
||||
if err != nil {
|
||||
return out, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
out[row.BrandID] = BrandStats{Runways: row.Runways, Images: row.Images}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *brandRepository) PopularWithCover(ctx context.Context, limit int) ([]PopularBrandRow, error) {
|
||||
if limit < 1 {
|
||||
limit = 20
|
||||
|
||||
338
internal/repository/crawl_repository.go
Normal file
338
internal/repository/crawl_repository.go
Normal file
@ -0,0 +1,338 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"gorm.io/gorm"
|
||||
|
||||
"fashionapi/internal/model"
|
||||
)
|
||||
|
||||
// CrawlRepository 品牌抓取任务（crawl_tasks）的数据访问。
|
||||
//
|
||||
// 任务表同时被两个方向使用：
|
||||
// - 后台（cookie 会话）往里加任务 / 取消 / 重试；
|
||||
// - spider 服务（HMAC 验签）来领取 / 回报结果。
|
||||
//
|
||||
// 领取用 `FOR UPDATE SKIP LOCKED`，多个 spider 实例并发领取不会抢到同一条。
|
||||
type CrawlRepository interface {
|
||||
// Enqueue 把单个品牌加入队列。
|
||||
// 返回 true 表示新入队或把已完成/失败的任务重置为待领取；false 表示该品牌已在队列中/正在跑。
|
||||
Enqueue(ctx context.Context, brandID uint32, maxShows int, priority int) (bool, error)
|
||||
// EnqueueBulk 按「范围 + 前缀」批量入队，返回本次实际新入队（或重置）的数量。
|
||||
EnqueueBulk(ctx context.Context, f BulkFilter) (int, error)
|
||||
// Claim 领取一个待执行任务并置为 running；无任务时返回 nil, nil。
|
||||
Claim(ctx context.Context, worker string) (*model.CrawlTaskRow, error)
|
||||
// Report spider 回报结果。failed 且未达重试上限时自动回到 queued。
|
||||
Report(ctx context.Context, id uint32, status string, shows int, errMsg string) error
|
||||
// List 后台列表（按 running → queued → 其它 排序）。
|
||||
List(ctx context.Context, offset, limit int) ([]model.CrawlTaskRow, int64, error)
|
||||
// SetStatus 后台手动改状态（取消 / 重试）。
|
||||
SetStatus(ctx context.Context, id uint32, status string) error
|
||||
// Delete 删除任务记录（不影响已入库的数据）。
|
||||
Delete(ctx context.Context, id uint32) error
|
||||
// Counts 各状态任务数，供导航角标使用。
|
||||
Counts(ctx context.Context) (map[string]int64, error)
|
||||
// BrandCrawlStates 返回这些品牌在抓取队列里的「最近状态 + 最近时间」。
|
||||
// 品牌列表用它展示「抓取状态 / 上次抓取」；没有任务记录的品牌不会出现在 map 里。
|
||||
BrandCrawlStates(ctx context.Context, brandIDs []uint32) (map[uint32]BrandCrawlState, error)
|
||||
}
|
||||
|
||||
// BrandCrawlState 品牌最近一次抓取任务的状态快照。
|
||||
type BrandCrawlState struct {
|
||||
Status string // queued / running / done / failed / canceled
|
||||
UpdatedAt time.Time // 该任务最后一次更新时间
|
||||
}
|
||||
|
||||
type crawlRepository struct {
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
// NewCrawlRepository 创建爬取任务仓储。
|
||||
func NewCrawlRepository(db *gorm.DB) CrawlRepository {
|
||||
return &crawlRepository{db: db}
|
||||
}
|
||||
|
||||
const crawlColumns = `t.id, t.brand_id, t.status, t.priority, t.max_shows, t.worker,
|
||||
t.locked_at, t.attempts, t.shows_submitted, t.last_error, t.created_at, t.updated_at,
|
||||
coalesce(b.name_en, '') AS brand_name`
|
||||
|
||||
func (r *crawlRepository) Enqueue(ctx context.Context, brandID uint32, maxShows int, priority int) (bool, error) {
|
||||
now := uint32(time.Now().Unix())
|
||||
// INSERT ... SELECT 保证品牌不存在（或已软删）时不会凭空建任务；
|
||||
// ON CONFLICT 的 WHERE 让「已在排队/正在跑」的任务不被重置，避免打断正在执行的工作。
|
||||
res := r.db.WithContext(ctx).Exec(`
|
||||
INSERT INTO crawl_tasks (brand_id, status, priority, max_shows, created_at, updated_at)
|
||||
SELECT b.id, ?, ?, ?, ?, ?
|
||||
FROM brands b
|
||||
WHERE b.id = ? AND b.is_deleted = 0
|
||||
ON CONFLICT (brand_id) DO UPDATE
|
||||
SET status = 'queued',
|
||||
priority = EXCLUDED.priority,
|
||||
max_shows = EXCLUDED.max_shows,
|
||||
attempts = 0,
|
||||
locked_at = 0,
|
||||
worker = '',
|
||||
shows_submitted = 0,
|
||||
last_error = '',
|
||||
updated_at = EXCLUDED.updated_at
|
||||
WHERE crawl_tasks.status IN ('done', 'failed', 'canceled')`,
|
||||
model.CrawlQueued, priority, maxShows, now, now, brandID)
|
||||
if res.Error != nil {
|
||||
return false, res.Error
|
||||
}
|
||||
return res.RowsAffected > 0, nil
|
||||
}
|
||||
|
||||
// BulkScope 批量入队的品牌范围。
|
||||
type BulkScope string
|
||||
|
||||
const (
|
||||
// ScopeAll 所有品牌（首次全站导入用，量大）。
|
||||
ScopeAll BulkScope = "all"
|
||||
// ScopeWithData 只挑「已有图集」的品牌 —— 用于增量更新：只补最新场次。
|
||||
ScopeWithData BulkScope = "with-data"
|
||||
// ScopeWithoutData 只挑「还没有任何图集」的品牌 —— 用于首次补全。
|
||||
ScopeWithoutData BulkScope = "without-data"
|
||||
)
|
||||
|
||||
// BulkFilter 批量入队条件。
|
||||
type BulkFilter struct {
|
||||
Prefix string // 品牌英文名前缀，空 = 不限
|
||||
Scope BulkScope // 范围（见上方常量）
|
||||
MaxShows int // 0 = 不限场次（全量），>0 = 只抓最近 N 场
|
||||
Limit int // 本次最多入队多少个品牌（防止一次灌爆）
|
||||
Priority int16 // 越大越先被领取
|
||||
}
|
||||
|
||||
func (r *crawlRepository) EnqueueBulk(ctx context.Context, f BulkFilter) (int, error) {
|
||||
now := uint32(time.Now().Unix())
|
||||
if f.Limit < 1 || f.Limit > 5000 {
|
||||
f.Limit = 500
|
||||
}
|
||||
if f.Scope == "" {
|
||||
f.Scope = ScopeAll
|
||||
}
|
||||
|
||||
where := []string{"b.is_deleted = 0", "b.name_en <> ''"}
|
||||
// 注意：参数顺序必须与 SQL 里 ? 出现的顺序严格一致（INSERT 的 4 个 + 每个 WHERE 条件 + LIMIT）。
|
||||
args := []any{model.CrawlQueued, f.Priority, f.MaxShows, now, now}
|
||||
|
||||
if f.Prefix != "" {
|
||||
where = append(where, `b.name_en ILIKE ? ESCAPE '\'`)
|
||||
args = append(args, escapeLike(f.Prefix)+"%")
|
||||
}
|
||||
switch f.Scope {
|
||||
case ScopeWithData:
|
||||
where = append(where, "EXISTS (SELECT 1 FROM brand_runways r WHERE r.brand_id = b.id AND r.is_deleted = 0)")
|
||||
case ScopeWithoutData:
|
||||
where = append(where, "NOT EXISTS (SELECT 1 FROM brand_runways r WHERE r.brand_id = b.id AND r.is_deleted = 0)")
|
||||
}
|
||||
args = append(args, f.Limit)
|
||||
|
||||
sql := `
|
||||
INSERT INTO crawl_tasks (brand_id, status, priority, max_shows, created_at, updated_at)
|
||||
SELECT b.id, ?, ?, ?, ?, ?
|
||||
FROM brands b
|
||||
WHERE ` + strings.Join(where, " AND ") + `
|
||||
ORDER BY EXISTS (SELECT 1 FROM crawl_tasks t WHERE t.brand_id = b.id) ASC, b.id ASC
|
||||
LIMIT ?
|
||||
ON CONFLICT (brand_id) DO UPDATE
|
||||
SET status = 'queued',
|
||||
priority = EXCLUDED.priority,
|
||||
max_shows = EXCLUDED.max_shows,
|
||||
attempts = 0,
|
||||
locked_at = 0,
|
||||
worker = '',
|
||||
shows_submitted = 0,
|
||||
last_error = '',
|
||||
updated_at = EXCLUDED.updated_at
|
||||
WHERE crawl_tasks.status IN ('done', 'failed', 'canceled')`
|
||||
|
||||
res := r.db.WithContext(ctx).Exec(sql, args...)
|
||||
if res.Error != nil {
|
||||
return 0, res.Error
|
||||
}
|
||||
return int(res.RowsAffected), nil
|
||||
}
|
||||
|
||||
// escapeLike 转义 LIKE 的通配符，让用户输入按字面量匹配。
|
||||
func escapeLike(s string) string {
|
||||
out := make([]rune, 0, len(s))
|
||||
for _, c := range s {
|
||||
switch c {
|
||||
case '%', '_', '\\':
|
||||
out = append(out, '\\')
|
||||
}
|
||||
out = append(out, c)
|
||||
}
|
||||
return string(out)
|
||||
}
|
||||
|
||||
func (r *crawlRepository) Claim(ctx context.Context, worker string) (*model.CrawlTaskRow, error) {
|
||||
now := uint32(time.Now().Unix())
|
||||
|
||||
tx := r.db.WithContext(ctx).Begin()
|
||||
if tx.Error != nil {
|
||||
return nil, tx.Error
|
||||
}
|
||||
defer func() {
|
||||
if tx.Error != nil {
|
||||
tx.Rollback()
|
||||
}
|
||||
}()
|
||||
|
||||
// 回收租约过期的 running（worker 崩溃 / 断电），否则任务永远卡着。
|
||||
if now > model.CrawlStuckTimeoutSec {
|
||||
if err := tx.Exec(`
|
||||
UPDATE crawl_tasks
|
||||
SET status = 'queued', locked_at = 0, worker = '', updated_at = ?
|
||||
WHERE status = 'running' AND locked_at > 0 AND locked_at < ?`,
|
||||
now, now-model.CrawlStuckTimeoutSec).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
|
||||
var id uint32
|
||||
if err := tx.Raw(`
|
||||
SELECT id FROM crawl_tasks
|
||||
WHERE status = 'queued'
|
||||
ORDER BY priority DESC, id ASC
|
||||
LIMIT 1
|
||||
FOR UPDATE SKIP LOCKED`).Scan(&id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if id == 0 {
|
||||
if err := tx.Commit().Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if err := tx.Exec(`
|
||||
UPDATE crawl_tasks
|
||||
SET status = 'running', worker = ?, locked_at = ?, attempts = attempts + 1, updated_at = ?
|
||||
WHERE id = ?`, worker, now, now, id).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := tx.Commit().Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
var row model.CrawlTaskRow
|
||||
if err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT `+crawlColumns+`
|
||||
FROM crawl_tasks t LEFT JOIN brands b ON b.id = t.brand_id
|
||||
WHERE t.id = ?`, id).Scan(&row).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &row, nil
|
||||
}
|
||||
|
||||
func (r *crawlRepository) Report(ctx context.Context, id uint32, status string, shows int, errMsg string) error {
|
||||
now := uint32(time.Now().Unix())
|
||||
return r.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
|
||||
var cur model.CrawlTask
|
||||
if err := tx.Where("id = ?", id).First(&cur).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
next := status
|
||||
// 失败但还有重试额度 → 回到队列，由任意 worker 重试（网络抖动不必人工介入）。
|
||||
if status == model.CrawlFailed && cur.Attempts < model.CrawlMaxAttempts {
|
||||
next = model.CrawlQueued
|
||||
}
|
||||
return tx.Model(&model.CrawlTask{}).Where("id = ?", id).
|
||||
Updates(map[string]any{
|
||||
"status": next,
|
||||
"shows_submitted": shows,
|
||||
"last_error": errMsg,
|
||||
"locked_at": 0,
|
||||
"worker": "",
|
||||
"updated_at": now,
|
||||
}).Error
|
||||
})
|
||||
}
|
||||
|
||||
func (r *crawlRepository) List(ctx context.Context, offset, limit int) ([]model.CrawlTaskRow, int64, error) {
|
||||
if limit < 1 {
|
||||
limit = 30
|
||||
}
|
||||
var total int64
|
||||
if err := r.db.WithContext(ctx).Table("crawl_tasks").Count(&total).Error; err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
var rows []model.CrawlTaskRow
|
||||
err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT `+crawlColumns+`
|
||||
FROM crawl_tasks t LEFT JOIN brands b ON b.id = t.brand_id
|
||||
ORDER BY CASE t.status WHEN 'running' THEN 0 WHEN 'queued' THEN 1 ELSE 2 END,
|
||||
t.priority DESC, t.id DESC
|
||||
LIMIT ? OFFSET ?`, limit, offset).Scan(&rows).Error
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
return rows, total, nil
|
||||
}
|
||||
|
||||
func (r *crawlRepository) SetStatus(ctx context.Context, id uint32, status string) error {
|
||||
return r.db.WithContext(ctx).Model(&model.CrawlTask{}).Where("id = ?", id).
|
||||
Updates(map[string]any{
|
||||
"status": status,
|
||||
"locked_at": 0,
|
||||
"worker": "",
|
||||
"attempts": 0,
|
||||
"last_error": "",
|
||||
"updated_at": uint32(time.Now().Unix()),
|
||||
}).Error
|
||||
}
|
||||
|
||||
func (r *crawlRepository) Delete(ctx context.Context, id uint32) error {
|
||||
return r.db.WithContext(ctx).Exec(`DELETE FROM crawl_tasks WHERE id = ?`, id).Error
|
||||
}
|
||||
|
||||
func (r *crawlRepository) BrandCrawlStates(ctx context.Context, brandIDs []uint32) (map[uint32]BrandCrawlState, error) {
|
||||
out := make(map[uint32]BrandCrawlState, len(brandIDs))
|
||||
if len(brandIDs) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
var rows []struct {
|
||||
BrandID uint32
|
||||
Status string
|
||||
UpdatedAt int64
|
||||
}
|
||||
// DISTINCT ON 取每个品牌最近一条任务（按 updated_at 倒序），一次查询拿全状态与时间。
|
||||
if err := r.db.WithContext(ctx).Raw(`
|
||||
SELECT DISTINCT ON (brand_id) brand_id, status, updated_at
|
||||
FROM crawl_tasks
|
||||
WHERE brand_id IN ?
|
||||
ORDER BY brand_id, updated_at DESC`, brandIDs).Scan(&rows).Error; err != nil {
|
||||
return out, err
|
||||
}
|
||||
for _, row := range rows {
|
||||
st := BrandCrawlState{Status: row.Status}
|
||||
if row.UpdatedAt > 0 {
|
||||
st.UpdatedAt = time.Unix(row.UpdatedAt, 0)
|
||||
}
|
||||
out[row.BrandID] = st
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (r *crawlRepository) Counts(ctx context.Context) (map[string]int64, error) {
|
||||
var rows []struct {
|
||||
Status string
|
||||
N int64
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Raw(
|
||||
`SELECT status, count(*) AS n FROM crawl_tasks GROUP BY status`).Scan(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]int64, len(rows))
|
||||
for _, r := range rows {
|
||||
out[r.Status] = r.N
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
@ -58,7 +58,24 @@ type DupStat struct {
|
||||
//
|
||||
// 近邻结果走 image_similar_pairs 预计算表：入库时（或 dbtool similarbackfill 回填时）
|
||||
// 算一次存表，后台页面只读结果。原因见 ReplaceSimilarPairs 注释。
|
||||
// PendingEmbedding 是一张「缺向量」的图片，供回填 worker 拉取。
|
||||
// 带 Key 而不是完整 URL：worker 自己按存储基址拼（后端可能换了域名，不该把地址写死在队列里）。
|
||||
type PendingEmbedding struct {
|
||||
Kind string // runway | street
|
||||
ImageID uint32
|
||||
Key string // 存储 key（S4）或本地相对路径（/uploads/...）
|
||||
}
|
||||
|
||||
type EmbeddingRepository interface {
|
||||
// ListPendingEmbeddings 取一批「还活着、没有向量、失败次数未达上限」的图片，按 id 升序。
|
||||
//
|
||||
// 媒体拆分后这条查询是向量回填的唯一入口：媒体机直传只写 key，后端不碰图片字节，
|
||||
// 向量必须靠外部 worker 回传，而 worker 需要知道「还有哪些图缺向量」。
|
||||
ListPendingEmbeddings(ctx context.Context, kind string, limit, maxAttempts int) ([]PendingEmbedding, error)
|
||||
// CountPendingEmbeddings 统计剩余待回填数量（worker 据此决定是否继续拉）。
|
||||
CountPendingEmbeddings(ctx context.Context, kind string, maxAttempts int) (int, error)
|
||||
// RecordEmbeddingFailure 记一次失败（attempts+1）。达到 maxAttempts 后该图不再进队列。
|
||||
RecordEmbeddingFailure(ctx context.Context, kind string, imageID uint32, reason string) error
|
||||
// UpsertEmbedding 写入或更新一张图的向量（(kind, image_id) 唯一）。
|
||||
UpsertEmbedding(ctx context.Context, kind string, imageID uint32, vec []float32) error
|
||||
// BulkUpsertEmbedding 批量写入；vecs[i] 对应 ids[i]，vecs[i] 为 nil 则跳过该行。
|
||||
@ -175,6 +192,85 @@ func (r *embeddingRepository) UpsertEmbedding(ctx context.Context, kind string,
|
||||
return r.BulkUpsertEmbedding(ctx, kind, []uint32{imageID}, [][]float32{vec})
|
||||
}
|
||||
|
||||
// ListPendingEmbeddings 取缺失向量的图片。
|
||||
//
|
||||
// 三个 LEFT JOIN 各有明确目的：
|
||||
// - e 判「已有没有向量」；
|
||||
// - f 判「失败次数是否已到上限」——没有它，一张永远编不出来的坏图会让每轮回填都空转；
|
||||
// - i.is_deleted = 0 跳过软删图，避免为已删内容白算向量。
|
||||
//
|
||||
// 表名/列名来自 ownerOf 常量，kind 走参数绑定，无注入面。
|
||||
func (r *embeddingRepository) ListPendingEmbeddings(ctx context.Context, kind string, limit, maxAttempts int) ([]PendingEmbedding, error) {
|
||||
owner, ok := ownerOf(kind)
|
||||
if !ok {
|
||||
return nil, nil
|
||||
}
|
||||
if limit <= 0 {
|
||||
limit = 50
|
||||
}
|
||||
if maxAttempts <= 0 {
|
||||
maxAttempts = 5
|
||||
}
|
||||
var rows []struct {
|
||||
ImageID uint32 `gorm:"column:image_id"`
|
||||
Image string `gorm:"column:image"`
|
||||
}
|
||||
if err := r.db.WithContext(ctx).Raw(
|
||||
"SELECT i.id AS image_id, i.image AS image "+
|
||||
"FROM "+owner.table+" AS i "+
|
||||
"LEFT JOIN image_embeddings AS e ON e.kind = ? AND e.image_id = i.id "+
|
||||
"LEFT JOIN image_embedding_failures AS f ON f.kind = ? AND f.image_id = i.id "+
|
||||
"WHERE i.is_deleted = 0 AND i.image <> '' AND e.image_id IS NULL "+
|
||||
" AND coalesce(f.attempts, 0) < ? "+
|
||||
"ORDER BY i.id LIMIT ?",
|
||||
kind, kind, maxAttempts, limit).Scan(&rows).Error; err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make([]PendingEmbedding, 0, len(rows))
|
||||
for _, row := range rows {
|
||||
out = append(out, PendingEmbedding{Kind: kind, ImageID: row.ImageID, Key: row.Image})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// CountPendingEmbeddings 统计剩余量（与 ListPendingEmbeddings 同口径，避免两边不一致）。
|
||||
func (r *embeddingRepository) CountPendingEmbeddings(ctx context.Context, kind string, maxAttempts int) (int, error) {
|
||||
owner, ok := ownerOf(kind)
|
||||
if !ok {
|
||||
return 0, nil
|
||||
}
|
||||
if maxAttempts <= 0 {
|
||||
maxAttempts = 5
|
||||
}
|
||||
var n int
|
||||
if err := r.db.WithContext(ctx).Raw(
|
||||
"SELECT count(*) FROM "+owner.table+" AS i "+
|
||||
"LEFT JOIN image_embeddings AS e ON e.kind = ? AND e.image_id = i.id "+
|
||||
"LEFT JOIN image_embedding_failures AS f ON f.kind = ? AND f.image_id = i.id "+
|
||||
"WHERE i.is_deleted = 0 AND i.image <> '' AND e.image_id IS NULL "+
|
||||
" AND coalesce(f.attempts, 0) < ?",
|
||||
kind, kind, maxAttempts).Scan(&n).Error; err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
|
||||
// RecordEmbeddingFailure 累加失败次数。
|
||||
// 用 upsert 而不是先查后写：worker 可能并发回传，两步写法会丢计数。
|
||||
func (r *embeddingRepository) RecordEmbeddingFailure(ctx context.Context, kind string, imageID uint32, reason string) error {
|
||||
if len(reason) > 300 {
|
||||
reason = reason[:300]
|
||||
}
|
||||
return r.db.WithContext(ctx).Exec(`
|
||||
INSERT INTO image_embedding_failures (kind, image_id, attempts, last_error, updated_at)
|
||||
VALUES (?, ?, 1, ?, ?)
|
||||
ON CONFLICT (kind, image_id) DO UPDATE
|
||||
SET attempts = image_embedding_failures.attempts + 1,
|
||||
last_error = EXCLUDED.last_error,
|
||||
updated_at = EXCLUDED.updated_at`,
|
||||
kind, imageID, reason, uint32(time.Now().Unix())).Error
|
||||
}
|
||||
|
||||
// BulkUpsertEmbedding 逐行写入而非 unnest 批量：日入库仅几百张，量级极小，
|
||||
// 换来的是可读性与低出错概率。ON CONFLICT 保证同一图重复入库不会写出多行。
|
||||
func (r *embeddingRepository) BulkUpsertEmbedding(ctx context.Context, kind string, ids []uint32, vecs [][]float32) error {
|
||||
|
||||
Reference in New Issue
Block a user