"""Discovery: хранилище задач поиска каналов, кандидатов, чёрного списка и лога. Единый слой доступа к disc_tasks / disc_candidates / disc_blacklist / disc_log (Task 1). Потребляется и API (Task 7), и фоновым воркером (Task 6). Соглашения модуля: - все обращения к БД только через `store.*` с параметрами (без конкатенации SQL); - время — миллисекунды (`time.time_ns() // 1_000_000`); - JSON-поля (keywords/marks/topics) в БД — VARCHAR, наружу всегда список (`json.dumps(..., ensure_ascii=False)` при записи, `json.loads` при чтении); - наружные dict-ы — camelCase (конвенция границы API проекта), поля совпадают с колонками БД: keywords/minSubscribers/sampleSize/planJoins/autoJoin, searchIdx/searchDone, fitRatio/autoJoined, createdAt/updatedAt, dialogId... - create/patch принимают ключи и в camelCase, и в snake_case (поле "min_subscribers" и т.п.) — на входе идёт нормализация к колонкам БД; - статусы кандидата: new -> review -> joined|rejected. Переводы в joined/rejected делаются ТОЛЬКО через mark_joined()/mark_rejected() (счётчики, чёрный список, лог); set_candidate_status() разрешает new/review. Бюджет авто-вступлений: сумма plan_joins задач со статусом NOT IN ('done','failed') плюс plan_joins новой/увеличиваемой задачи не должна превышать discJoinLimit. """ from __future__ import annotations import json import time from ..db import store # префиксы id (конвенция: видно, из какой таблицы запись) _ID_TASK = "dt_" _ID_LOG = "dl_" # статусы кандидата, при которых повторное добавление источника запрещено _ACTIVE_CANDIDATE = {"new", "review", "joined"} # статусы задачи, которые НЕ занимают бюджет plan_joins _DONE_TASK = {"done", "failed"} # алиасы полей задачи: колонка БД -> набор имён в payload/патче _TASK_ALIASES = { "min_subscribers": {"minSubscribers", "min_subscribers"}, "sample_size": {"sampleSize", "sample_size"}, "plan_joins": {"planJoins", "plan_joins"}, "auto_join": {"autoJoin", "auto_join"}, # остальные поля называются одинаково: name, description, keywords, lang, threshold } # поля кандидата, которые можно менять через set_candidate() _CANDIDATE_FIELDS = { "name": "name", "username": "username", "kind": "kind", "hue": "hue", "participants": "participants", "lang_ru": "lang_ru", "langRu": "lang_ru", "marks": "marks", "topics": "topics", "fit_ratio": "fit_ratio", "fitRatio": "fit_ratio", "auto_joined": "auto_joined", "autoJoined": "auto_joined", } def _now() -> int: return time.time_ns() // 1_000_000 def _json(value) -> str: return json.dumps(value or [], ensure_ascii=False) def _loads(raw, default: list | None = None) -> list: try: return json.loads(raw or "[]") except (TypeError, ValueError): return list(default or []) def _plan_limit() -> int: """Верхняя граница plan_joins: текущий суточный лимит discJoinLimit.""" return max(1, int(store.get_setting("discJoinLimit") or 0)) def _active_plan_sum(exclude_id: str | None = None) -> int: sql = "SELECT coalesce(sum(plan_joins), 0) FROM disc_tasks WHERE status NOT IN ('done', 'failed')" params: list = [] if exclude_id: sql += " AND id <> ?" params.append(exclude_id) return int(store.scalar(sql, params) or 0) def _assert_plan(plan_joins: int) -> None: """plan_joins >= 1 и не больше суточного лимита (настраивается в UI).""" if plan_joins < 1: raise ValueError("plan_joins должен быть не меньше 1") limit = _plan_limit() if plan_joins > limit: raise ValueError(f"plan_joins {plan_joins} больше суточного лимита авто-вступлений ({limit})") def _assert_budget(plan_joins: int, exclude_id: str | None = None) -> None: """Правило бюджета: сумма планов активных задач + новая <= discJoinLimit.""" limit = _plan_limit() used = _active_plan_sum(exclude_id) if used + plan_joins > limit: raise ValueError( f"Бюджет авто-вступлений исчерпан: задачи уже занимают {used} из {limit} в сутки, " f"ещё {plan_joins} не влезает" ) # ─── view-слои (наружу camelCase) ────────────────────────────────────────── def _task_view(row: dict) -> dict: return { "id": row["id"], "name": row["name"], "description": row["description"], "keywords": _loads(row["keywords"]), "minSubscribers": int(row["min_subscribers"]), "lang": row["lang"], "threshold": int(row["threshold"]), "sampleSize": int(row["sample_size"]), "planJoins": int(row["plan_joins"]), "autoJoin": bool(row["auto_join"]), "status": row["status"], "searchIdx": int(row["search_idx"]), "searchDone": bool(row["search_done"]), "found": int(row["found"]), "evaluated": int(row["evaluated"]), "joined": int(row["joined"]), "rejected": int(row["rejected"]), "createdAt": row["created_at"], "updatedAt": row["updated_at"], } def _candidate_view(row: dict) -> dict: return { "dialogId": row["dialog_id"], "taskId": row["task_id"], "name": row["name"], "username": row["username"], "kind": row["kind"], "hue": row["hue"], "participants": row["participants"], "langRu": row["lang_ru"], "marks": _loads(row["marks"]), "topics": _loads(row["topics"]), "fitRatio": row["fit_ratio"], "status": row["status"], "autoJoined": bool(row["auto_joined"]), "joinFailures": int(row.get("join_failures") or 0), "createdAt": row["created_at"], "updatedAt": row["updated_at"], } def _blacklist_view(row: dict) -> dict: return { "dialogId": row["dialog_id"], "name": row["name"], "reason": row["reason"], "createdAt": row["created_at"], } def _log_view(row: dict) -> dict: return { "id": row["id"], "taskId": row["task_id"], "event": row["event"], "text": row["text"], "createdAt": row["created_at"], } # ─── задачи ──────────────────────────────────────────────────────────────── def _task_or_raise(task_id: str) -> dict: row = store.query_one("SELECT * FROM disc_tasks WHERE id = ?", [task_id]) if row is None: raise KeyError(task_id) return row def list_tasks() -> list[dict]: """Все задачи, старые первыми (воркер берёт самую старую running).""" return [_task_view(r) for r in store.query("SELECT * FROM disc_tasks ORDER BY created_at ASC")] def get_task(task_id: str) -> dict | None: row = store.query_one("SELECT * FROM disc_tasks WHERE id = ?", [task_id]) return _task_view(row) if row else None def _norm_task_values(patch: dict) -> dict: """Нормализация payload/патча задачи (camel/snake алиасы) к колонкам БД.""" out: dict = {} for key, value in patch.items(): if key in ("name", "description", "keywords", "lang", "threshold"): col = key else: col = next((c for c, aliases in _TASK_ALIASES.items() if key in aliases), None) if col is None: continue # неизвестное поле игнорируем out[col] = value return out def _validate_task_values(cols: dict) -> None: """Проверка границ значений задачи (после нормализации).""" if "threshold" in cols: cols["threshold"] = max(1, min(100, int(cols["threshold"]))) if "sample_size" in cols: cols["sample_size"] = max(1, int(cols["sample_size"])) if "min_subscribers" in cols: cols["min_subscribers"] = max(0, int(cols["min_subscribers"])) if "lang" in cols and cols["lang"] not in ("ru", "any"): cols["lang"] = "ru" if "auto_join" in cols: cols["auto_join"] = bool(cols["auto_join"]) if "keywords" in cols: keywords = cols["keywords"] if isinstance(cols["keywords"], list) else [cols["keywords"]] cols["keywords"] = [str(k).strip() for k in keywords if str(k).strip()] if "description" in cols: cols["description"] = str(cols["description"] or "") if "name" in cols: cols["name"] = str(cols["name"] or "").strip() def create_task(payload: dict) -> dict: """Создать задачу поиска. Валидация: name непустое; plan_joins 1..discJoinLimit; правило бюджета (сумма plan_joins активных задач + новая <= discJoinLimit) — иначе ValueError. """ cols = _norm_task_values(payload) name = str(cols.get("name") or "").strip() if not name: raise ValueError("Укажите название задачи") plan_joins = int(cols.get("plan_joins", 1)) _assert_plan(plan_joins) _assert_budget(plan_joins) values = { "name": name, "description": str(cols.get("description") or ""), "keywords": cols.get("keywords", []), "min_subscribers": max(0, int(cols.get("min_subscribers", 0))), "lang": cols.get("lang", "ru"), "threshold": max(1, min(100, int(cols.get("threshold", int(store.get_setting("discEvalThreshold") or 40))))), "sample_size": max(1, int(cols.get("sample_size", int(store.get_setting("discEvalSample") or 10)))), "plan_joins": plan_joins, "auto_join": bool(cols.get("auto_join", False)), } _validate_task_values(values) task_id = store.uid(_ID_TASK) now = _now() store.execute( "INSERT INTO disc_tasks(id, name, description, keywords, min_subscribers, lang, threshold, " "sample_size, plan_joins, auto_join, status, search_idx, search_done, found, evaluated, " "joined, rejected, created_at, updated_at) " "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 'draft', 0, FALSE, 0, 0, 0, 0, ?, ?)", [ task_id, values["name"], values["description"], _json(values["keywords"]), values["min_subscribers"], values["lang"], values["threshold"], values["sample_size"], values["plan_joins"], values["auto_join"], now, now, ], ) return get_task(task_id) # type: ignore[return-value] def patch_task(task_id: str, patch: dict) -> dict: """Обновить поля задачи. Увеличение plan_joins — с проверкой бюджета.""" row = _task_or_raise(task_id) cols = _norm_task_values(patch) if not cols: return get_task(task_id) # type: ignore[return-value] old_plan = int(row["plan_joins"]) if "plan_joins" in cols: new_plan = int(cols["plan_joins"]) _assert_plan(new_plan) if new_plan > old_plan: _assert_budget(new_plan, exclude_id=task_id) cols["plan_joins"] = new_plan _validate_task_values(cols) sets = ["updated_at = ?"] params: list = [_now()] for col, value in cols.items(): sets.append(f"{col} = ?") params.append(_json(value) if col == "keywords" else value) params.append(task_id) store.execute( f"UPDATE disc_tasks SET {', '.join(sets)} WHERE id = ?", params, ) return get_task(task_id) # type: ignore[return-value] def delete_task(task_id: str) -> None: """Удалить задачу вместе с её кандидатами и логом (чёрный список общий).""" store.execute("DELETE FROM disc_tasks WHERE id = ?", [task_id]) store.execute("DELETE FROM disc_candidates WHERE task_id = ?", [task_id]) store.execute("DELETE FROM disc_log WHERE task_id = ?", [task_id]) def start_task(task_id: str) -> dict: """Запустить поиск: keywords непустые; status=running. При повторном запуске завершённой/упавшей задачи прогресс поиска обнуляется (свежий проход по ключам); при продолжении из paused — сохраняется. """ row = _task_or_raise(task_id) keywords = _loads(row["keywords"]) if not keywords: raise ValueError("Нет ключевых слов для поиска — добавьте их в задачу") now = _now() reset = row["status"] in _DONE_TASK if reset: # повторный прогон завершённой/упавшей задачи — свежий проход по ключам store.execute( "UPDATE disc_tasks SET status = 'running', search_idx = 0, search_done = FALSE, " "found = 0, evaluated = 0, joined = 0, rejected = 0, updated_at = ? WHERE id = ?", [now, task_id], ) else: # старт из draft или продолжение из paused — прогресс поиска сохраняется store.execute( "UPDATE disc_tasks SET status = 'running', updated_at = ? WHERE id = ?", [now, task_id], ) return get_task(task_id) # type: ignore[return-value] def pause_task(task_id: str) -> dict: """Поставить задачу на паузу.""" _task_or_raise(task_id) store.execute( "UPDATE disc_tasks SET status = 'paused', updated_at = ? WHERE id = ?", [_now(), task_id], ) return get_task(task_id) # type: ignore[return-value] def bump_counter(task_id: str, field: str, n: int = 1) -> None: """Увеличить счётчик задачи: found|evaluated|joined|rejected.""" if field not in ("found", "evaluated", "joined", "rejected"): raise ValueError(f"Неизвестный счётчик задачи: {field}") row = _task_or_raise(task_id) row[field] = int(row[field]) + max(0, int(n)) store.execute( f"UPDATE disc_tasks SET {field} = ?, updated_at = ? WHERE id = ?", [row[field], _now(), task_id], ) def advance_search(task_id: str) -> None: """Перейти к следующему ключу; когда search_idx >= len(keywords) — search_done=True.""" row = _task_or_raise(task_id) keywords = _loads(row["keywords"]) new_idx = int(row["search_idx"]) + 1 done = new_idx >= len(keywords) store.execute( "UPDATE disc_tasks SET search_idx = ?, search_done = ?, updated_at = ? WHERE id = ?", [new_idx, done, _now(), task_id], ) # ─── кандидаты ───────────────────────────────────────────────────────────── def list_candidates(task_id: str, status: str | None = None) -> list[dict]: """Кандидаты задачи (marks/topics уже списки); status — фильтр.""" if status: rows = store.query( "SELECT * FROM disc_candidates WHERE task_id = ? AND status = ? ORDER BY created_at ASC", [task_id, status], ) else: rows = store.query( "SELECT * FROM disc_candidates WHERE task_id = ? ORDER BY created_at ASC", [task_id], ) return [_candidate_view(r) for r in rows] def _get_candidate(dialog_id: str) -> dict | None: row = store.query_one("SELECT * FROM disc_candidates WHERE dialog_id = ?", [dialog_id]) return row def _skip(task_id: str, reason: str) -> None: add_log(task_id, "skip", reason) def add_candidate(task_id: str, dialog_id: str, name: str, username: str, kind: str, hue: str) -> dict | None: """Добавить найденный источник как кандидата задачи. None (с логом skip), если источник уже мониторится (есть в dialogs), в чёрном списке или уже добавлен в статусе new/review/joined. Прежняя запись со статусом rejected (например, после remove_blacklist) заменяется новой. """ _task_or_raise(task_id) dialog_id = str(dialog_id) if store.scalar("SELECT 1 FROM dialogs WHERE id = ? LIMIT 1", [dialog_id]): _skip(task_id, f"пропущен {dialog_id}: источник уже мониторится (мы состоим)") return None if store.scalar("SELECT 1 FROM disc_blacklist WHERE dialog_id = ? LIMIT 1", [dialog_id]): _skip(task_id, f"пропущен {dialog_id}: источник в чёрном списке") return None existing = store.query_one( "SELECT status FROM disc_candidates WHERE dialog_id = ?", [dialog_id], ) if existing and existing["status"] in _ACTIVE_CANDIDATE: _skip(task_id, f"пропущен {dialog_id}: кандидат уже есть (статус {existing['status']})") return None if existing: # устаревшая rejected-запись: перезаписываем как новый кандидат store.execute("DELETE FROM disc_candidates WHERE dialog_id = ?", [dialog_id]) now = _now() store.execute( "INSERT INTO disc_candidates(dialog_id, task_id, name, username, kind, hue, participants, " "lang_ru, marks, topics, fit_ratio, status, auto_joined, created_at, updated_at) " "VALUES (?, ?, ?, ?, ?, ?, NULL, NULL, '[]', '[]', NULL, 'new', FALSE, ?, ?)", [ dialog_id, task_id, str(name or "").strip() or dialog_id, str(username or "").strip(), str(kind or "channel"), str(hue or "#666"), now, now, ], ) bump_counter(task_id, "found") row = _get_candidate(dialog_id) return _candidate_view(row) if row else None def set_candidate(task_id: str, dialog_id: str, patch: dict) -> dict: """Обновить поля кандидата задачи (например, по результатам оценки). patch — значения в нотации кандидата (camelCase или snake_case): participants, kind, langRu/lang_ru, marks, topics, fitRatio/fit_ratio, autoJoined/auto_joined, name, username, hue. """ _task_or_raise(task_id) row = _get_candidate(dialog_id) if row is None or row["task_id"] != task_id: raise KeyError(dialog_id) cols: dict = {} for key, value in patch.items(): col = _CANDIDATE_FIELDS.get(key) if col is None: continue cols[col] = value if not cols: return _candidate_view(row) if "marks" in cols: cols["marks"] = _json([str(m) for m in cols["marks"]]) if "topics" in cols: cols["topics"] = _json(cols["topics"]) for col in ("name", "username", "kind", "hue"): if col in cols: cols[col] = str(cols[col] or "").strip() or row[col] if "participants" in cols and cols["participants"] is not None: cols["participants"] = int(cols["participants"]) sets = ["updated_at = ?"] params: list = [_now()] for col, value in cols.items(): sets.append(f"{col} = ?") params.append(value) params.append(dialog_id) store.execute(f"UPDATE disc_candidates SET {', '.join(sets)} WHERE dialog_id = ?", params) updated = _get_candidate(dialog_id) return _candidate_view(updated) if updated else _candidate_view(row) def set_candidate_status(dialog_id: str, status: str) -> dict: """Перевести кандидата в new/review (+ лог review). joined/rejected меняются только через mark_joined()/mark_rejected() — там счётчики задачи, чёрный список и лог join/reject. """ if status not in ("new", "review"): raise ValueError(f"Статус {status} выставляется через mark_joined/mark_rejected") row = _get_candidate(dialog_id) if row is None: raise KeyError(dialog_id) store.execute( "UPDATE disc_candidates SET status = ?, updated_at = ? WHERE dialog_id = ?", [status, _now(), dialog_id], ) if status == "review": add_log(row["task_id"], "review", f"кандидат {dialog_id} переведён в review") updated = _get_candidate(dialog_id) return _candidate_view(updated) if updated else _candidate_view(row) def delete_candidate(dialog_id: str) -> None: """Удалить кандидата (skip-ветки воркера; повторный вызов безопасен).""" store.execute("DELETE FROM disc_candidates WHERE dialog_id = ?", [dialog_id]) def mark_joined(dialog_id: str, auto: bool) -> dict: """Источник вступил: status=joined, счётчик joined задачи, лог join_auto/join_manual.""" row = _get_candidate(dialog_id) if row is None: raise KeyError(dialog_id) if row["status"] == "joined": return _candidate_view(row) # идемпотентно: повторно не считаем now = _now() store.execute( "UPDATE disc_candidates SET status = 'joined', auto_joined = ?, updated_at = ? WHERE dialog_id = ?", [bool(auto), now, dialog_id], ) bump_counter(row["task_id"], "joined") add_log(row["task_id"], "join_auto" if auto else "join_manual", f"вступили в {dialog_id}") updated = _get_candidate(dialog_id) return _candidate_view(updated) if updated else _candidate_view(row) def mark_rejected(dialog_id: str, reason: str = "") -> dict: """Отклонить кандидата: status=rejected, счётчик rejected, лог reject, чёрный список. Повторный вызов для уже отклонённого — идемпотентен: счётчик/лог/чёрный список не трогаются (кандидата мог отклонить и человек, и воркер). """ row = _get_candidate(dialog_id) if row is None: raise KeyError(dialog_id) if row["status"] == "joined": raise ValueError("Нельзя отклонить источник, в который уже вступили") if row["status"] == "rejected": return _candidate_view(row) # идемпотентно: повторно не считаем и не логируем now = _now() store.execute( "UPDATE disc_candidates SET status = 'rejected', updated_at = ? WHERE dialog_id = ?", [now, dialog_id], ) bump_counter(row["task_id"], "rejected") add_log(row["task_id"], "reject", reason or f"отклонён {dialog_id}") add_blacklist(dialog_id, row["name"], reason or "") updated = _get_candidate(dialog_id) return _candidate_view(updated) if updated else _candidate_view(row) # ─── чёрный список ───────────────────────────────────────────────────────── def add_blacklist(dialog_id: str, name: str = "", reason: str = "") -> dict: """Пометить источник в чёрном списке (существующая запись обновляется).""" dialog_id = str(dialog_id) now = _now() store.execute( "INSERT INTO disc_blacklist(dialog_id, name, reason, created_at) VALUES (?, ?, ?, ?) " "ON CONFLICT(dialog_id) DO UPDATE SET name = excluded.name, reason = excluded.reason", [dialog_id, str(name or "").strip() or dialog_id, str(reason or ""), now], ) row = store.query_one("SELECT * FROM disc_blacklist WHERE dialog_id = ?", [dialog_id]) return _blacklist_view(row) if row else {"dialogId": dialog_id, "name": name, "reason": reason, "createdAt": now} def remove_blacklist(dialog_id: str) -> None: store.execute("DELETE FROM disc_blacklist WHERE dialog_id = ?", [dialog_id]) def list_blacklist() -> list[dict]: return [ _blacklist_view(r) for r in store.query("SELECT * FROM disc_blacklist ORDER BY created_at DESC") ] # ─── лог ─────────────────────────────────────────────────────────────────── def add_log(task_id: str, event: str, text: str = "") -> None: store.execute( "INSERT INTO disc_log(id, task_id, event, text, created_at) VALUES (?, ?, ?, ?, ?)", [store.uid(_ID_LOG), task_id, str(event), str(text or ""), _now()], ) def task_log(task_id: str, limit: int = 100) -> list[dict]: """Последние события задачи, новые сверху.""" limit = max(1, min(500, int(limit))) rows = store.query( "SELECT * FROM disc_log WHERE task_id = ? ORDER BY created_at DESC LIMIT ?", [task_id, limit], ) return [_log_view(r) for r in rows]