From 50347038e97c398442c57bf7d18be74ec29a9d5d Mon Sep 17 00:00:00 2001 From: Nikolay Tatarinov Date: Wed, 18 Feb 2026 01:00:31 +0300 Subject: [PATCH] Implement group duty pinning and user phone management - Added functionality to pin duty messages in group chats, including scheduling updates and handling bot add/remove events. - Introduced a new `GroupDutyPin` model to store pinned message details and a `phone` field in the `User` model for user contact information. - Implemented commands for users to set or clear their phone numbers in private chats. - Enhanced the repository with functions to manage group duty pins and user phone data. - Updated handlers to register new commands and manage duty pin updates effectively. --- .env.example | 3 + .../004_group_duty_pins_and_user_phone.py | 35 +++ config.py | 3 + db/models.py | 10 + db/repository.py | 102 ++++++- handlers/__init__.py | 5 +- handlers/commands.py | 53 +++- handlers/group_duty_pin.py | 280 ++++++++++++++++++ main.py | 11 +- 9 files changed, 495 insertions(+), 7 deletions(-) create mode 100644 alembic/versions/004_group_duty_pins_and_user_phone.py create mode 100644 handlers/group_duty_pin.py diff --git a/.env.example b/.env.example index afc129a..98b7f07 100644 --- a/.env.example +++ b/.env.example @@ -12,3 +12,6 @@ ADMIN_USERNAMES=admin1,admin2 # Optional: URL of a public ICS calendar (e.g. holidays). Days from this calendar are highlighted on the duty grid; click "i" for summary. # EXTERNAL_CALENDAR_ICS_URL=https://example.com/holidays.ics + +# Timezone for the pinned duty message in groups (e.g. Europe/Moscow). +# DUTY_DISPLAY_TZ=Europe/Moscow diff --git a/alembic/versions/004_group_duty_pins_and_user_phone.py b/alembic/versions/004_group_duty_pins_and_user_phone.py new file mode 100644 index 0000000..28b53e7 --- /dev/null +++ b/alembic/versions/004_group_duty_pins_and_user_phone.py @@ -0,0 +1,35 @@ +"""Group duty pins table and User.phone + +Revision ID: 004 +Revises: 003 +Create Date: 2025-02-18 + +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa + +revision: str = "004" +down_revision: Union[str, None] = "003" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.add_column( + "users", + sa.Column("phone", sa.Text(), nullable=True), + ) + op.create_table( + "group_duty_pins", + sa.Column("chat_id", sa.BigInteger(), nullable=False), + sa.Column("message_id", sa.Integer(), nullable=False), + sa.PrimaryKeyConstraint("chat_id"), + ) + + +def downgrade() -> None: + op.drop_table("group_duty_pins") + op.drop_column("users", "phone") diff --git a/config.py b/config.py index 3cde1dd..31bc23e 100644 --- a/config.py +++ b/config.py @@ -44,6 +44,9 @@ CORS_ORIGINS = ( # Optional: URL of a public ICS calendar (e.g. holidays). Empty = no external calendar; /api/calendar-events returns []. EXTERNAL_CALENDAR_ICS_URL = os.getenv("EXTERNAL_CALENDAR_ICS_URL", "").strip() +# Timezone for displaying duty times in the pinned group message (e.g. Europe/Moscow). +DUTY_DISPLAY_TZ = os.getenv("DUTY_DISPLAY_TZ", "Europe/Moscow").strip() or "Europe/Moscow" + def is_admin(username: str) -> bool: """True if the given Telegram username (no @, any case) is in ADMIN_USERNAMES.""" diff --git a/db/models.py b/db/models.py index 68e4419..8154fa4 100644 --- a/db/models.py +++ b/db/models.py @@ -21,6 +21,7 @@ class User(Base): username: Mapped[str | None] = mapped_column(Text, nullable=True) first_name: Mapped[str | None] = mapped_column(Text, nullable=True) last_name: Mapped[str | None] = mapped_column(Text, nullable=True) + phone: Mapped[str | None] = mapped_column(Text, nullable=True) duties: Mapped[list["Duty"]] = relationship("Duty", back_populates="user") @@ -39,3 +40,12 @@ class Duty(Base): event_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="duty") user: Mapped["User"] = relationship("User", back_populates="duties") + + +class GroupDutyPin(Base): + """Stores which message to update in each group for the pinned duty notice.""" + + __tablename__ = "group_duty_pins" + + chat_id: Mapped[int] = mapped_column(BigInteger, primary_key=True) + message_id: Mapped[int] = mapped_column(Integer, nullable=False) diff --git a/db/repository.py b/db/repository.py index 0001005..46b67f3 100644 --- a/db/repository.py +++ b/db/repository.py @@ -1,10 +1,10 @@ -"""Repository: get_or_create_user, get_duties, insert_duty.""" +"""Repository: get_or_create_user, get_duties, insert_duty, get_current_duty, group_duty_pins.""" from datetime import datetime, timedelta from sqlalchemy.orm import Session -from db.models import User, Duty +from db.models import User, Duty, GroupDutyPin def get_or_create_user( @@ -116,3 +116,101 @@ def insert_duty( session.commit() session.refresh(duty) return duty + + +def get_current_duty( + session: Session, at_utc: datetime +) -> tuple[Duty, User] | None: + """Return the duty (and user) for which start_at <= at_utc < end_at, event_type='duty'. + at_utc is in UTC (naive or aware); comparison uses ISO strings.""" + from datetime import timezone + if at_utc.tzinfo is not None: + at_utc = at_utc.astimezone(timezone.utc) + now_iso = at_utc.strftime("%Y-%m-%dT%H:%M:%S") + "Z" + row = ( + session.query(Duty, User) + .join(User, Duty.user_id == User.id) + .filter( + Duty.event_type == "duty", + Duty.start_at <= now_iso, + Duty.end_at > now_iso, + ) + .first() + ) + if row is None: + return None + return (row[0], row[1]) + + +def get_next_shift_end(session: Session, after_utc: datetime) -> datetime | None: + """Return the end_at of the current duty (if after_utc is inside one) or of the next duty. + For scheduling the next pin update. Returns naive UTC datetime.""" + from datetime import timezone + if after_utc.tzinfo is not None: + after_utc = after_utc.astimezone(timezone.utc) + after_iso = after_utc.strftime("%Y-%m-%dT%H:%M:%S") + "Z" + # Current duty: start_at <= after_iso < end_at → use this end_at + current = ( + session.query(Duty) + .filter( + Duty.event_type == "duty", + Duty.start_at <= after_iso, + Duty.end_at > after_iso, + ) + .first() + ) + if current: + return datetime.fromisoformat(current.end_at.replace("Z", "+00:00")).replace(tzinfo=None) + # Next future duty: start_at > after_iso, order by start_at + next_duty = ( + session.query(Duty) + .filter(Duty.event_type == "duty", Duty.start_at > after_iso) + .order_by(Duty.start_at) + .first() + ) + if next_duty: + return datetime.fromisoformat(next_duty.end_at.replace("Z", "+00:00")).replace(tzinfo=None) + return None + + +def get_group_duty_pin(session: Session, chat_id: int) -> GroupDutyPin | None: + """Get the pinned message record for a chat, if any.""" + return session.query(GroupDutyPin).filter(GroupDutyPin.chat_id == chat_id).first() + + +def save_group_duty_pin( + session: Session, chat_id: int, message_id: int +) -> GroupDutyPin: + """Save or update the pinned message for a chat.""" + pin = session.query(GroupDutyPin).filter(GroupDutyPin.chat_id == chat_id).first() + if pin: + pin.message_id = message_id + else: + pin = GroupDutyPin(chat_id=chat_id, message_id=message_id) + session.add(pin) + session.commit() + session.refresh(pin) + return pin + + +def delete_group_duty_pin(session: Session, chat_id: int) -> None: + """Remove the pinned message record when the bot leaves the group.""" + session.query(GroupDutyPin).filter(GroupDutyPin.chat_id == chat_id).delete() + session.commit() + + +def get_all_group_duty_pin_chat_ids(session: Session) -> list[int]: + """Return all chat_ids that have a pinned duty message (for restoring jobs on startup).""" + rows = session.query(GroupDutyPin.chat_id).all() + return [r[0] for r in rows] + + +def set_user_phone(session: Session, telegram_user_id: int, phone: str | None) -> User | None: + """Set phone for user by telegram_user_id. Returns User or None if not found.""" + user = session.query(User).filter(User.telegram_user_id == telegram_user_id).first() + if not user: + return None + user.phone = phone + session.commit() + session.refresh(user) + return user diff --git a/handlers/__init__.py b/handlers/__init__.py index 0a2539b..3a70a3a 100644 --- a/handlers/__init__.py +++ b/handlers/__init__.py @@ -2,13 +2,16 @@ from telegram.ext import Application -from . import commands, errors, import_duty_schedule +from . import commands, errors, group_duty_pin, import_duty_schedule def register_handlers(app: Application) -> None: app.add_handler(commands.start_handler) app.add_handler(commands.help_handler) + app.add_handler(commands.set_phone_handler) app.add_handler(import_duty_schedule.import_duty_schedule_handler) app.add_handler(import_duty_schedule.handover_time_handler) app.add_handler(import_duty_schedule.duty_schedule_document_handler) + app.add_handler(group_duty_pin.group_duty_pin_handler) + app.add_handler(group_duty_pin.pin_duty_handler) app.add_error_handler(errors.error_handler) diff --git a/handlers/commands.py b/handlers/commands.py index 8879931..bdafe2e 100644 --- a/handlers/commands.py +++ b/handlers/commands.py @@ -7,7 +7,7 @@ from telegram import Update from telegram.ext import CommandHandler, ContextTypes from db.session import get_session -from db.repository import get_or_create_user +from db.repository import get_or_create_user, set_user_phone async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: @@ -45,6 +45,54 @@ async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: await update.message.reply_text(text) +async def set_phone(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Set or clear phone for the current user (private chat only).""" + if not update.message or not update.effective_user: + return + if update.effective_chat and update.effective_chat.type != "private": + await update.message.reply_text("Команда /set_phone доступна только в личке.") + return + # Optional: restrict to allowed usernames; plan says "or without restrictions" + args = (context.args or []) + phone = " ".join(args).strip() if args else None + telegram_user_id = update.effective_user.id + + def do_set_phone() -> str: + session = get_session(config.DATABASE_URL) + try: + full_name = ( + " ".join( + filter( + None, + [ + update.effective_user.first_name or "", + update.effective_user.last_name or "", + ], + ) + ).strip() + or "User" + ) + get_or_create_user( + session, + telegram_user_id=telegram_user_id, + full_name=full_name, + username=update.effective_user.username, + first_name=update.effective_user.first_name, + last_name=update.effective_user.last_name, + ) + user = set_user_phone(session, telegram_user_id, phone or None) + if user is None: + return "Ошибка сохранения." + if phone: + return f"Телефон сохранён: {phone}" + return "Телефон очищен." + finally: + session.close() + + result = await asyncio.get_running_loop().run_in_executor(None, do_set_phone) + await update.message.reply_text(result) + + async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: if not update.message or not update.effective_user: return @@ -52,6 +100,8 @@ async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: "Доступные команды:", "/start — Начать", "/help — Показать эту справку", + "/set_phone — Указать или очистить телефон для отображения в дежурстве", + "/pin_duty — В группе: закрепить сообщение о дежурстве (нужны права админа у бота)", ] if config.is_admin(update.effective_user.username or ""): lines.append("/import_duty_schedule — Импорт расписания дежурств (JSON)") @@ -60,3 +110,4 @@ async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: start_handler = CommandHandler("start", start) help_handler = CommandHandler("help", help_cmd) +set_phone_handler = CommandHandler("set_phone", set_phone) diff --git a/handlers/group_duty_pin.py b/handlers/group_duty_pin.py new file mode 100644 index 0000000..d3cf5cc --- /dev/null +++ b/handlers/group_duty_pin.py @@ -0,0 +1,280 @@ +"""Pinned duty message in groups: handle bot add/remove, schedule updates at shift end.""" + +import asyncio +import logging +from datetime import datetime, timezone +from zoneinfo import ZoneInfo + +import config +from telegram import Update +from telegram.constants import ChatMemberStatus +from telegram.error import BadRequest, Forbidden +from telegram.ext import ChatMemberHandler, CommandHandler, ContextTypes + +from db.session import get_session +from db.repository import ( + get_current_duty, + get_next_shift_end, + get_group_duty_pin, + save_group_duty_pin, + delete_group_duty_pin, + get_all_group_duty_pin_chat_ids, +) + +logger = logging.getLogger(__name__) + +JOB_NAME_PREFIX = "duty_pin_" +RETRY_WHEN_NO_DUTY_MINUTES = 15 + + +def _format_duty_message(duty, user, tz_name: str) -> str: + """Build the text for the pinned message. duty, user may be None.""" + if duty is None or user is None: + return "Сейчас дежурства нет." + try: + tz = ZoneInfo(tz_name) + except Exception: + tz = ZoneInfo("Europe/Moscow") + tz_name = "Europe/Moscow" + start_dt = datetime.fromisoformat(duty.start_at.replace("Z", "+00:00")) + end_dt = datetime.fromisoformat(duty.end_at.replace("Z", "+00:00")) + start_local = start_dt.astimezone(tz) + end_local = end_dt.astimezone(tz) + # Показать смещение (UTC+3) чтобы было понятно, в каком поясе время + offset_sec = start_local.utcoffset().total_seconds() if start_local.utcoffset() else 0 + sign = "+" if offset_sec >= 0 else "-" + h, r = divmod(abs(int(offset_sec)), 3600) + m = r // 60 + tz_hint = f"UTC{sign}{h:d}:{m:02d}, {tz_name}" + time_range = f"{start_local.strftime('%d.%m.%Y %H:%M')} — {end_local.strftime('%d.%m.%Y %H:%M')} ({tz_hint})" + lines = [ + f"🕐 Дежурство: {time_range}", + f"👤 {user.full_name}", + ] + if user.phone: + lines.append(f"📞 {user.phone}") + if user.username: + lines.append(f"@{user.username}") + return "\n".join(lines) + + +def _get_duty_message_text() -> str: + """Get current duty from DB and return formatted message (sync, for run_in_executor).""" + session = get_session(config.DATABASE_URL) + try: + now = datetime.now(timezone.utc) + result = get_current_duty(session, now) + if result is None: + return "Сейчас дежурства нет." + duty, user = result + return _format_duty_message(duty, user, config.DUTY_DISPLAY_TZ) + finally: + session.close() + + +def _get_next_shift_end_utc(): + """Return next shift end as naive UTC datetime for job scheduling (sync).""" + session = get_session(config.DATABASE_URL) + try: + return get_next_shift_end(session, datetime.now(timezone.utc)) + finally: + session.close() + + +def _sync_save_pin(chat_id: int, message_id: int) -> None: + session = get_session(config.DATABASE_URL) + try: + save_group_duty_pin(session, chat_id, message_id) + finally: + session.close() + + +def _sync_delete_pin(chat_id: int) -> None: + session = get_session(config.DATABASE_URL) + try: + delete_group_duty_pin(session, chat_id) + finally: + session.close() + + +def _sync_get_message_id(chat_id: int) -> int | None: + session = get_session(config.DATABASE_URL) + try: + pin = get_group_duty_pin(session, chat_id) + return pin.message_id if pin else None + finally: + session.close() + + +async def _schedule_next_update( + application, chat_id: int, when_utc: datetime | None +) -> None: + """Schedule run_once for update_group_pin. Remove existing job with same name first.""" + job_queue = application.job_queue + if job_queue is None: + logger.warning("Job queue not available, cannot schedule pin update") + return + name = f"{JOB_NAME_PREFIX}{chat_id}" + for job in job_queue.get_jobs_by_name(name): + job.schedule_removal() + if when_utc is not None: + delay = when_utc - datetime.utcnow() + if delay.total_seconds() < 1: + delay = 1 + job_queue.run_once( + update_group_pin, + when=delay, + data={"chat_id": chat_id}, + name=name, + ) + logger.info("Scheduled pin update for chat_id=%s at %s", chat_id, when_utc) + else: + from datetime import timedelta + job_queue.run_once( + update_group_pin, + when=timedelta(minutes=RETRY_WHEN_NO_DUTY_MINUTES), + data={"chat_id": chat_id}, + name=name, + ) + logger.info( + "No next shift for chat_id=%s; scheduled retry in %s min", + chat_id, + RETRY_WHEN_NO_DUTY_MINUTES, + ) + + +async def update_group_pin(context: ContextTypes.DEFAULT_TYPE) -> None: + """Job callback: edit pinned message with current duty and schedule next update.""" + chat_id = context.job.data.get("chat_id") + if chat_id is None: + return + loop = asyncio.get_running_loop() + message_id = await loop.run_in_executor(None, _sync_get_message_id, chat_id) + if message_id is None: + logger.info("No pin record for chat_id=%s, skipping update", chat_id) + return + text = await loop.run_in_executor(None, _get_duty_message_text) + try: + await context.bot.edit_message_text( + chat_id=chat_id, + message_id=message_id, + text=text, + ) + except (BadRequest, Forbidden) as e: + logger.warning("Failed to edit pinned message chat_id=%s: %s", chat_id, e) + next_end = await loop.run_in_executor(None, _get_next_shift_end_utc) + await _schedule_next_update(context.application, chat_id, next_end) + + +async def my_chat_member_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """On bot added to group: send, pin, save, schedule. On removed: delete pin, cancel job.""" + if not update.my_chat_member or not update.effective_user: + return + old = update.my_chat_member.old_chat_member + new = update.my_chat_member.new_chat_member + chat = update.effective_chat + if not chat or chat.type not in ("group", "supergroup"): + return + if new.user.id != context.bot.id: + return + chat_id = chat.id + + # Bot added to group + if new.status in (ChatMemberStatus.MEMBER, ChatMemberStatus.ADMINISTRATOR) and old.status in ( + ChatMemberStatus.LEFT, + ChatMemberStatus.BANNED, + ): + loop = asyncio.get_running_loop() + text = await loop.run_in_executor(None, _get_duty_message_text) + try: + msg = await context.bot.send_message(chat_id=chat_id, text=text) + except (BadRequest, Forbidden) as e: + logger.warning("Failed to send duty message in chat_id=%s: %s", chat_id, e) + return + # Make it a pinned post (bot must be admin with "Pin messages" right) + pinned = False + try: + await context.bot.pin_chat_message( + chat_id=chat_id, + message_id=msg.message_id, + disable_notification=True, + ) + pinned = True + except (BadRequest, Forbidden) as e: + logger.warning("Failed to pin message in chat_id=%s: %s", chat_id, e) + await loop.run_in_executor(None, _sync_save_pin, chat_id, msg.message_id) + if not pinned: + try: + await context.bot.send_message( + chat_id=chat_id, + text="Сообщение о дежурстве отправлено, но закрепить его не удалось. " + "Сделайте бота администратором с правом «Закреплять сообщения» (Pin messages), " + "затем отправьте в чат команду /pin_duty — текущее сообщение будет закреплено.", + ) + except (BadRequest, Forbidden): + pass + next_end = await loop.run_in_executor(None, _get_next_shift_end_utc) + await _schedule_next_update(context.application, chat_id, next_end) + return + + # Bot removed from group + if new.status in (ChatMemberStatus.LEFT, ChatMemberStatus.BANNED): + await asyncio.get_running_loop().run_in_executor(None, _sync_delete_pin, chat_id) + name = f"{JOB_NAME_PREFIX}{chat_id}" + if context.application.job_queue: + for job in context.application.job_queue.get_jobs_by_name(name): + job.schedule_removal() + logger.info("Bot left chat_id=%s, removed pin record and jobs", chat_id) + + +def _get_all_pin_chat_ids_sync() -> list[int]: + session = get_session(config.DATABASE_URL) + try: + return get_all_group_duty_pin_chat_ids(session) + finally: + session.close() + + +async def restore_group_pin_jobs(application) -> None: + """On startup: for each chat_id in group_duty_pins, schedule next update at shift end.""" + loop = asyncio.get_running_loop() + chat_ids = await loop.run_in_executor(None, _get_all_pin_chat_ids_sync) + for chat_id in chat_ids: + next_end = await loop.run_in_executor(None, _get_next_shift_end_utc) + await _schedule_next_update(application, chat_id, next_end) + logger.info("Restored %s group pin jobs", len(chat_ids)) + + +async def pin_duty_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: + """Pin the current duty message (use after granting the bot 'Pin messages' right).""" + if not update.message or not update.effective_chat: + return + chat = update.effective_chat + if chat.type not in ("group", "supergroup"): + await update.message.reply_text("Команда /pin_duty работает только в группах.") + return + chat_id = chat.id + loop = asyncio.get_running_loop() + message_id = await loop.run_in_executor(None, _sync_get_message_id, chat_id) + if message_id is None: + await update.message.reply_text("В этом чате ещё нет сообщения о дежурстве. Добавьте бота в группу — оно создастся автоматически.") + return + try: + await context.bot.pin_chat_message( + chat_id=chat_id, + message_id=message_id, + disable_notification=True, + ) + await update.message.reply_text("Сообщение о дежурстве закреплено.") + except (BadRequest, Forbidden) as e: + logger.warning("pin_duty failed chat_id=%s: %s", chat_id, e) + await update.message.reply_text( + "Не удалось закрепить. Убедитесь, что бот — администратор с правом «Закреплять сообщения»." + ) + + +group_duty_pin_handler = ChatMemberHandler( + my_chat_member_handler, + ChatMemberHandler.MY_CHAT_MEMBER, +) +pin_duty_handler = CommandHandler("pin_duty", pin_duty_cmd) diff --git a/main.py b/main.py index 19e133a..d659879 100644 --- a/main.py +++ b/main.py @@ -9,7 +9,7 @@ import urllib.request import config from telegram.ext import ApplicationBuilder -from handlers import register_handlers +from handlers import group_duty_pin, register_handlers logging.basicConfig( format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", @@ -63,7 +63,12 @@ def _run_uvicorn(web_app, port: int) -> None: def main() -> None: # Menu button (Календарь) and inline Calendar button are disabled; users open the app by link if needed. # _set_default_menu_button_webapp() - app = ApplicationBuilder().token(config.BOT_TOKEN).build() + app = ( + ApplicationBuilder() + .token(config.BOT_TOKEN) + .post_init(group_duty_pin.restore_group_pin_jobs) + .build() + ) register_handlers(app) from api.app import app as web_app @@ -76,7 +81,7 @@ def main() -> None: t.start() logger.info("Bot starting (polling)... HTTP API on port %s", config.HTTP_PORT) - app.run_polling(allowed_updates=["message"]) + app.run_polling(allowed_updates=["message", "my_chat_member"]) if __name__ == "__main__":