From 5250a0ce54c48bd72372ab1acf34363879b2be06 Mon Sep 17 00:00:00 2001 From: Johnny Gear Date: Thu, 19 Mar 2026 12:38:17 -0500 Subject: [PATCH] Skip Days - Order Pool specific --- db/models.py | 8 +++++ db/queries.py | 13 ++++--- migrations/027_skip_day_add_pool.py | 51 ++++++++++++++++++++++++++ scheduling.py | 8 +++-- telegram/commands.py | 56 +++++++++++++++++++++++++---- telegram/telegram.py | 40 +++++++++++++++++++-- 6 files changed, 161 insertions(+), 15 deletions(-) create mode 100644 migrations/027_skip_day_add_pool.py diff --git a/db/models.py b/db/models.py index 2bbba87..b9eed6a 100644 --- a/db/models.py +++ b/db/models.py @@ -210,6 +210,14 @@ class SkipDay(BaseModel): backref='skip_days' ) + pool = ForeignKeyField( + column_name='orders_pool_id', + field='id', + model=OrdersPool, + backref='skip_days', + null=True + ) + class Meta: table_name = 'skip_day' diff --git a/db/queries.py b/db/queries.py index 4d75065..8e37e20 100644 --- a/db/queries.py +++ b/db/queries.py @@ -121,8 +121,12 @@ def repeat_clear(id): q = Repeat.delete() q.execute() -def skip_day_put(user, date): - return SkipDay.create(user=user, date=date) +def skip_day_put(user, date, order_pool=None): + return SkipDay.create( + user=user, + date=date, + pool=order_pool + ) def skip_day_delete(user, date): q = SkipDay.delete().where((SkipDay.user == user) & (SkipDay.date == date)) @@ -134,9 +138,8 @@ def skip_days_upcoming(user): .order_by(SkipDay.date) .limit(10)) -def skip_day_contains(user, date): - q = SkipDay.select().where((SkipDay.user == user) & (SkipDay.date == date)) - return len(q) > 0 +def skip_day_get(user, date): + return SkipDay.select().where((SkipDay.user == user) & (SkipDay.date == date)).get_or_none() def order_status_put(orders_pool, user, mastodon_id, created_at, due_at, text, punishment_for=None, verify_at=None): return OrderStatus.create( diff --git a/migrations/027_skip_day_add_pool.py b/migrations/027_skip_day_add_pool.py new file mode 100644 index 0000000..d0ac246 --- /dev/null +++ b/migrations/027_skip_day_add_pool.py @@ -0,0 +1,51 @@ +"""Peewee migrations -- 027_skip_day_add_pool.py. + +Some examples (model - class or model name):: + + > Model = migrator.orm['table_name'] # Return model in current state by name + > Model = migrator.ModelClass # Return model in current state by name + + > migrator.sql(sql) # Run custom SQL + > migrator.run(func, *args, **kwargs) # Run python function with the given args + > migrator.create_model(Model) # Create a model (could be used as decorator) + > migrator.remove_model(model, cascade=True) # Remove a model + > migrator.add_fields(model, **fields) # Add fields to a model + > migrator.change_fields(model, **fields) # Change fields + > migrator.remove_fields(model, *field_names, cascade=True) + > migrator.rename_field(model, old_field_name, new_field_name) + > migrator.rename_table(model, new_table_name) + > migrator.add_index(model, *col_names, unique=False) + > migrator.add_not_null(model, *field_names) + > migrator.add_default(model, field_name, default) + > migrator.add_constraint(model, name, sql) + > migrator.drop_index(model, *col_names) + > migrator.drop_not_null(model, *field_names) + > migrator.drop_constraints(model, *constraints) + +""" + +from contextlib import suppress + +import peewee as pw +from peewee_migrate import Migrator + + +with suppress(ImportError): + import playhouse.postgres_ext as pw_pext + + +def migrate(migrator: Migrator, database: pw.Database, *, fake=False): + """Write your migrations here.""" + + migrator.add_fields( + 'skip_day', + + pool=pw.ForeignKeyField(column_name='orders_pool_id', field='id', model=migrator.orm['orders_pool'], null=True)) + + +def rollback(migrator: Migrator, database: pw.Database, *, fake=False): + """Write your rollback migrations here.""" + + migrator.remove_fields('skip_day', 'pool') + + migrator.drop_index('skip_day', 'pool') diff --git a/scheduling.py b/scheduling.py index e48034b..916f59a 100644 --- a/scheduling.py +++ b/scheduling.py @@ -7,7 +7,7 @@ from scheduler.asyncio import Scheduler from db.constants import TIMELINE_ORDER_ISSUED, TIMELINE_ORDER_NOT_ISSUED from settings import TIMEZONE, SCHEDULE_SYNC_INTERVAL from orders import order_issue, order_check -from db.queries import orders_pool_by_id, orders_pool_scheduled, orders_pool_since, skip_day_contains, order_status_outstanding, timeline_event_put +from db.queries import orders_pool_by_id, orders_pool_scheduled, orders_pool_since, skip_day_get, order_status_outstanding, timeline_event_put from util import order_time, sqlite_time logger = logging.getLogger(__name__) @@ -76,7 +76,11 @@ class OrderScheduler(): # Skip stored dates today = datetime.datetime.now(tz=self.tz).strftime("%Y-%m-%d") - if (skip_day_contains(orders_pool.user, today)): + skip_day = skip_day_get(orders_pool.user, today) + if (skip_day is not None and ( + skip_day.pool is None or + skip_day.pool.id == orders_pool.id + )): logger.info(f'{orders_pool} Today is a skip day') timeline_event_put( TIMELINE_ORDER_NOT_ISSUED, diff --git a/telegram/commands.py b/telegram/commands.py index 64fd957..fce58a1 100644 --- a/telegram/commands.py +++ b/telegram/commands.py @@ -4,11 +4,14 @@ import peewee import datetime from db.queries import skip_day_put, skip_day_delete, skip_days_upcoming, user_add, user_get, domsubusers_add, domsubusers_delete +from db.models import OrdersPool from .telegram import Telegram, TelegramCommand from settings import FLASK_URL logger = logging.getLogger(__name__) +CALLBACK_PREFIX_SKIP_DAY_ADD = "skip_day_add_" + class StartCommand(TelegramCommand): command_text = '/start' description = "Start a session with the bot" @@ -88,14 +91,48 @@ class SkipDayAddCommand(TelegramCommand): try: yield "Please enter a skip day" - response = await forResponse() + skip_day_str = await forResponse() + + try: + skip_day = datetime.date.fromisoformat(skip_day_str) + except: + yield "The date you entered was not valid" + return user = user_get(update['message']['chat']['username']) - skip_day = datetime.date.fromisoformat(response) - skip_day_put(user, skip_day) + inline_keyboard = [{ + 'text': 'All', + 'callback_data': f"{CALLBACK_PREFIX_SKIP_DAY_ADD}_all" + }] + [ + { + 'text': pool.name, + 'callback_data': f"{CALLBACK_PREFIX_SKIP_DAY_ADD}_{pool.id}" + } for pool in user.orders_pools + ] - yield f"Skip day {response} has been added" + yield "Please choose an order pool", { + "inline_keyboard": [inline_keyboard] + } + + pool_str = await forResponse() + + if not pool_str.startswith("skip_day_add_"): + yield "There was a problem with your command" + return + + try: + order_pool_id = int(pool_str[len(CALLBACK_PREFIX_SKIP_DAY_ADD)+1:]) + order_pool = user.orders_pools.where(OrdersPool.id == order_pool_id).get() + except: + order_pool = None + + skip_day_put(user, skip_day, order_pool) + + if order_pool is not None: + yield f"Skip day {skip_day_str} for pool {order_pool.name} has been added" + else: + yield f"Skip day {skip_day_str} has been added" except ValueError: yield "That date was not valid" except peewee.IntegrityError: @@ -130,8 +167,15 @@ class SkipDaysCommand(TelegramCommand): async def exec_inner(self, text, update, session, forResponse=None, reply=None): user = user_get(update['message']['chat']['username']) dates = skip_days_upcoming(user) - yield ("Upcoming skip days -\n" + - "\n".join([d.date.isoformat() for d in dates])) + response = "Upcoming skip days -\n" + + for date in dates: + if date.pool is None: + response += date.date.isoformat() + "\n" + else: + response += f"{date.date.isoformat()} - {date.pool.name}\n" + + yield response commands = [ StartCommand(), diff --git a/telegram/telegram.py b/telegram/telegram.py index 9175fcf..91c98af 100644 --- a/telegram/telegram.py +++ b/telegram/telegram.py @@ -39,14 +39,24 @@ class Telegram: return await response.json() - async def message_send(self, chat_id, text): + async def message_send(self, chat_id, text, reply_markup=None): data = { 'chat_id': chat_id, 'text': text } + + if reply_markup is not None: + data['reply_markup'] = reply_markup return await self.post('sendMessage', data=data) + async def message_edit_reply_markup(self, chat_id, message_id, reply_markup): + return await self.post('editMessageReplyMarkup', data={ + 'chat_id': chat_id, + 'message_id': message_id, + 'reply_markup': reply_markup + }) + async def set_commands(self, commands): data = { 'commands': commands @@ -139,7 +149,12 @@ async def handle_commands(commands=[], loop=None): session, forResponse=forResponse ): - await t.message_send(chat_id, reply) + match reply.__class__.__name__: + case 'str': + await t.message_send(chat_id, reply) + case 'tuple': + (text, reply_markup) = reply + await t.message_send(chat_id, text, reply_markup=reply_markup) except TimeoutError: if chat_id in command_futures: del command_futures[chat_id] @@ -157,6 +172,27 @@ async def handle_commands(commands=[], loop=None): task.add_done_callback(command_tasks.discard) break + elif('callback_query' in update): + chat_id = update['callback_query']['message']['chat']['id'] + + if(update['callback_query']['message']['chat']['type'] != "private" or + update['callback_query']['message']['chat']['username'] not in TELEGRAM_ALLOWLIST): + continue + + if(chat_id in command_futures): + command_futures[chat_id].set_result( + update['callback_query']['data'] + ) + del command_futures[chat_id] + + # Clear inline keyboard + await t.message_edit_reply_markup( + chat_id, + update['callback_query']['message']['message_id'], + { "inline_keyboard": [] } + ) + continue + offset = update['update_id'] + 1 else: logger.warning("getUpdates was not ok {}".format(update))