scheduleSon/backend/database.py

from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import declarative_base
from sqlalchemy import Column, Integer, String, Boolean, text
import os

Base = declarative_base()

DATABASE_URL = os.getenv("DATABASE_PATH", "data/schedule.db")
# SQLite через aiosqlite требует sqlite+aiosqlite:///
DATABASE_URL_ASYNC = f"sqlite+aiosqlite:///{DATABASE_URL}"

engine = create_async_engine(DATABASE_URL_ASYNC, echo=False)
AsyncSessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


class Task(Base):
    __tablename__ = "tasks"
    
    id = Column(Integer, primary_key=True, index=True)
    date = Column(String, index=True)  # YYYY-MM-DD
    title = Column(String, nullable=False)
    repeat_weekly = Column(Boolean, default=False)
    weekday = Column(Integer)  # 0=Monday, 6=Sunday, только для weekly tasks


class Event(Base):
    __tablename__ = "events"
    
    id = Column(Integer, primary_key=True, index=True)
    date = Column(String, index=True)  # YYYY-MM-DD
    start_time = Column(String)  # HH:MM
    duration_min = Column(Integer, nullable=False)
    title = Column(String, nullable=False)
    repeat_weekly = Column(Boolean, default=False)
    weekday = Column(Integer)  # 0=Monday, 6=Sunday, только для weekly events


class WeeklyTaskException(Base):
    __tablename__ = "weekly_task_exceptions"
    
    id = Column(Integer, primary_key=True, index=True)
    task_id = Column(Integer, index=True)
    weekday = Column(Integer, nullable=False)  # 0=Monday, 6=Sunday
    date = Column(String, nullable=False)  # YYYY-MM-DD
    action = Column(String, nullable=False)  # "delete" или "replace"
    replacement_title = Column(String)  # для action="replace"
    replacement_date = Column(String)  # для переноса/замены одного вхождения


class EventException(Base):
    __tablename__ = "event_exceptions"

    id = Column(Integer, primary_key=True, index=True)
    event_id = Column(Integer, index=True, nullable=False)
    date = Column(String, nullable=False)
    action = Column(String, nullable=False)  # "delete" или "replace"
    replacement_title = Column(String)
    replacement_date = Column(String)
    replacement_start_time = Column(String)
    replacement_duration_min = Column(Integer)


class TelegramUser(Base):
    __tablename__ = "telegram_users"
    
    id = Column(Integer, primary_key=True, index=True)
    chat_id = Column(String, unique=True, nullable=False, index=True)


async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        await _run_migrations(conn)


async def _run_migrations(conn):
    table_columns = await _get_table_columns(conn)

    task_columns = table_columns.get("tasks", set())
    if "weekday" not in task_columns:
        await conn.execute(text("ALTER TABLE tasks ADD COLUMN weekday INTEGER"))

    event_columns = table_columns.get("events", set())
    if "repeat_weekly" not in event_columns:
        await conn.execute(text("ALTER TABLE events ADD COLUMN repeat_weekly BOOLEAN DEFAULT 0"))
    if "weekday" not in event_columns:
        await conn.execute(text("ALTER TABLE events ADD COLUMN weekday INTEGER"))

    exception_columns = table_columns.get("weekly_task_exceptions", set())
    if "task_id" not in exception_columns:
        await conn.execute(text("ALTER TABLE weekly_task_exceptions ADD COLUMN task_id INTEGER"))
    if "replacement_date" not in exception_columns:
        await conn.execute(text("ALTER TABLE weekly_task_exceptions ADD COLUMN replacement_date TEXT"))


async def _get_table_columns(conn):
    table_names = ["tasks", "events", "weekly_task_exceptions", "event_exceptions"]
    columns = {}
    for table_name in table_names:
        result = await conn.exec_driver_sql(f"PRAGMA table_info({table_name})")
        rows = result.fetchall()
        columns[table_name] = {row[1] for row in rows}
    return columns


async def get_db():
    async with AsyncSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()