You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
kmftools/bookscraper/scraper/abort.py

125 lines
3.3 KiB

import os
import redis
from scraper.logger_decorators import logcall
from scraper.ui_log import push_ui
# ---------------------------------------------------------
# Default Redis connection (Docker workers)
# ---------------------------------------------------------
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
r = redis.Redis.from_url(REDIS_URL, decode_responses=True)
# Debug mode (optional)
ABORT_DEBUG = os.getenv("ABORT_DEBUG", "1") == "1"
# Avoid duplicate spam
_seen_debug_keys = set()
# =========================================================
# INTERNAL DEBUGGING
# =========================================================
def _debug(msg: str):
print(msg)
push_ui(msg)
# =========================================================
# ABORT FLAG — unified book_idx
# =========================================================
def set_abort(book_idx: str):
"""Enable abort mode for book_idx."""
key = f"abort:{book_idx}"
r.set(key, "1")
if ABORT_DEBUG:
_debug(f"[ABORT] SET {key}")
def clear_abort(book_idx: str):
"""Clear abort flag."""
key = f"abort:{book_idx}"
r.delete(key)
if ABORT_DEBUG:
_debug(f"[ABORT] CLEAR {key}")
def abort_requested(book_idx: str, redis_client=None) -> bool:
"""
Check whether abort flag is active for book_idx.
redis_client:
- Docker workers → None → use default Redis (r)
- Local macOS audio worker → passes Redis(host=127.0.0.1)
"""
client = redis_client or r
key = f"abort:{book_idx}"
try:
exists = client.exists(key)
if ABORT_DEBUG:
# Log only once per book
if key not in _seen_debug_keys:
try:
conn = client.connection_pool.connection_kwargs
host = conn.get("host")
port = conn.get("port")
db = conn.get("db")
_debug(
# f"[ABORT_DEBUG] first check book_idx={book_idx} "
f"redis={host}:{port} db={db}"
)
except Exception:
_debug(f"[ABORT_DEBUG] first check book_idx={book_idx}")
_seen_debug_keys.add(key)
# Log ACTIVE state
if exists == 1:
_debug(f"[ABORT] ACTIVE for {book_idx}")
return exists == 1
except Exception as e:
if ABORT_DEBUG:
_debug(f"[ABORT_DEBUG] ERROR checking {key}: {e}")
return False
# =========================================================
# PER-CHAPTER STATE — unified book_idx
# =========================================================
def mark_chapter_started(book_idx: str, chapter_num: int):
key = f"started:{book_idx}:{chapter_num}"
r.set(key, "1")
def chapter_started(book_idx: str, chapter_num: int) -> bool:
key = f"started:{book_idx}:{chapter_num}"
return r.exists(key) == 1
# =========================================================
# RESET STATE FOR BOOK_IDX
# =========================================================
def reset_book_state(book_idx: str):
"""
Remove abort flag and all per-chapter started markers.
"""
# abort flag
r.delete(f"abort:{book_idx}")
# chapter markers
pattern = f"started:{book_idx}:*"
for k in r.scan_iter(pattern):
r.delete(k)