Nikolay Tatarinov
5331fac334
- Introduced a new configuration file `.cursorrules` to define coding standards, error handling, testing requirements, and project-specific guidelines. - Refactored `config.py` to implement a `Settings` dataclass for better management of environment variables, improving testability and maintainability. - Updated the import duty schedule handler to utilize session management with `session_scope`, ensuring proper database session handling. - Enhanced the import service to streamline the duty schedule import process, improving code organization and readability. - Added new service layer functions to encapsulate business logic related to group duty pinning and duty schedule imports. - Updated README documentation to reflect the new configuration structure and improved import functionality.
236 lines
9.0 KiB
Python
236 lines
9.0 KiB
Python
"""Pinned duty message in groups: handle bot add/remove, schedule updates at shift end."""
|
||
|
||
import asyncio
|
||
import logging
|
||
from datetime import datetime, timezone
|
||
|
||
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 session_scope
|
||
from services.group_duty_pin_service import (
|
||
get_duty_message_text,
|
||
get_next_shift_end_utc,
|
||
save_pin,
|
||
delete_pin,
|
||
get_message_id,
|
||
get_all_pin_chat_ids,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
JOB_NAME_PREFIX = "duty_pin_"
|
||
RETRY_WHEN_NO_DUTY_MINUTES = 15
|
||
|
||
|
||
def _get_duty_message_text_sync() -> str:
|
||
"""Get current duty message (sync, for run_in_executor)."""
|
||
with session_scope(config.DATABASE_URL) as session:
|
||
return get_duty_message_text(session, config.DUTY_DISPLAY_TZ)
|
||
|
||
|
||
def _get_next_shift_end_sync():
|
||
"""Return next shift end as naive UTC (sync, for run_in_executor)."""
|
||
with session_scope(config.DATABASE_URL) as session:
|
||
return get_next_shift_end_utc(session)
|
||
|
||
|
||
def _sync_save_pin(chat_id: int, message_id: int) -> None:
|
||
with session_scope(config.DATABASE_URL) as session:
|
||
save_pin(session, chat_id, message_id)
|
||
|
||
|
||
def _sync_delete_pin(chat_id: int) -> None:
|
||
with session_scope(config.DATABASE_URL) as session:
|
||
delete_pin(session, chat_id)
|
||
|
||
|
||
def _sync_get_message_id(chat_id: int) -> int | None:
|
||
with session_scope(config.DATABASE_URL) as session:
|
||
return get_message_id(session, chat_id)
|
||
|
||
|
||
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:
|
||
now_utc = datetime.now(timezone.utc).replace(tzinfo=None)
|
||
delay = when_utc - now_utc
|
||
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_sync)
|
||
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_sync)
|
||
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_sync)
|
||
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_sync)
|
||
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]:
|
||
with session_scope(config.DATABASE_URL) as session:
|
||
return get_all_pin_chat_ids(session)
|
||
|
||
|
||
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_sync)
|
||
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)
|