feat: add support for channels and supergroups (#289)

Signed-off-by: Eiko Wagenknecht <[email protected]>

README.md

@@ -31,6 +31,8 @@ For our mobile gamers:
 
 Want to receive only the offers *you* want in a single chat? Subscribe directly to the source: The [Telegram LootScraperBot](https://t.me/LootScraperBot) will happily send you push notifications for new offers. You can choose which categories you want to subscribe to.
 
+If you want, you can even add the bot to your own groups (including threaded groups) and channels. Just make sure to give it the neccessary permissions (admin rights work best).
+
 This is what it currently looks like in Telegram:
 
 ![image](https://user-images.githubusercontent.com/1475672/166058823-98e2beb9-7eb5-403d-94c7-7e17966fe9b7.png)

alembic.ini

@@ -60,12 +60,6 @@ version_path_separator = os  # Use os.pathsep. Default configuration used for ne
 # on newly generated revision scripts.  See the documentation for further
 # detail and examples
 
-# format using "black" - use the console_scripts runner, against the "black" entrypoint
-hooks = black
-black.type = console_scripts
-black.entrypoint = black
-black.options = -l 79 REVISION_SCRIPT_FILENAME
-
 # Logging configuration
 [loggers]
 keys = root,sqlalchemy,alembic

scripts/db_refresh.py

@@ -0,0 +1,43 @@
+# Script to copy over all data from the old database to the new one.
+# The new one should be generated with a run of lootscraper without an existing
+# database file.
+import sqlite3
+
+source_conn = sqlite3.connect("lootscraper.db")
+target_conn = sqlite3.connect("lootscraper_refreshed.db")
+
+source_cursor = source_conn.cursor()
+target_cursor = target_conn.cursor()
+
+# Disable foreign keys
+source_cursor.execute("PRAGMA foreign_keys = OFF;")
+target_cursor.execute("PRAGMA foreign_keys = OFF;")
+
+# Get the list of all tables in source database
+source_cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
+tables = source_cursor.fetchall()
+
+# Copy data from source to target database
+for table in tables:
+    table_name = table[0]
+    source_cursor.execute(f"SELECT * FROM {table_name};")  # noqa: S608
+
+    rows = source_cursor.fetchall()
+    if not rows:
+        continue
+
+    placeholders = ", ".join(["?"] * len(rows[0]))
+    target_cursor.executemany(
+        f"INSERT INTO {table_name} VALUES ({placeholders});",  # noqa: S608
+        rows,
+    )
+
+# Enable foreign keys
+source_cursor.execute("PRAGMA foreign_keys = ON;")
+target_cursor.execute("PRAGMA foreign_keys = ON;")
+
+source_conn.commit()
+target_conn.commit()
+
+source_conn.close()
+target_conn.close()

src/lootscraper/alembic/versions/20230202_100720_initial_sqlalchemy_2_0.py

@@ -174,13 +174,3 @@ def upgrade() -> None:
         ),
         sa.PrimaryKeyConstraint("id"),
     )
-
-
-def downgrade() -> None:
-    op.drop_table("offers")
-    op.drop_table("telegram_subscriptions")
-    op.drop_table("games")
-    op.drop_table("users")
-    op.drop_table("steam_info")
-    op.drop_table("igdb_info")
-    op.drop_table("announcements")

src/lootscraper/alembic/versions/20231027_115622_revamp_telegram.py

@@ -0,0 +1,150 @@
+"""
+Revamp Telegram.
+
+Revision ID: 8338a761b831
+Revises: 023117ae895e
+Create Date: 2023-10-27 11:56:22.363419+00:00
+
+"""
+import json
+import logging
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import orm
+from telegram.constants import ChatType
+
+from lootscraper.database import AwareDateTime, Base, TelegramChat
+
+# revision identifiers, used by Alembic.
+revision = "8338a761b831"
+down_revision = "023117ae895e"
+branch_labels = None
+depends_on = None
+
+
+class TempUser(Base):
+    __tablename__ = "users"
+
+    id = sa.Column(sa.Integer, primary_key=True)  # noqa: A003
+    registration_date = sa.Column(AwareDateTime, nullable=False)
+    telegram_id = sa.Column(sa.String, nullable=True)
+    telegram_chat_id = sa.Column(sa.String)
+    telegram_user_details = sa.Column(sa.JSON)
+    timezone_offset = sa.Column(sa.Integer)
+    inactive = sa.Column(sa.String, default=None)
+    offers_received_count = sa.Column(sa.Integer, default=0)
+    last_announcement_id = sa.Column(sa.Integer, default=0)
+
+
+def upgrade() -> None:
+    # First create the telegram_chats table
+    logging.info("Creating new telegram_chats table")
+    op.create_table(
+        "telegram_chats",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("registration_date", AwareDateTime(), nullable=False),
+        sa.Column("chat_type", sa.Enum(ChatType), nullable=False),
+        sa.Column("chat_id", sa.Integer(), nullable=False),
+        sa.Column("user_id", sa.Integer(), nullable=True),
+        sa.Column("thread_id", sa.Integer(), nullable=True),
+        sa.Column("chat_details", sa.JSON(), nullable=True),
+        sa.Column("user_details", sa.JSON(), nullable=True),
+        sa.Column("timezone_offset", sa.Integer(), nullable=False),
+        sa.Column("active", sa.Boolean(), nullable=False),
+        sa.Column("inactive_reason", sa.String(), nullable=True),
+        sa.Column("offers_received_count", sa.Integer(), nullable=False),
+        sa.Column("last_announcement_id", sa.Integer(), nullable=False),
+        sa.PrimaryKeyConstraint("id"),
+    )
+
+    # Fill the new table with data from the old one
+    bind = op.get_bind()
+    with orm.Session(bind=bind) as session:
+        seen_chat_ids = set()
+        for user in session.query(TempUser):
+            if user.telegram_chat_id in seen_chat_ids:
+                logging.warning(f"Skipping duplicate chat {user.telegram_chat_id}")
+                continue
+            seen_chat_ids.add(user.telegram_chat_id)
+            logging.info(f"Migrating user {user.id}")
+
+            # Determine chat type
+            if "Channel user created by admin" in user.telegram_user_details:
+                chat_type = ChatType.CHANNEL
+                user_details = None
+            elif int(user.telegram_chat_id) < 0:
+                chat_type = ChatType.GROUP
+                user_details = user.telegram_user_details
+            else:
+                chat_type = ChatType.PRIVATE
+                user_details = user.telegram_user_details
+
+            user_id = int(user.telegram_id) if int(user.telegram_id) > 0 else None
+
+            if user_details and not isinstance(user_details, dict):
+                user_details = json.loads(user_details)  # type: ignore
+
+            new_chat = TelegramChat(
+                registration_date=user.registration_date,
+                user_id=user_id,
+                chat_id=int(user.telegram_chat_id),
+                user_details=user_details,
+                chat_details=None,
+                timezone_offset=user.timezone_offset,
+                active=user.inactive is None,
+                inactive_reason=user.inactive,
+                offers_received_count=user.offers_received_count,
+                last_announcement_id=user.last_announcement_id,
+                chat_type=chat_type,
+            )
+            # Keep the primary key!
+            new_chat.id = user.id
+            session.add(new_chat)
+
+        session.commit()
+
+    # Drop all foreign keys workaround. They can't be dropped directly because
+    # they have no name.
+    conn = op.get_bind()
+    # Disable foreign key constraint temporarily
+    conn.execute(sa.text("PRAGMA foreign_keys=off;"))
+    conn.execute(
+        sa.text(
+            """
+            CREATE TABLE "new_telegram_subscriptions" (
+                "id"	INTEGER NOT NULL,
+                "chat_id"	INTEGER NOT NULL,
+                "source"	VARCHAR(7) NOT NULL,
+                "type"	VARCHAR(4) NOT NULL,
+                "last_offer_id"	INTEGER NOT NULL,
+                "duration"	VARCHAR(9) NOT NULL,
+                CONSTRAINT "fk_telegram_subscriptions_chat_id_telegram_chats"
+                FOREIGN KEY("chat_id")
+                REFERENCES "telegram_chats"("id"),
+                CONSTRAINT "pk_telegram_subscriptions"
+                PRIMARY KEY("id")
+                );
+            """,
+        ),
+    )
+    conn.execute(
+        sa.text(
+            """
+                INSERT INTO new_telegram_subscriptions
+                (id, chat_id, source, type, last_offer_id, duration)
+                SELECT id, user_id, source, type, last_offer_id, duration
+                FROM telegram_subscriptions;
+            """,
+        ),
+    )
+    conn.execute(sa.text("DROP TABLE telegram_subscriptions;"))
+    conn.execute(
+        sa.text(
+            "ALTER TABLE new_telegram_subscriptions RENAME TO telegram_subscriptions;",
+        ),
+    )
+    conn.execute(sa.text("PRAGMA foreign_keys=on;"))  # Enable back the foreign keys
+
+    # At last, drop the old users table
+    op.drop_table("users")

src/lootscraper/database.py

@@ -18,8 +18,15 @@
     scoped_session,
     sessionmaker,
 )
-
-from lootscraper.common import Category, Channel, OfferDuration, OfferType, Source
+from telegram.constants import ChatType
+
+from lootscraper.common import (
+    Category,
+    Channel,
+    OfferDuration,
+    OfferType,
+    Source,
+)
 from lootscraper.config import Config
 from lootscraper.utils import calc_real_valid_to
 
@@ -31,12 +38,23 @@
 logger = logging.getLogger(__name__)
 
 
-# mapper_registry = orm.registry()
+# Naming convention for keys and constraints
+convention = {
+    "ix": "ix_%(column_0_label)s",
+    "uq": "uq_%(table_name)s_%(column_0_name)s",
+    "ck": "ck_%(table_name)s_%(constraint_name)s",
+    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+    "pk": "pk_%(table_name)s",
+}
+
+metadata_obj = sa.MetaData(naming_convention=convention)
 
 
 class Base(MappedAsDataclass, DeclarativeBase):
     """Subclasses will be converted to dataclasses."""
 
+    metadata = metadata_obj
+
 
 class AwareDateTime(sa.TypeDecorator):  # type: ignore
     """Results returned as aware datetimes, not naive ones."""
@@ -205,45 +223,49 @@ def real_valid_to(self) -> datetime | None:
         return calc_real_valid_to(self.seen_last, self.valid_to)
 
 
-class User(Base):
-    """A user of the application."""
+class TelegramChat(Base):
+    """A Telegram chat. Can be a single user, a group or a channel."""
 
-    __tablename__ = "users"
+    __tablename__ = "telegram_chats"
 
-    telegram_subscriptions: Mapped[list[TelegramSubscription]] = relationship(
+    subscriptions: Mapped[list[TelegramSubscription]] = relationship(
         "TelegramSubscription",
-        back_populates="user",
+        back_populates="chat",
         cascade="all, delete-orphan",
         init=False,
     )
 
-    id: Mapped[int] = mapped_column(  # noqa: A003
-        init=False,
-        primary_key=True,
-    )
+    id: Mapped[int] = mapped_column(init=False, primary_key=True)  # noqa: A003
     registration_date: Mapped[datetime] = mapped_column(AwareDateTime)
-    telegram_id: Mapped[str | None]
-    telegram_chat_id: Mapped[str]
-    telegram_user_details: Mapped[dict[str, Any] | None] = mapped_column(sa.JSON)
+    chat_type: Mapped[ChatType] = mapped_column(sa.Enum(ChatType))
+    chat_id: Mapped[int]
+    user_id: Mapped[int | None] = mapped_column(default=None)
+    thread_id: Mapped[int | None] = mapped_column(default=None)
+    chat_details: Mapped[dict[str, Any] | None] = mapped_column(sa.JSON, default=None)
+    user_details: Mapped[dict[str, Any] | None] = mapped_column(sa.JSON, default=None)
     timezone_offset: Mapped[int] = mapped_column(default=0)
-    inactive: Mapped[str | None] = mapped_column(default=None)
+    active: Mapped[bool] = mapped_column(default=True)
+    inactive_reason: Mapped[str | None] = mapped_column(default=None)
     offers_received_count: Mapped[int] = mapped_column(default=0)
     last_announcement_id: Mapped[int] = mapped_column(default=0)
 
 
 class TelegramSubscription(Base):
-    """Subscription of a user to a category for Telegram notifications."""
+    """Subscription of a chat to a category for Telegram notifications."""
 
     __tablename__ = "telegram_subscriptions"
 
-    user: Mapped[User] = relationship("User", back_populates="telegram_subscriptions")
+    chat: Mapped[TelegramChat] = relationship(
+        "TelegramChat",
+        back_populates="subscriptions",
+    )
 
     id: Mapped[int] = mapped_column(  # noqa: A003
         init=False,
         primary_key=True,
     )
-    user_id: Mapped[int] = mapped_column(
-        sa.ForeignKey("users.id"),
+    chat_id: Mapped[int] = mapped_column(
+        sa.ForeignKey("telegram_chats.id"),
         init=False,
     )
     source: Mapped[Source] = mapped_column(sa.Enum(Source))

src/lootscraper/main.py

@@ -14,7 +14,7 @@
 from lootscraper import __version__
 from lootscraper.browser import get_browser_context
 from lootscraper.common import TIMESTAMP_LONG
-from lootscraper.config import Config
+from lootscraper.config import Config, TelegramLogLevel
 from lootscraper.database import LootDatabase
 from lootscraper.processing import (
     action_generate_feed,
@@ -138,7 +138,12 @@ def setup_logging() -> None:
     handlers.append(stream_handler)
 
     # Create a rotating log file, size: 5 MB, keep 10 files
-    file_handler = RotatingFileHandler(filename, maxBytes=5 * 1024**2, backupCount=10)
+    file_handler = RotatingFileHandler(
+        filename,
+        maxBytes=5 * 1024**2,
+        backupCount=10,
+        encoding="utf-8",
+    )
     file_handler.setFormatter(logging.Formatter(LOGFORMAT))
     handlers.append(file_handler)
 
@@ -162,11 +167,21 @@ async def run_telegram_bot(
     async with TelegramBot(Config.get(), db) as bot:
         # The bot is running now and will stop when the context exits
         try:
-            telegram_handler = TelegramLoggingHandler(bot)
-            # Only log errors and above to Telegram. TODO: Make this configurable.
-            telegram_handler.setLevel(logging.ERROR)
-            telegram_handler.setFormatter(logging.Formatter(LOGFORMAT))
-            logger.addHandler(telegram_handler)
+            lvl = Config.get().telegram_log_level
+            if lvl != TelegramLogLevel.DISABLED:
+                telegram_handler = TelegramLoggingHandler(bot)
+                if lvl == TelegramLogLevel.DEBUG:
+                    telegram_handler.setLevel(logging.DEBUG)
+                elif lvl == TelegramLogLevel.INFO:
+                    telegram_handler.setLevel(logging.INFO)
+                elif lvl == TelegramLogLevel.WARNING:
+                    telegram_handler.setLevel(logging.WARNING)
+                elif lvl == TelegramLogLevel.ERROR:
+                    telegram_handler.setLevel(logging.ERROR)
+
+                telegram_handler.setLevel(logging.ERROR)
+                telegram_handler.setFormatter(logging.Formatter(LOGFORMAT))
+                logger.addHandler(telegram_handler)
         except Exception:
             logger.exception("Could not add Telegram logging handler.")