From 27b0edcc48a9df5442469d3fd61dd7b459eb7f98 Mon Sep 17 00:00:00 2001 From: Alex Date: Thu, 25 Jun 2026 01:02:14 +0300 Subject: [PATCH] feat: add multi-user support using contextvars --- miband_tracker/bot/app.py | 115 ++++++++++++++++++++++++-------------- miband_tracker/config.py | 53 +++++++++++------- miband_tracker/sync.py | 18 ++++-- tests/test_config.py | 20 ++++--- 4 files changed, 133 insertions(+), 73 deletions(-) diff --git a/miband_tracker/bot/app.py b/miband_tracker/bot/app.py index bda2f34..5220fc7 100644 --- a/miband_tracker/bot/app.py +++ b/miband_tracker/bot/app.py @@ -70,6 +70,8 @@ except Exception: # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- +import contextvars + SETTINGS = Settings.from_env() BOT_TOKEN = SETTINGS.telegram_bot_token ALLOWED_USER_ID = SETTINGS.telegram_allowed_user_id @@ -78,54 +80,79 @@ SYNC_LOCK = asyncio.Lock() AUTH_LOCK = asyncio.Lock() AUTO_MENU_REFRESH_INTERVAL = max(5, int(os.getenv("AUTO_MENU_REFRESH_INTERVAL", "30"))) +current_user_id_var = contextvars.ContextVar("current_user_id", default=None) + STEP_GOAL = 10_000 # можно вынести в env при желании +# --------------------------------------------------------------------------- # --------------------------------------------------------------------------- # Auth helpers # --------------------------------------------------------------------------- +def get_current_user_id() -> int | None: + uid = current_user_id_var.get() + if uid is not None: + return uid + allowed_ids = SETTINGS.telegram_allowed_user_ids + if allowed_ids: + return allowed_ids[0] + return None + + def is_allowed(update: Update) -> bool: - global ALLOWED_USER_ID + global SETTINGS, ALLOWED_USER_ID uid = update.effective_user.id if update.effective_user else None if uid is None: return False - if ALLOWED_USER_ID is None: + + allowed_ids = SETTINGS.telegram_allowed_user_ids + if not allowed_ids: ALLOWED_USER_ID = uid try: allowed_user_file = SETTINGS.data_dir / "allowed_user.id" SETTINGS.data_dir.mkdir(parents=True, exist_ok=True) allowed_user_file.write_text(str(uid), encoding="utf-8") logger.info("🎉 Бот успешно привязан к первому пользователю (ID: %s)!", uid) + SETTINGS = Settings.from_env() + allowed_ids = SETTINGS.telegram_allowed_user_ids except Exception as e: logger.error("Не удалось сохранить ID владельца в файл: %s", e) return True - return uid == ALLOWED_USER_ID + return uid in allowed_ids def with_user_context(func): @wraps(func) async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs): - return await func(update, context, *args, **kwargs) + uid = update.effective_user.id if update.effective_user else None + token = current_user_id_var.set(uid) + try: + return await func(update, context, *args, **kwargs) + finally: + current_user_id_var.reset(token) return wrapper def get_user_db_path() -> str: - if ALLOWED_USER_ID is None: + uid = get_current_user_id() + if uid is None: return DB_PATH - return str(SETTINGS.user_db_path(ALLOWED_USER_ID)) + return str(SETTINGS.user_db_path(uid)) def get_user_status_path() -> str: - if ALLOWED_USER_ID is None: + uid = get_current_user_id() + if uid is None: return str(SETTINGS.status_path) - return str(SETTINGS.user_status_path(ALLOWED_USER_ID)) + return str(SETTINGS.user_status_path(uid)) def get_xiaomi_token_path() -> Path | None: - if ALLOWED_USER_ID is None: + uid = get_current_user_id() + if uid is None: return None try: - return SETTINGS.token_path(ALLOWED_USER_ID) + return SETTINGS.token_path(uid) except ConfigError: return None @@ -220,7 +247,7 @@ def daily_tip(steps: sqlite3.Row | None, sleep: sqlite3.Row | None, hr: sqlite3. # DB: health # --------------------------------------------------------------------------- def health_db_exists() -> bool: - return storage.health_db_exists(SETTINGS, ALLOWED_USER_ID) + return storage.health_db_exists(SETTINGS, get_current_user_id()) def health_conn() -> sqlite3.Connection: @@ -231,11 +258,11 @@ def health_conn() -> sqlite3.Connection: def fetch_one(query: str, params: tuple = ()) -> sqlite3.Row | None: - return storage.fetch_one(SETTINGS, query, params, ALLOWED_USER_ID) + return storage.fetch_one(SETTINGS, query, params, get_current_user_id()) def fetch_all(query: str, params: tuple = ()) -> list[sqlite3.Row]: - return storage.fetch_all(SETTINGS, query, params, ALLOWED_USER_ID) + return storage.fetch_all(SETTINGS, query, params, get_current_user_id()) # --------------------------------------------------------------------------- @@ -340,28 +367,32 @@ async def update_menu( async def auto_refresh_main_menu_loop(app: Application) -> None: - """Refresh the pinned main menu after the sync daemon writes a new status file.""" - if ALLOWED_USER_ID is None: - return - - last_seen_mtime: float | None = None + """Refresh the pinned main menu for all allowed users after the sync daemon writes their status file.""" + last_seen_mtimes: dict[int, float] = {} while True: try: - status_path = SETTINGS.user_status_path(ALLOWED_USER_ID) - if status_path.exists(): - current_mtime = status_path.stat().st_mtime - if last_seen_mtime is None: - last_seen_mtime = current_mtime - elif current_mtime > last_seen_mtime: - last_seen_mtime = current_mtime - if get_user_menu_msg_id(ALLOWED_USER_ID): - await send_or_update_menu( - app.bot, - ALLOWED_USER_ID, - main_menu_text(), - main_keyboard(), - ) - logger.info("Auto-refreshed main menu for user %s", ALLOWED_USER_ID) + allowed_ids = SETTINGS.telegram_allowed_user_ids + for uid in allowed_ids: + status_path = SETTINGS.user_status_path(uid) + if status_path.exists(): + current_mtime = status_path.stat().st_mtime + last_seen_mtime = last_seen_mtimes.get(uid) + if last_seen_mtime is None: + last_seen_mtimes[uid] = current_mtime + elif current_mtime > last_seen_mtime: + last_seen_mtimes[uid] = current_mtime + if get_user_menu_msg_id(uid): + token = current_user_id_var.set(uid) + try: + await send_or_update_menu( + app.bot, + uid, + main_menu_text(), + main_keyboard(), + ) + logger.info("Auto-refreshed main menu for user %s", uid) + finally: + current_user_id_var.reset(token) except asyncio.CancelledError: raise except Exception as exc: @@ -391,7 +422,7 @@ async def stop_background_tasks(app: Application) -> None: # Data queries # --------------------------------------------------------------------------- def read_status_file() -> dict: - return storage.read_status_file(SETTINGS, ALLOWED_USER_ID) + return storage.read_status_file(SETTINGS, get_current_user_id()) def latest_steps() -> sqlite3.Row | None: @@ -1318,7 +1349,7 @@ def db_status_text() -> str: # Export # --------------------------------------------------------------------------- def zip_export() -> io.BytesIO: - return storage.zip_export(SETTINGS, ALLOWED_USER_ID) + return storage.zip_export(SETTINGS, get_current_user_id()) # --------------------------------------------------------------------------- @@ -1441,7 +1472,7 @@ async def run_initial_sync_after_login(update: Update, context: ContextTypes.DEF return async with SYNC_LOCK: - result = await run_sync(ALLOWED_USER_ID, SETTINGS) + result = await run_sync(get_current_user_id(), SETTINGS) if result.success: await show_main_menu(update, context) @@ -1550,7 +1581,7 @@ async def run_manual_sync( ) async with SYNC_LOCK: try: - result = await run_sync(ALLOWED_USER_ID, SETTINGS) + result = await run_sync(get_current_user_id(), SETTINGS) except Exception as e: logger.exception("Manual sync failed") await update_menu( @@ -1774,10 +1805,12 @@ def main() -> None: sys.exit(1) BOT_TOKEN = SETTINGS.telegram_bot_token ALLOWED_USER_ID = SETTINGS.telegram_allowed_user_id - if ALLOWED_USER_ID is not None: - DB_PATH = str(SETTINGS.user_db_path(ALLOWED_USER_ID)) - print(f"Запуск бота для пользователя ID {ALLOWED_USER_ID}...") - storage.init_health_db(Path(DB_PATH)) + allowed_ids = SETTINGS.telegram_allowed_user_ids + if allowed_ids: + print(f"Запуск бота для пользователей: {allowed_ids}...") + for uid in allowed_ids: + db_p = SETTINGS.user_db_path(uid) + storage.init_health_db(db_p) else: DB_PATH = str(SETTINGS.db_path) print("Бот запущен. Отправьте /start в Telegram чтобы привязать аккаунт.") diff --git a/miband_tracker/config.py b/miband_tracker/config.py index 6b62afd..7c2d129 100644 --- a/miband_tracker/config.py +++ b/miband_tracker/config.py @@ -4,7 +4,7 @@ from __future__ import annotations import os -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path @@ -59,10 +59,17 @@ class Settings: status_path: Path bot_state_db_path: Path telegram_bot_token: str - telegram_allowed_user_id: int | None - sync_interval: int - query_duration: int - enable_fds_sleep_details: bool + telegram_allowed_user_ids: list[int] = field(default_factory=list) + sync_interval: int = 900 + query_duration: int = 2 + enable_fds_sleep_details: bool = True + telegram_allowed_user_id: int | None = None + + def __post_init__(self) -> None: + if self.telegram_allowed_user_id is not None and not self.telegram_allowed_user_ids: + object.__setattr__(self, "telegram_allowed_user_ids", [self.telegram_allowed_user_id]) + elif self.telegram_allowed_user_ids and self.telegram_allowed_user_id is None: + object.__setattr__(self, "telegram_allowed_user_id", self.telegram_allowed_user_ids[0]) @classmethod def from_env(cls, *, require_bot: bool = False) -> Settings: @@ -78,14 +85,17 @@ class Settings: else "/opt/miband-tracker/data" ) data_dir = Path(os.environ.get("DATA_DIR", _default_data)) - allowed_user_id = parse_single_user_id( - os.environ.get("TELEGRAM_ALLOWED_USER_ID", ""), required=False - ) + + # Read from TELEGRAM_ALLOWED_USER_IDS or legacy TELEGRAM_ALLOWED_USER_ID + raw_ids = os.environ.get("TELEGRAM_ALLOWED_USER_IDS", os.environ.get("TELEGRAM_ALLOWED_USER_ID", "")) + allowed_user_ids = parse_user_ids(raw_ids, required=False) + # Если ID не задан в env, пробуем загрузить из файла allowed_user.id allowed_user_file = data_dir / "allowed_user.id" - if allowed_user_id is None and allowed_user_file.exists(): + if not allowed_user_ids and allowed_user_file.exists(): try: - allowed_user_id = int(allowed_user_file.read_text(encoding="utf-8").strip()) + raw_file = allowed_user_file.read_text(encoding="utf-8").strip() + allowed_user_ids = parse_user_ids(raw_file) except Exception: pass @@ -103,7 +113,7 @@ class Settings: ) ), telegram_bot_token=bot_token, - telegram_allowed_user_id=allowed_user_id, + telegram_allowed_user_ids=allowed_user_ids, sync_interval=_env_int("SYNC_INTERVAL", 900, min_value=0), query_duration=_env_int("QUERY_DURATION", 2, min_value=1), enable_fds_sleep_details=_env_bool("ENABLE_FDS_SLEEP_DETAILS", default=True), @@ -112,7 +122,7 @@ class Settings: def require_user_id(self, user_id: int | None = None) -> int: resolved = user_id if user_id is not None else self.telegram_allowed_user_id if resolved is None: - raise ConfigError("TELEGRAM_ALLOWED_USER_ID должен содержать ровно один user id") + raise ConfigError("Нет доступных пользователей (TELEGRAM_ALLOWED_USER_IDS пуст)") return int(resolved) def token_path(self, user_id: int | None = None) -> Path: @@ -144,15 +154,16 @@ class Settings: return self.data_dir / f"status_{self.require_user_id(user_id)}.json" -def parse_single_user_id(raw: str, *, required: bool = False) -> int | None: +def parse_user_ids(raw: str, *, required: bool = False) -> list[int]: values = [item.strip() for item in raw.split(",") if item.strip()] if not values: if required: - raise ConfigError("TELEGRAM_ALLOWED_USER_ID не задан или пуст") - return None - if len(values) > 1: - raise ConfigError("TELEGRAM_ALLOWED_USER_ID должен содержать ровно один user id") - try: - return int(values[0]) - except ValueError as exc: - raise ConfigError("TELEGRAM_ALLOWED_USER_ID должен быть целым числом") from exc + raise ConfigError("TELEGRAM_ALLOWED_USER_IDS не задан или пуст") + return [] + res = [] + for val in values: + try: + res.append(int(val)) + except ValueError as exc: + raise ConfigError("Каждый ID в TELEGRAM_ALLOWED_USER_IDS должен быть целым числом") from exc + return res diff --git a/miband_tracker/sync.py b/miband_tracker/sync.py index 38f95c8..954761e 100644 --- a/miband_tracker/sync.py +++ b/miband_tracker/sync.py @@ -597,14 +597,23 @@ async def _sync_workouts( async def daemon_main(settings: Settings | None = None) -> int: settings = settings or Settings.from_env() if settings.sync_interval <= 0: - result = await run_sync(settings=settings) - return 0 if result.success else 1 + allowed_ids = settings.telegram_allowed_user_ids + if not allowed_ids: + log("Нет разрешенных пользователей для синхронизации.") + return 1 + success = True + for uid in allowed_ids: + result = await run_sync(user_id=uid, settings=settings) + if not result.success: + success = False + return 0 if success else 1 _waiting_logged = False while True: try: current_settings = Settings.from_env() - if current_settings.telegram_allowed_user_id is None: + allowed_ids = current_settings.telegram_allowed_user_ids + if not allowed_ids: if not _waiting_logged: log("Синхронизатор ожидает привязки аккаунта через Telegram (/start)...") _waiting_logged = True @@ -612,7 +621,8 @@ async def daemon_main(settings: Settings | None = None) -> int: continue _waiting_logged = False # Reset so we log again if user unregisters - await run_sync(settings=current_settings) + for uid in allowed_ids: + await run_sync(user_id=uid, settings=current_settings) except Exception as exc: log(f"Unhandled error in main loop: {exc}") diff --git a/tests/test_config.py b/tests/test_config.py index c0d9c4e..047c8ec 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -5,12 +5,17 @@ from pathlib import Path import pytest -from miband_tracker.config import ConfigError, Settings, parse_single_user_id +from miband_tracker.config import ConfigError, Settings, parse_user_ids -def test_parse_single_user_id_rejects_multiple_values() -> None: +def test_parse_user_ids_accepts_multiple_values() -> None: + assert parse_user_ids("1,2", required=True) == [1, 2] + assert parse_user_ids("1", required=True) == [1] + + +def test_parse_user_ids_rejects_invalid_values() -> None: with pytest.raises(ConfigError): - parse_single_user_id("1,2", required=True) + parse_user_ids("1,invalid", required=True) def test_settings_user_paths_prefer_existing_user_files(tmp_path: Path) -> None: @@ -25,12 +30,13 @@ def test_settings_user_paths_prefer_existing_user_files(tmp_path: Path) -> None: status_path=tmp_path / "status.json", bot_state_db_path=tmp_path / "fitness_bot_state.db", telegram_bot_token="token", - telegram_allowed_user_id=user_id, + telegram_allowed_user_ids=[user_id], sync_interval=900, query_duration=2, enable_fds_sleep_details=True, ) + assert settings.telegram_allowed_user_id == user_id assert settings.user_db_path() == tmp_path / f"miband_{user_id}.db" assert settings.user_status_path() == tmp_path / f"status_{user_id}.json" assert settings.token_path() == tmp_path / f"token_{user_id}.json" @@ -43,7 +49,7 @@ def test_settings_falls_back_to_legacy_db_and_status(tmp_path: Path) -> None: status_path=tmp_path / "status.json", bot_state_db_path=tmp_path / "fitness_bot_state.db", telegram_bot_token="token", - telegram_allowed_user_id=123, + telegram_allowed_user_ids=[123], sync_interval=900, query_duration=2, enable_fds_sleep_details=True, @@ -55,7 +61,7 @@ def test_settings_falls_back_to_legacy_db_and_status(tmp_path: Path) -> None: def test_settings_from_env_rejects_invalid_interval(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("TELEGRAM_ALLOWED_USER_ID", "123") + monkeypatch.setenv("TELEGRAM_ALLOWED_USER_IDS", "123") monkeypatch.setenv("SYNC_INTERVAL", "soon") with pytest.raises(ConfigError, match="SYNC_INTERVAL"): @@ -63,7 +69,7 @@ def test_settings_from_env_rejects_invalid_interval(monkeypatch: pytest.MonkeyPa def test_settings_from_env_rejects_zero_query_duration(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv("TELEGRAM_ALLOWED_USER_ID", "123") + monkeypatch.setenv("TELEGRAM_ALLOWED_USER_IDS", "123") monkeypatch.setenv("QUERY_DURATION", "0") with pytest.raises(ConfigError, match="QUERY_DURATION"):