mirror of
https://github.com/alexgetmancom/miband-bot.git
synced 2026-09-05 18:16:18 +03:00
feat: add multi-user support using contextvars
This commit is contained in:
+74
-41
@@ -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 чтобы привязать аккаунт.")
|
||||
|
||||
+32
-21
@@ -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
|
||||
|
||||
+14
-4
@@ -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}")
|
||||
|
||||
|
||||
+13
-7
@@ -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"):
|
||||
|
||||
Reference in New Issue
Block a user