"""需求搜索队列 MySQL 读写 + 视频搜索词 PG 查询 —— 仅负责 DB I/O,不含业务逻辑""" from typing import Dict, List from src.infra.database.mysql.manager import MysqlManager from src.infra.database.postgresql.manager import PgManager from ._const import DemandConst class DemandQueueMapper(DemandConst): """demand_search_queue 表的数据访问层 职责: SQL 拼接 + 参数绑定 + 返回原生 dict。 不负责: 业务校验、格式转换(由上层 service 负责)。 """ def __init__(self, pool: MysqlManager): self.pool = pool # ═══════════════════════════════════════════════════════════════ # 写入 # ═══════════════════════════════════════════════════════════════ async def insert_batch(self, items: List[Dict]) -> int: """批量写入队列,返回受影响行数 items 每条须包含: account, search_key, dt, video_id, channel_name, key_type, experiment_id, drive_type """ if not items: return 0 query = f""" INSERT IGNORE INTO {self.QUEUE_TABLE} (account, search_key, dt, video_id, channel_name, category, key_type, experiment_id, status, drive_type) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """ params = [ ( it["account"], it["search_key"], it["dt"], it["video_id"], it["channel_name"], it.get("category", ""), it["key_type"], it["experiment_id"], self.QueueStatus.INIT, it["drive_type"], ) for it in items ] return await self.pool.save(query=query, params=params, batch=True) # ═══════════════════════════════════════════════════════════════ # 读取(供下游 consumer 使用) # ═══════════════════════════════════════════════════════════════ async def fetch_pending( self, dt: str, *, limit: int = 500, offset: int = 0, ) -> List[Dict]: """查询待搜索的队列项(按 dt)""" query = f""" SELECT * FROM {self.QUEUE_TABLE} WHERE dt = %s AND status = %s ORDER BY id ASC LIMIT %s OFFSET %s """ return await self.pool.fetch( query=query, params=(dt, self.QueueStatus.INIT, limit, offset), ) async def count_by_status(self, dt: str, status: int) -> int: """统计指定状态的队列项数量""" query = f""" SELECT COUNT(*) AS cnt FROM {self.QUEUE_TABLE} WHERE dt = %s AND status = %s """ row = await self.pool.fetch_one(query=query, params=(dt, status)) return row["cnt"] if row else 0 # ═══════════════════════════════════════════════════════════════ # 状态更新 # ═══════════════════════════════════════════════════════════════ async def mark_processing(self, item_id: int) -> int: query = f""" UPDATE {self.QUEUE_TABLE} SET status = %s WHERE id = %s AND status = %s """ return await self.pool.save( query=query, params=(self.QueueStatus.PROCESSING, item_id, self.QueueStatus.INIT), ) async def mark_processing_batch(self, ids: List[int]) -> int: """批量标记为处理中""" if not ids: return 0 placeholders = ",".join(["%s"] * len(ids)) query = f""" UPDATE {self.QUEUE_TABLE} SET status = %s WHERE id IN ({placeholders}) """ return await self.pool.save( query=query, params=(self.QueueStatus.PROCESSING, *ids), ) async def mark_done(self, item_id: int) -> int: query = f""" UPDATE {self.QUEUE_TABLE} SET status = %s WHERE id = %s """ return await self.pool.save( query=query, params=(self.QueueStatus.SUCCESS, item_id), ) async def mark_failed(self, item_id: int, reason: str = "") -> int: query = f""" UPDATE {self.QUEUE_TABLE} SET status = %s, fail_reason = %s WHERE id = %s """ return await self.pool.save( query=query, params=(self.QueueStatus.FAIL, reason[:512], item_id), ) class VideoDeconstructMapper: """video_vectors 表的数据访问层(PostgreSQL) 职责: 按 video_id 查询搜索词,每行 (config_code, text) 为一个搜索词。 config_code 直接对应 DemandConst.KeyType。 """ COLUMNS = ("video_id", "config_code", "text") def __init__(self, pg: PgManager): self.pg = pg async def find_by_video_ids(self, video_ids: list[int]) -> dict[int, list[Dict]]: """批量查询 video_vectors,一次网络往返 Returns: {video_id: [{"config_code": "VIDEO_KEYPOINT", "text": "xx"}, ...], ...} 不存在的 video_id 不会出现在返回 key 中。 """ if not video_ids: return {} cols = ", ".join(self.COLUMNS) rows = await self.pg.fetch( f"SELECT {cols} FROM video_vectors WHERE video_id = ANY($1)", params=(video_ids,), ) result: dict[int, list[Dict]] = {} for row in rows: vid = int(row["video_id"]) config_code = (row.get("config_code") or "").strip() text_val = (row.get("text") or "").strip() if config_code and text_val: if vid not in result: result[vid] = [] result[vid].append({"config_code": config_code, "text": text_val}) return result __all__ = ["DemandQueueMapper", "VideoDeconstructMapper"]