Merge pull request #2 from Ovler-Young/More-Topics

Ovler-Young · web-flow · commit 9de5140422fd · 2025-05-12T08:01:56.000-04:00
More topics
diff --git a/efb_telegram_master/bot_manager.py b/efb_telegram_master/bot_manager.py
@@ -546,6 +546,11 @@ def answer_callback_query(self, *args, prefix="", suffix="", text=None,
             *args, text=prefix + text + suffix, **kwargs
         )
 
+    @Decorators.retry_on_timeout
+    @Decorators.retry_on_chat_migration
+    def get_chat_info(self, *args, **kwargs):
+        return self.updater.bot.get_chat(*args, **kwargs)
+
     @Decorators.retry_on_timeout
     @Decorators.retry_on_chat_migration
     def create_forum_topic(self, *args, **kwargs):
diff --git a/efb_telegram_master/chat_binding.py b/efb_telegram_master/chat_binding.py
@@ -223,6 +223,22 @@ def link_chat_show_list(self, update: Update, context: CallbackContext):
                 self.link_handler.conversations[storage_id] = Flags.LINK_EXEC
                 self.msg_storage[storage_id] = ChatListStorage([chat])
                 return self.build_link_action_message(chat, tg_chat_id, tg_msg_id)
+            if message.message_thread_id:
+                topic = message.message_thread_id
+                if topic:
+                    slave_origin_uid = self.db.get_topic_slave(
+                        topic_chat_id=TelegramChatID(message.chat_id),
+                        message_thread_id=topic
+                    )
+                    if slave_origin_uid:
+                        channel_id, chat_id, _ = utils.chat_id_str_to_id(slave_origin_uid)
+                        chat: ETMChatType = self.chat_manager.get_chat(channel_id, chat_id, build_dummy=True)
+                        tg_chat_id = TelegramChatID(message.chat_id)
+                        tg_msg_id = TelegramMessageID(message.reply_text(self._("Processing...")).message_id)
+                        storage_id: Tuple[TelegramChatID, TelegramMessageID] = (tg_chat_id, tg_msg_id)
+                        self.link_handler.conversations[storage_id] = Flags.LINK_EXEC
+                        self.msg_storage[storage_id] = ChatListStorage([chat])
+                        return self.build_link_action_message(chat, tg_chat_id, tg_msg_id)
 
         if message.chat.type != telegram.Chat.PRIVATE:
             links = self.db.get_chat_assoc(
@@ -567,15 +583,28 @@ def link_chat(self, update: Update, args: Optional[List[str]]):
         msg = self.bot.send_message(tg_chat_to_link, text=txt)
 
         chat.link(self.channel.channel_id, ChatID(str(tg_chat_to_link)), self.channel.flag("multiple_slave_chats"))
-        try:
-            self.db.remove_topic_assoc(
-                topic_chat_id=self.channel.topic_group,
-                slave_uid=chat_uid,
-            )
-            #TODO delete or close the topic after link?
-        except Exception as e:
-            self.logger.warn("Error occurred while remove topic assoc.\nError: %s\n%s",
-                              repr(e), traceback.format_exc())
+        self.db.remove_topic_assoc(
+            slave_uid=chat_uid,
+        )
+
+        # chat_id = utils.chat_id_to_str(self.channel.channel_id, ChatID(str(tg_chat_to_link)))
+        # links = self.db.get_chat_assoc(master_uid=chat_id)
+        # if len(links) > 1 and self.channel.topic_group:
+        #     try:
+        #         topic: ForumTopic = self.bot.create_forum_topic(
+        #             chat_id=chat_id,
+        #             name=chat.chat_title
+        #         )
+        #         thread_id = topic.message_thread_id
+        #         self.db.add_topic_assoc(
+        #             topic_chat_id=chat_id,
+        #             message_thread_id=thread_id,
+        #             slave_uid=chat_uid,
+        #         )
+        #         return thread_id
+        #     except telegram.error.BadRequest as e:
+        #         self.logger.error('Failed to create topic, Reason: %s', e)
+        #         return None
 
         txt = self._("Chat {0} is now linked.").format(chat_display_name)
         self.bot.edit_message_text(text=txt, chat_id=msg.chat.id, message_id=msg.message_id)
diff --git a/efb_telegram_master/db.py b/efb_telegram_master/db.py
@@ -22,7 +22,7 @@
 from .message import ETMMsg
 from .msg_type import TGMsgType
 from .utils import TelegramChatID, EFBChannelChatIDStr, TgChatMsgIDStr, message_id_to_str, \
-    chat_id_to_str, OldMsgID, chat_id_str_to_id, TelegramMessageID
+    chat_id_to_str, OldMsgID, chat_id_str_to_id, TelegramMessageID, TelegramTopicID
 
 if TYPE_CHECKING:
     from . import TelegramChannel
@@ -72,8 +72,6 @@ class MsgLog(BaseModel):
     """Editable message ID from Telegram if ``master_msg_id`` is not editable
     and a separate one is sent.
     """
-    master_message_thread_id = TextField(null=True)
-    """Message thread ID from Telegram"""
     slave_message_id = TextField()
     """Message from slave channel."""
     text = TextField()
@@ -204,8 +202,6 @@ def __init__(self, channel: 'TelegramChannel'):
                 self._migrate(2)
             elif "file_unique_id" not in msg_log_columns:
                 self._migrate(3)
-            elif "master_message_thread_id" not in msg_log_columns:
-                self._migrate(4)
         self.logger.debug("Database migration finished...")
 
     def stop_worker(self):
@@ -256,12 +252,6 @@ def _migrate(i: int):
             migrate(
                 migrator.add_column("msglog", "file_unique_id", MsgLog.file_unique_id)
             )
-        if i <= 4:
-            # Migration 4: Add column for message thread ID to message log table
-            # 2025APR12
-            migrate(
-                migrator.add_column("msglog", "master_message_thread_id", MsgLog.master_message_thread_id)
-            )
 
     def add_chat_assoc(self, master_uid: EFBChannelChatIDStr,
                        slave_uid: EFBChannelChatIDStr,
@@ -389,80 +379,111 @@ def get_chat_assoc(master_uid: Optional[EFBChannelChatIDStr] = None,
         except DoesNotExist:
             return []
 
-    def add_topic_assoc(self, message_thread_id: EFBChannelChatIDStr,
-                       slave_uid: EFBChannelChatIDStr, 
-                       topic_chat_id: int):
+    def add_topic_assoc(self, topic_chat_id: TelegramChatID,
+                       message_thread_id: EFBChannelChatIDStr,
+                       slave_uid: EFBChannelChatIDStr, ):
         """
         Add topic associations (topic links).
         One Master channel with many Slave channel.
 
         Args:
-            message_thread_id (str): thread UID in topic
-            slave_uid (str): Slave channel UID ("%(channel_id)s.%(chat_id)s")
+            topic_chat_id (TelegramChatID): The topic group chat ID
+            message_thread_id (EFBChannelChatIDStr): The topic thread ID
+            slave_uid (EFBChannelChatIDStr): Slave channel UID ("%(channel_id)s.%(chat_id)s")
         """
         return TopicAssoc.create(topic_chat_id=topic_chat_id, message_thread_id=message_thread_id, slave_uid=slave_uid)
 
     @staticmethod
-    def get_topic_thread_id(topic_chat_id: int,
-                       slave_uid: Optional[EFBChannelChatIDStr]
-                        ) -> int:
+    def get_topic_thread_id(slave_uid: EFBChannelChatIDStr) -> TelegramTopicID:
         """
         Get topic association (topic link) information.
         Only one parameter is to be provided.
 
         Args:
-            topic_chat_id (int): The topic UID
-            slave_uid (str): Slave channel UID ("%(channel_id)s.%(chat_id)s")
+            slave_uid (EFBChannelChatIDStr): Slave channel UID ("%(channel_id)s.%(chat_id)s")
 
         Returns:
             The message thread_id
         """
         try:
             assoc = TopicAssoc.select(TopicAssoc.message_thread_id)\
-                .where(TopicAssoc.slave_uid == slave_uid, TopicAssoc.topic_chat_id == topic_chat_id)\
+                .where(TopicAssoc.slave_uid == slave_uid)\
                 .order_by(TopicAssoc.id.desc()).first()
             if assoc:
                 return int(assoc.message_thread_id)
         except DoesNotExist:
             return None
 
     @staticmethod
-    def get_topic_slave(topic_chat_id: int,
-                        message_thread_id: int
+    def get_topic_slave(topic_chat_id: TelegramTopicID,
+                        message_thread_id: Optional[EFBChannelChatIDStr] = None,
                         ) -> Optional[EFBChannelChatIDStr]:
         """
         Get topic association (topic link) information.
         Only one parameter is to be provided.
 
         Args:
-            topic_chat_id (int): The topic UID
-            message_thread_id (int): The message thread ID
+            topic_chat_id (TelegramTopicID): The topic UID
+            message_thread_id (TelegramTopicID): The message thread ID
 
         Returns:
             Slave channel UID ("%(channel_id)s.%(chat_id)s")
         """
         try:
-            return TopicAssoc.select(TopicAssoc.slave_uid)\
-                .where(TopicAssoc.message_thread_id == message_thread_id, TopicAssoc.topic_chat_id == topic_chat_id).first().slave_uid
+            if message_thread_id:
+                return TopicAssoc.select(TopicAssoc.slave_uid)\
+                    .where(TopicAssoc.message_thread_id == message_thread_id, TopicAssoc.topic_chat_id == topic_chat_id).first().slave_uid
+            else:
+                return TopicAssoc.select(TopicAssoc.slave_uid)\
+                    .where(TopicAssoc.topic_chat_id == topic_chat_id).first().slave_uid
         except DoesNotExist:
             return None
-        except AttributeError: # Handle case where .slave_uid doesn't exist on the result
+        except AttributeError:
             return None
 
     @staticmethod
-    def remove_topic_assoc(topic_chat_id: int, slave_uid: Optional[EFBChannelChatIDStr] = None):
+    def get_topic_slaves(topic_chat_id: TelegramChatID) -> Optional[List[Tuple[EFBChannelChatIDStr, TelegramTopicID]]]:
+        """
+        Get topic association (topic link) information.
+        Only one parameter is to be provided.
+
+        Args:
+            topic_chat_id (TelegramChatID): The topic UID
+
+        Returns:
+            List[Tuple[EFBChannelChatIDStr, TelegramTopicID]]: A list of tuples containing slave channel UID and message thread ID
+        """
+        try:
+            query = TopicAssoc.select(TopicAssoc.slave_uid, TopicAssoc.message_thread_id)\
+                .where(TopicAssoc.topic_chat_id == topic_chat_id).order_by(TopicAssoc.id.desc())
+            return [(row.slave_uid, int(row.message_thread_id)) for row in query]
+        except DoesNotExist:
+            return None
+        except AttributeError:
+            return None
+
+    @staticmethod
+    def remove_topic_assoc(topic_chat_id: Optional[TelegramChatID] = None,
+                           message_thread_id: Optional[EFBChannelChatIDStr] = None,
+                           slave_uid: Optional[EFBChannelChatIDStr] = None):
         """
         Remove topic association (topic link).
 
         Args:
-            topic_chat_id (int): The topic group chat ID
-            slave_uid (str): Slave channel UID ("%(channel_id)s.%(chat_id)s")
+            topic_chat_id (TelegramChatID): The topic group chat ID
+            message_thread_id (EFBChannelChatIDStr): The topic thread ID
+            slave_uid (EFBChannelChatIDStr): Slave channel UID ("%(channel_id)s.%(chat_id)s")
         """
         try:
-            return TopicAssoc.delete().where(
-                (TopicAssoc.topic_chat_id == str(topic_chat_id)) &
-                (TopicAssoc.slave_uid == str(slave_uid))
-            ).execute()
+            if bool(topic_chat_id and message_thread_id) == bool(slave_uid):
+                raise ValueError("Please provide either topic_chat_id and message_thread_id or slave_uid.")
+            elif topic_chat_id and message_thread_id:
+                return TopicAssoc.delete().where(
+                    (TopicAssoc.topic_chat_id == str(topic_chat_id)) &
+                    (TopicAssoc.message_thread_id == str(message_thread_id))
+                ).execute()
+            elif slave_uid:
+                return TopicAssoc.delete().where(TopicAssoc.slave_uid == slave_uid).execute()
         except DoesNotExist:
             return 0
 
@@ -494,7 +515,6 @@ def add_or_update_message_log(self,
 
         row.master_msg_id = master_msg_id
         row.master_msg_id_alt = master_msg_id_alt
-        row.master_message_thread_id = str(master_message.message_thread_id) if master_message.message_thread_id else None
         row.text = msg.text
         row.slave_origin_uid = chat_id_to_str(chat=msg.chat)
         row.slave_member_uid = chat_id_to_str(chat=msg.author)
diff --git a/efb_telegram_master/master_message.py b/efb_telegram_master/master_message.py
@@ -160,22 +160,36 @@ def msg(self, update: Update, context: CallbackContext):
             if destination:
                 quote = message.reply_to_message is not None
                 self.logger.debug("[%s] Chat %s is singly-linked to %s", mid, message.chat, destination)
+                if message.chat.is_forum:
+                    ideal_thread_id = self.db.get_topic_thread_id(slave_uid=destination)
+                    if ideal_thread_id and ideal_thread_id != message.message_thread_id:
+                        self.logger.debug("[%s] Chat %s is singly-linked to %s, but the thread ID is not matching.", mid, message.chat, destination)
+                        destination = None
+                        quote = False
+                        return
 
         if destination is None:
-            thread_id = message.message_thread_id
-            if thread_id:
-                if TelegramChatID(update.effective_chat.id) == self.channel.topic_group:
-                    destination = self.db.get_topic_slave(message_thread_id=thread_id, topic_chat_id=self.channel.topic_group)
-                    if destination:
-                        quote = message.reply_to_message.message_id != message.reply_to_message.message_thread_id
-                        if not quote:
-                            message.reply_to_message = None
-                    else:
+            if message.chat.is_forum:
+                topic_destinations = self.db.get_topic_slaves(topic_chat_id=message.chat.id)
+                thread_id = message.message_thread_id
+                if thread_id:
+                    for (destination, topic_id) in topic_destinations:
+                        if topic_id == thread_id:
+                            self.logger.debug("[%s] Chat %s is singly-linked to %s in topic %s", mid, message.chat, destination, topic_id)
+                            destination = destination
+                            quote = message.reply_to_message.message_id != message.reply_to_message.message_thread_id
+                            if quote:
+                                message.reply_to_message = None
+                            break
+                    if destination is None:
                         self.logger.debug("[%s] Ignored message as it's a topic which wasn't created by this bot", mid)
                         return
                 else:
-                    self.logger.debug("[%s] Ignored message as it's a invalid topic group.", mid)
-                    return
+                    self.logger.debug("[%s] Chat %s is a forum, but no thread ID is found.", mid, message.chat)
+                    destinations = self.db.get_chat_assoc(master_uid=utils.chat_id_to_str(self.channel_id, ChatID(str(message.chat.id))))
+                    if len(destinations) == len(topic_destinations):
+                        self.logger.debug("[%s] Chat %s is a forum, and all destinations are in topics. The new message is not in any topic, so ignore it.", mid, message.chat)
+                        return
 
         if destination is None:  # not singly linked
             quote = False
diff --git a/efb_telegram_master/slave_message.py b/efb_telegram_master/slave_message.py
diff --git a/efb_telegram_master/utils.py b/efb_telegram_master/utils.py
