feat: add multi-user support using contextvars

This commit is contained in:
Alex
2026-06-25 01:02:14 +03:00
parent ce2ba35d5d
commit 27b0edcc48
4 changed files with 133 additions and 73 deletions
+74 -41
View File
@@ -70,6 +70,8 @@ except Exception:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Config # Config
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
import contextvars
SETTINGS = Settings.from_env() SETTINGS = Settings.from_env()
BOT_TOKEN = SETTINGS.telegram_bot_token BOT_TOKEN = SETTINGS.telegram_bot_token
ALLOWED_USER_ID = SETTINGS.telegram_allowed_user_id ALLOWED_USER_ID = SETTINGS.telegram_allowed_user_id
@@ -78,54 +80,79 @@ SYNC_LOCK = asyncio.Lock()
AUTH_LOCK = asyncio.Lock() AUTH_LOCK = asyncio.Lock()
AUTO_MENU_REFRESH_INTERVAL = max(5, int(os.getenv("AUTO_MENU_REFRESH_INTERVAL", "30"))) 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 при желании STEP_GOAL = 10_000 # можно вынести в env при желании
# ---------------------------------------------------------------------------
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Auth helpers # 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: 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 uid = update.effective_user.id if update.effective_user else None
if uid is None: if uid is None:
return False return False
if ALLOWED_USER_ID is None:
allowed_ids = SETTINGS.telegram_allowed_user_ids
if not allowed_ids:
ALLOWED_USER_ID = uid ALLOWED_USER_ID = uid
try: try:
allowed_user_file = SETTINGS.data_dir / "allowed_user.id" allowed_user_file = SETTINGS.data_dir / "allowed_user.id"
SETTINGS.data_dir.mkdir(parents=True, exist_ok=True) SETTINGS.data_dir.mkdir(parents=True, exist_ok=True)
allowed_user_file.write_text(str(uid), encoding="utf-8") allowed_user_file.write_text(str(uid), encoding="utf-8")
logger.info("🎉 Бот успешно привязан к первому пользователю (ID: %s)!", uid) logger.info("🎉 Бот успешно привязан к первому пользователю (ID: %s)!", uid)
SETTINGS = Settings.from_env()
allowed_ids = SETTINGS.telegram_allowed_user_ids
except Exception as e: except Exception as e:
logger.error("Не удалось сохранить ID владельца в файл: %s", e) logger.error("Не удалось сохранить ID владельца в файл: %s", e)
return True return True
return uid == ALLOWED_USER_ID return uid in allowed_ids
def with_user_context(func): def with_user_context(func):
@wraps(func) @wraps(func)
async def wrapper(update: Update, context: ContextTypes.DEFAULT_TYPE, *args, **kwargs): 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 return wrapper
def get_user_db_path() -> str: 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 DB_PATH
return str(SETTINGS.user_db_path(ALLOWED_USER_ID)) return str(SETTINGS.user_db_path(uid))
def get_user_status_path() -> str: 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.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: def get_xiaomi_token_path() -> Path | None:
if ALLOWED_USER_ID is None: uid = get_current_user_id()
if uid is None:
return None return None
try: try:
return SETTINGS.token_path(ALLOWED_USER_ID) return SETTINGS.token_path(uid)
except ConfigError: except ConfigError:
return None return None
@@ -220,7 +247,7 @@ def daily_tip(steps: sqlite3.Row | None, sleep: sqlite3.Row | None, hr: sqlite3.
# DB: health # DB: health
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def health_db_exists() -> bool: 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: def health_conn() -> sqlite3.Connection:
@@ -231,11 +258,11 @@ def health_conn() -> sqlite3.Connection:
def fetch_one(query: str, params: tuple = ()) -> sqlite3.Row | None: 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]: 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: async def auto_refresh_main_menu_loop(app: Application) -> None:
"""Refresh the pinned main menu after the sync daemon writes a new status file.""" """Refresh the pinned main menu for all allowed users after the sync daemon writes their status file."""
if ALLOWED_USER_ID is None: last_seen_mtimes: dict[int, float] = {}
return
last_seen_mtime: float | None = None
while True: while True:
try: try:
status_path = SETTINGS.user_status_path(ALLOWED_USER_ID) allowed_ids = SETTINGS.telegram_allowed_user_ids
if status_path.exists(): for uid in allowed_ids:
current_mtime = status_path.stat().st_mtime status_path = SETTINGS.user_status_path(uid)
if last_seen_mtime is None: if status_path.exists():
last_seen_mtime = current_mtime current_mtime = status_path.stat().st_mtime
elif current_mtime > last_seen_mtime: last_seen_mtime = last_seen_mtimes.get(uid)
last_seen_mtime = current_mtime if last_seen_mtime is None:
if get_user_menu_msg_id(ALLOWED_USER_ID): last_seen_mtimes[uid] = current_mtime
await send_or_update_menu( elif current_mtime > last_seen_mtime:
app.bot, last_seen_mtimes[uid] = current_mtime
ALLOWED_USER_ID, if get_user_menu_msg_id(uid):
main_menu_text(), token = current_user_id_var.set(uid)
main_keyboard(), try:
) await send_or_update_menu(
logger.info("Auto-refreshed main menu for user %s", ALLOWED_USER_ID) 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: except asyncio.CancelledError:
raise raise
except Exception as exc: except Exception as exc:
@@ -391,7 +422,7 @@ async def stop_background_tasks(app: Application) -> None:
# Data queries # Data queries
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def read_status_file() -> dict: 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: def latest_steps() -> sqlite3.Row | None:
@@ -1318,7 +1349,7 @@ def db_status_text() -> str:
# Export # Export
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
def zip_export() -> io.BytesIO: 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 return
async with SYNC_LOCK: async with SYNC_LOCK:
result = await run_sync(ALLOWED_USER_ID, SETTINGS) result = await run_sync(get_current_user_id(), SETTINGS)
if result.success: if result.success:
await show_main_menu(update, context) await show_main_menu(update, context)
@@ -1550,7 +1581,7 @@ async def run_manual_sync(
) )
async with SYNC_LOCK: async with SYNC_LOCK:
try: try:
result = await run_sync(ALLOWED_USER_ID, SETTINGS) result = await run_sync(get_current_user_id(), SETTINGS)
except Exception as e: except Exception as e:
logger.exception("Manual sync failed") logger.exception("Manual sync failed")
await update_menu( await update_menu(
@@ -1774,10 +1805,12 @@ def main() -> None:
sys.exit(1) sys.exit(1)
BOT_TOKEN = SETTINGS.telegram_bot_token BOT_TOKEN = SETTINGS.telegram_bot_token
ALLOWED_USER_ID = SETTINGS.telegram_allowed_user_id ALLOWED_USER_ID = SETTINGS.telegram_allowed_user_id
if ALLOWED_USER_ID is not None: allowed_ids = SETTINGS.telegram_allowed_user_ids
DB_PATH = str(SETTINGS.user_db_path(ALLOWED_USER_ID)) if allowed_ids:
print(f"Запуск бота для пользователя ID {ALLOWED_USER_ID}...") print(f"Запуск бота для пользователей: {allowed_ids}...")
storage.init_health_db(Path(DB_PATH)) for uid in allowed_ids:
db_p = SETTINGS.user_db_path(uid)
storage.init_health_db(db_p)
else: else:
DB_PATH = str(SETTINGS.db_path) DB_PATH = str(SETTINGS.db_path)
print("Бот запущен. Отправьте /start в Telegram чтобы привязать аккаунт.") print("Бот запущен. Отправьте /start в Telegram чтобы привязать аккаунт.")
+32 -21
View File
@@ -4,7 +4,7 @@
from __future__ import annotations from __future__ import annotations
import os import os
from dataclasses import dataclass from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
@@ -59,10 +59,17 @@ class Settings:
status_path: Path status_path: Path
bot_state_db_path: Path bot_state_db_path: Path
telegram_bot_token: str telegram_bot_token: str
telegram_allowed_user_id: int | None telegram_allowed_user_ids: list[int] = field(default_factory=list)
sync_interval: int sync_interval: int = 900
query_duration: int query_duration: int = 2
enable_fds_sleep_details: bool 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 @classmethod
def from_env(cls, *, require_bot: bool = False) -> Settings: def from_env(cls, *, require_bot: bool = False) -> Settings:
@@ -78,14 +85,17 @@ class Settings:
else "/opt/miband-tracker/data" else "/opt/miband-tracker/data"
) )
data_dir = Path(os.environ.get("DATA_DIR", _default_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 # Если ID не задан в env, пробуем загрузить из файла allowed_user.id
allowed_user_file = data_dir / "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: 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: except Exception:
pass pass
@@ -103,7 +113,7 @@ class Settings:
) )
), ),
telegram_bot_token=bot_token, 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), sync_interval=_env_int("SYNC_INTERVAL", 900, min_value=0),
query_duration=_env_int("QUERY_DURATION", 2, min_value=1), query_duration=_env_int("QUERY_DURATION", 2, min_value=1),
enable_fds_sleep_details=_env_bool("ENABLE_FDS_SLEEP_DETAILS", default=True), 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: 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 resolved = user_id if user_id is not None else self.telegram_allowed_user_id
if resolved is None: if resolved is None:
raise ConfigError("TELEGRAM_ALLOWED_USER_ID должен содержать ровно один user id") raise ConfigError("Нет доступных пользователей (TELEGRAM_ALLOWED_USER_IDS пуст)")
return int(resolved) return int(resolved)
def token_path(self, user_id: int | None = None) -> Path: 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" 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()] values = [item.strip() for item in raw.split(",") if item.strip()]
if not values: if not values:
if required: if required:
raise ConfigError("TELEGRAM_ALLOWED_USER_ID не задан или пуст") raise ConfigError("TELEGRAM_ALLOWED_USER_IDS не задан или пуст")
return None return []
if len(values) > 1: res = []
raise ConfigError("TELEGRAM_ALLOWED_USER_ID должен содержать ровно один user id") for val in values:
try: try:
return int(values[0]) res.append(int(val))
except ValueError as exc: except ValueError as exc:
raise ConfigError("TELEGRAM_ALLOWED_USER_ID должен быть целым числом") from exc raise ConfigError("Каждый ID в TELEGRAM_ALLOWED_USER_IDS должен быть целым числом") from exc
return res
+14 -4
View File
@@ -597,14 +597,23 @@ async def _sync_workouts(
async def daemon_main(settings: Settings | None = None) -> int: async def daemon_main(settings: Settings | None = None) -> int:
settings = settings or Settings.from_env() settings = settings or Settings.from_env()
if settings.sync_interval <= 0: if settings.sync_interval <= 0:
result = await run_sync(settings=settings) allowed_ids = settings.telegram_allowed_user_ids
return 0 if result.success else 1 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 _waiting_logged = False
while True: while True:
try: try:
current_settings = Settings.from_env() 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: if not _waiting_logged:
log("Синхронизатор ожидает привязки аккаунта через Telegram (/start)...") log("Синхронизатор ожидает привязки аккаунта через Telegram (/start)...")
_waiting_logged = True _waiting_logged = True
@@ -612,7 +621,8 @@ async def daemon_main(settings: Settings | None = None) -> int:
continue continue
_waiting_logged = False # Reset so we log again if user unregisters _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: except Exception as exc:
log(f"Unhandled error in main loop: {exc}") log(f"Unhandled error in main loop: {exc}")
+13 -7
View File
@@ -5,12 +5,17 @@ from pathlib import Path
import pytest 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): 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: 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", status_path=tmp_path / "status.json",
bot_state_db_path=tmp_path / "fitness_bot_state.db", bot_state_db_path=tmp_path / "fitness_bot_state.db",
telegram_bot_token="token", telegram_bot_token="token",
telegram_allowed_user_id=user_id, telegram_allowed_user_ids=[user_id],
sync_interval=900, sync_interval=900,
query_duration=2, query_duration=2,
enable_fds_sleep_details=True, 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_db_path() == tmp_path / f"miband_{user_id}.db"
assert settings.user_status_path() == tmp_path / f"status_{user_id}.json" assert settings.user_status_path() == tmp_path / f"status_{user_id}.json"
assert settings.token_path() == tmp_path / f"token_{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", status_path=tmp_path / "status.json",
bot_state_db_path=tmp_path / "fitness_bot_state.db", bot_state_db_path=tmp_path / "fitness_bot_state.db",
telegram_bot_token="token", telegram_bot_token="token",
telegram_allowed_user_id=123, telegram_allowed_user_ids=[123],
sync_interval=900, sync_interval=900,
query_duration=2, query_duration=2,
enable_fds_sleep_details=True, 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: 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") monkeypatch.setenv("SYNC_INTERVAL", "soon")
with pytest.raises(ConfigError, match="SYNC_INTERVAL"): 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: 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") monkeypatch.setenv("QUERY_DURATION", "0")
with pytest.raises(ConfigError, match="QUERY_DURATION"): with pytest.raises(ConfigError, match="QUERY_DURATION"):