Skip to content

Commit

Permalink
chore: cleanup
Browse files Browse the repository at this point in the history
  • Loading branch information
jackstar12 committed Jan 16, 2025
1 parent 0166c0c commit 2fcbc51
Show file tree
Hide file tree
Showing 3 changed files with 5 additions and 41 deletions.
6 changes: 3 additions & 3 deletions bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@

import aiohttp
from pydantic import ValidationError
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from telegram import Update, Bot
from telegram.ext import Application, CommandHandler, ContextTypes
from settings import Settings

from db import init_db, add_subscriber, remove_subscriber, get_subscribers, Base
from db import add_subscriber, remove_subscriber, get_subscribers, Base

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(message)s")

Expand Down Expand Up @@ -71,7 +71,7 @@ def main():
try:
settings = Settings()

engine = init_db(settings.database_url)
engine = create_async_engine(settings.database_url)
async_session = async_sessionmaker(engine, expire_on_commit=False)

application = Application.builder().token(settings.telegram_bot_token).build()
Expand Down
39 changes: 1 addition & 38 deletions db.py
Original file line number Diff line number Diff line change
@@ -1,39 +1,17 @@
from sqlalchemy import Column, Integer
from sqlalchemy import select
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.ext.declarative import declarative_base

# SQLAlchemy setup
Base = declarative_base()


class Subscriber(Base):
"""
SQLAlchemy model for the subscribers table.
"""

__tablename__ = "subscribers"
chat_id = Column(Integer, primary_key=True, unique=True)


def init_db(database_url: str):
"""
Initialize the database and create the subscribers table if it doesn't exist.
"""
engine = create_async_engine(database_url, echo=True)
# async with engine.begin() as conn:
# await conn.run_sync(Base.metadata.create_all)
return engine


# Database operations
async def add_subscriber(session: AsyncSession, chat_id: int) -> bool:
"""
Add a subscriber to the database.
Returns:
bool: True if the subscriber was added, False if already subscribed.
"""
result = await session.get(Subscriber, chat_id)
if result:
return False
Expand All @@ -44,12 +22,6 @@ async def add_subscriber(session: AsyncSession, chat_id: int) -> bool:


async def remove_subscriber(session: AsyncSession, chat_id: int) -> bool:
"""
Remove a subscriber from the database.
Returns:
bool: True if the subscriber was removed, False if they were not subscribed.
"""
subscriber = await session.get(Subscriber, chat_id)
if not subscriber:
return False
Expand All @@ -59,14 +31,5 @@ async def remove_subscriber(session: AsyncSession, chat_id: int) -> bool:


async def get_subscribers(session: AsyncSession) -> list[int]:
"""
Retrieve all subscribers from the database.
Returns:
list[int]: A list of chat IDs.
"""

result = (await session.execute(select(Subscriber))).scalars().all()
return [subscriber.chat_id for subscriber in result]


1 change: 1 addition & 0 deletions settings.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from pydantic import Field
from pydantic_settings import BaseSettings


class Settings(BaseSettings):
telegram_bot_token: str = Field(
..., env="TELEGRAM_BOT_TOKEN", description="Telegram bot token"
Expand Down

0 comments on commit 2fcbc51

Please sign in to comment.