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.
83 lines
2.2 KiB
83 lines
2.2 KiB
import os
|
|
import redis
|
|
|
|
# ---------------------------------------------------------
|
|
# Redis connection
|
|
# ---------------------------------------------------------
|
|
REDIS_URL = os.getenv("REDIS_URL", "redis://redis:6379/0")
|
|
r = redis.Redis.from_url(REDIS_URL, decode_responses=True)
|
|
|
|
|
|
# =========================================================
|
|
# ABORT FLAG
|
|
# =========================================================
|
|
|
|
|
|
def set_abort(book_id: str):
|
|
"""
|
|
Enable abort mode for this book.
|
|
All download tasks that haven't started yet will immediately exit.
|
|
"""
|
|
r.set(f"abort:{book_id}", "1")
|
|
|
|
|
|
def clear_abort(book_id: str):
|
|
"""
|
|
Clear abort flag so future runs are unaffected.
|
|
"""
|
|
r.delete(f"abort:{book_id}")
|
|
|
|
|
|
def abort_requested(book_id: str) -> bool:
|
|
"""
|
|
True if abort flag is set for this book.
|
|
"""
|
|
return r.exists(f"abort:{book_id}") == 1
|
|
|
|
|
|
# =========================================================
|
|
# PER-CHAPTER STATE
|
|
# =========================================================
|
|
# We mark a chapter "started" once its download task begins.
|
|
# If abort is activated AFTER download start:
|
|
# → download must complete
|
|
# → parse must complete
|
|
# → save must complete
|
|
# All subsequent chapters will skip.
|
|
|
|
|
|
def mark_chapter_started(book_id: str, chapter_num: int):
|
|
"""
|
|
Mark this chapter as started. Parse/save will always run after this,
|
|
even if abort has been activated afterwards.
|
|
"""
|
|
key = f"started:{book_id}:{chapter_num}"
|
|
r.set(key, "1")
|
|
|
|
|
|
def chapter_started(book_id: str, chapter_num: int) -> bool:
|
|
"""
|
|
Return True if this chapter has already started downloading.
|
|
"""
|
|
key = f"started:{book_id}:{chapter_num}"
|
|
return r.exists(key) == 1
|
|
|
|
|
|
# =========================================================
|
|
# UTILITY: RESET FOR A BOOK
|
|
# =========================================================
|
|
|
|
|
|
def reset_book_state(book_id: str):
|
|
"""
|
|
Optional utility: remove abort flag and all started-chapter markers.
|
|
Useful during testing or manual cleanup.
|
|
"""
|
|
# Remove abort flag
|
|
r.delete(f"abort:{book_id}")
|
|
|
|
# Remove all "started:*" keys for this book
|
|
pattern = f"started:{book_id}:*"
|
|
for key in r.scan_iter(pattern):
|
|
r.delete(key)
|