Refactor project structure and enhance Docker configuration
- Updated `.dockerignore` to exclude test and development artifacts, optimizing the Docker image size. - Refactored `main.py` to delegate execution to `duty_teller.run.main()`, simplifying the entry point. - Introduced a new `duty_teller` package to encapsulate core functionality, improving modularity and organization. - Enhanced `pyproject.toml` to define a script for running the application, streamlining the execution process. - Updated README documentation to reflect changes in project structure and usage instructions. - Improved Alembic environment configuration to utilize the new package structure for database migrations.
This commit is contained in:
52
duty_teller/db/__init__.py
Normal file
52
duty_teller/db/__init__.py
Normal file
@@ -0,0 +1,52 @@
|
||||
"""Database layer: SQLAlchemy models, Pydantic schemas, repository, init."""
|
||||
|
||||
from duty_teller.db.models import Base, User, Duty
|
||||
from duty_teller.db.schemas import (
|
||||
UserCreate,
|
||||
UserInDb,
|
||||
DutyCreate,
|
||||
DutyInDb,
|
||||
DutyWithUser,
|
||||
)
|
||||
from duty_teller.db.session import (
|
||||
get_engine,
|
||||
get_session_factory,
|
||||
get_session,
|
||||
session_scope,
|
||||
)
|
||||
from duty_teller.db.repository import (
|
||||
delete_duties_in_range,
|
||||
get_or_create_user,
|
||||
get_or_create_user_by_full_name,
|
||||
get_duties,
|
||||
insert_duty,
|
||||
set_user_phone,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"Base",
|
||||
"User",
|
||||
"Duty",
|
||||
"UserCreate",
|
||||
"UserInDb",
|
||||
"DutyCreate",
|
||||
"DutyInDb",
|
||||
"DutyWithUser",
|
||||
"get_engine",
|
||||
"get_session_factory",
|
||||
"get_session",
|
||||
"session_scope",
|
||||
"delete_duties_in_range",
|
||||
"get_or_create_user",
|
||||
"get_or_create_user_by_full_name",
|
||||
"get_duties",
|
||||
"insert_duty",
|
||||
"set_user_phone",
|
||||
"init_db",
|
||||
]
|
||||
|
||||
|
||||
def init_db(database_url: str) -> None:
|
||||
"""Create tables from metadata (Alembic migrations handle schema in production)."""
|
||||
engine = get_engine(database_url)
|
||||
Base.metadata.create_all(bind=engine)
|
||||
51
duty_teller/db/models.py
Normal file
51
duty_teller/db/models.py
Normal file
@@ -0,0 +1,51 @@
|
||||
"""SQLAlchemy ORM models for users and duties."""
|
||||
|
||||
from sqlalchemy import ForeignKey, Integer, BigInteger, Text
|
||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
|
||||
|
||||
|
||||
class Base(DeclarativeBase):
|
||||
"""Declarative base for all models."""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class User(Base):
|
||||
__tablename__ = "users"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
telegram_user_id: Mapped[int | None] = mapped_column(
|
||||
BigInteger, unique=True, nullable=True
|
||||
)
|
||||
full_name: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
username: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
first_name: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
last_name: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
phone: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
|
||||
duties: Mapped[list["Duty"]] = relationship("Duty", back_populates="user")
|
||||
|
||||
|
||||
class Duty(Base):
|
||||
__tablename__ = "duties"
|
||||
|
||||
id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id: Mapped[int] = mapped_column(
|
||||
Integer, ForeignKey("users.id"), nullable=False
|
||||
)
|
||||
# UTC, ISO 8601 with Z suffix (e.g. 2025-01-15T09:00:00Z)
|
||||
start_at: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
end_at: Mapped[str] = mapped_column(Text, nullable=False)
|
||||
# duty | unavailable | vacation
|
||||
event_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="duty")
|
||||
|
||||
user: Mapped["User"] = relationship("User", back_populates="duties")
|
||||
|
||||
|
||||
class GroupDutyPin(Base):
|
||||
"""Stores which message to update in each group for the pinned duty notice."""
|
||||
|
||||
__tablename__ = "group_duty_pins"
|
||||
|
||||
chat_id: Mapped[int] = mapped_column(BigInteger, primary_key=True)
|
||||
message_id: Mapped[int] = mapped_column(Integer, nullable=False)
|
||||
213
duty_teller/db/repository.py
Normal file
213
duty_teller/db/repository.py
Normal file
@@ -0,0 +1,213 @@
|
||||
"""Repository: get_or_create_user, get_duties, insert_duty, get_current_duty, group_duty_pins."""
|
||||
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from duty_teller.db.models import User, Duty, GroupDutyPin
|
||||
|
||||
|
||||
def get_or_create_user(
|
||||
session: Session,
|
||||
telegram_user_id: int,
|
||||
full_name: str,
|
||||
username: str | None = None,
|
||||
first_name: str | None = None,
|
||||
last_name: str | None = None,
|
||||
) -> User:
|
||||
user = session.query(User).filter(User.telegram_user_id == telegram_user_id).first()
|
||||
if user:
|
||||
user.full_name = full_name
|
||||
user.username = username
|
||||
user.first_name = first_name
|
||||
user.last_name = last_name
|
||||
session.commit()
|
||||
session.refresh(user)
|
||||
return user
|
||||
user = User(
|
||||
telegram_user_id=telegram_user_id,
|
||||
full_name=full_name,
|
||||
username=username,
|
||||
first_name=first_name,
|
||||
last_name=last_name,
|
||||
)
|
||||
session.add(user)
|
||||
session.commit()
|
||||
session.refresh(user)
|
||||
return user
|
||||
|
||||
|
||||
def get_or_create_user_by_full_name(session: Session, full_name: str) -> User:
|
||||
"""Find user by exact full_name or create one with telegram_user_id=None (for duty-schedule import)."""
|
||||
user = session.query(User).filter(User.full_name == full_name).first()
|
||||
if user:
|
||||
return user
|
||||
user = User(
|
||||
telegram_user_id=None,
|
||||
full_name=full_name,
|
||||
username=None,
|
||||
first_name=None,
|
||||
last_name=None,
|
||||
)
|
||||
session.add(user)
|
||||
session.commit()
|
||||
session.refresh(user)
|
||||
return user
|
||||
|
||||
|
||||
def delete_duties_in_range(
|
||||
session: Session,
|
||||
user_id: int,
|
||||
from_date: str,
|
||||
to_date: str,
|
||||
) -> int:
|
||||
"""Delete all duties of the user that overlap [from_date, to_date] (YYYY-MM-DD). Returns count deleted."""
|
||||
to_next = (
|
||||
datetime.fromisoformat(to_date + "T00:00:00") + timedelta(days=1)
|
||||
).strftime("%Y-%m-%d")
|
||||
q = session.query(Duty).filter(
|
||||
Duty.user_id == user_id,
|
||||
Duty.start_at < to_next,
|
||||
Duty.end_at >= from_date,
|
||||
)
|
||||
count = q.count()
|
||||
q.delete(synchronize_session=False)
|
||||
session.commit()
|
||||
return count
|
||||
|
||||
|
||||
def get_duties(
|
||||
session: Session,
|
||||
from_date: str,
|
||||
to_date: str,
|
||||
) -> list[tuple[Duty, str]]:
|
||||
"""Return list of (Duty, full_name) overlapping the given date range."""
|
||||
to_date_next = (
|
||||
datetime.fromisoformat(to_date + "T00:00:00") + timedelta(days=1)
|
||||
).strftime("%Y-%m-%d")
|
||||
q = (
|
||||
session.query(Duty, User.full_name)
|
||||
.join(User, Duty.user_id == User.id)
|
||||
.filter(Duty.start_at < to_date_next, Duty.end_at >= from_date)
|
||||
)
|
||||
return list(q.all())
|
||||
|
||||
|
||||
def insert_duty(
|
||||
session: Session,
|
||||
user_id: int,
|
||||
start_at: str,
|
||||
end_at: str,
|
||||
event_type: str = "duty",
|
||||
) -> Duty:
|
||||
"""Create a duty. start_at and end_at must be UTC, ISO 8601 with Z."""
|
||||
duty = Duty(
|
||||
user_id=user_id,
|
||||
start_at=start_at,
|
||||
end_at=end_at,
|
||||
event_type=event_type,
|
||||
)
|
||||
session.add(duty)
|
||||
session.commit()
|
||||
session.refresh(duty)
|
||||
return duty
|
||||
|
||||
|
||||
def get_current_duty(session: Session, at_utc: datetime) -> tuple[Duty, User] | None:
|
||||
"""Return the duty (and user) for which start_at <= at_utc < end_at, event_type='duty'."""
|
||||
from datetime import timezone
|
||||
|
||||
if at_utc.tzinfo is not None:
|
||||
at_utc = at_utc.astimezone(timezone.utc)
|
||||
now_iso = at_utc.strftime("%Y-%m-%dT%H:%M:%S") + "Z"
|
||||
row = (
|
||||
session.query(Duty, User)
|
||||
.join(User, Duty.user_id == User.id)
|
||||
.filter(
|
||||
Duty.event_type == "duty",
|
||||
Duty.start_at <= now_iso,
|
||||
Duty.end_at > now_iso,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if row is None:
|
||||
return None
|
||||
return (row[0], row[1])
|
||||
|
||||
|
||||
def get_next_shift_end(session: Session, after_utc: datetime) -> datetime | None:
|
||||
"""Return the end_at of the current duty or of the next duty. Naive UTC."""
|
||||
from datetime import timezone
|
||||
|
||||
if after_utc.tzinfo is not None:
|
||||
after_utc = after_utc.astimezone(timezone.utc)
|
||||
after_iso = after_utc.strftime("%Y-%m-%dT%H:%M:%S") + "Z"
|
||||
current = (
|
||||
session.query(Duty)
|
||||
.filter(
|
||||
Duty.event_type == "duty",
|
||||
Duty.start_at <= after_iso,
|
||||
Duty.end_at > after_iso,
|
||||
)
|
||||
.first()
|
||||
)
|
||||
if current:
|
||||
return datetime.fromisoformat(current.end_at.replace("Z", "+00:00")).replace(
|
||||
tzinfo=None
|
||||
)
|
||||
next_duty = (
|
||||
session.query(Duty)
|
||||
.filter(Duty.event_type == "duty", Duty.start_at > after_iso)
|
||||
.order_by(Duty.start_at)
|
||||
.first()
|
||||
)
|
||||
if next_duty:
|
||||
return datetime.fromisoformat(next_duty.end_at.replace("Z", "+00:00")).replace(
|
||||
tzinfo=None
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def get_group_duty_pin(session: Session, chat_id: int) -> GroupDutyPin | None:
|
||||
"""Get the pinned message record for a chat, if any."""
|
||||
return session.query(GroupDutyPin).filter(GroupDutyPin.chat_id == chat_id).first()
|
||||
|
||||
|
||||
def save_group_duty_pin(
|
||||
session: Session, chat_id: int, message_id: int
|
||||
) -> GroupDutyPin:
|
||||
"""Save or update the pinned message for a chat."""
|
||||
pin = session.query(GroupDutyPin).filter(GroupDutyPin.chat_id == chat_id).first()
|
||||
if pin:
|
||||
pin.message_id = message_id
|
||||
else:
|
||||
pin = GroupDutyPin(chat_id=chat_id, message_id=message_id)
|
||||
session.add(pin)
|
||||
session.commit()
|
||||
session.refresh(pin)
|
||||
return pin
|
||||
|
||||
|
||||
def delete_group_duty_pin(session: Session, chat_id: int) -> None:
|
||||
"""Remove the pinned message record when the bot leaves the group."""
|
||||
session.query(GroupDutyPin).filter(GroupDutyPin.chat_id == chat_id).delete()
|
||||
session.commit()
|
||||
|
||||
|
||||
def get_all_group_duty_pin_chat_ids(session: Session) -> list[int]:
|
||||
"""Return all chat_ids that have a pinned duty message (for restoring jobs on startup)."""
|
||||
rows = session.query(GroupDutyPin.chat_id).all()
|
||||
return [r[0] for r in rows]
|
||||
|
||||
|
||||
def set_user_phone(
|
||||
session: Session, telegram_user_id: int, phone: str | None
|
||||
) -> User | None:
|
||||
"""Set phone for user by telegram_user_id. Returns User or None if not found."""
|
||||
user = session.query(User).filter(User.telegram_user_id == telegram_user_id).first()
|
||||
if not user:
|
||||
return None
|
||||
user.phone = phone
|
||||
session.commit()
|
||||
session.refresh(user)
|
||||
return user
|
||||
58
duty_teller/db/schemas.py
Normal file
58
duty_teller/db/schemas.py
Normal file
@@ -0,0 +1,58 @@
|
||||
"""Pydantic schemas for API and validation."""
|
||||
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
|
||||
|
||||
class UserBase(BaseModel):
|
||||
full_name: str
|
||||
username: str | None = None
|
||||
first_name: str | None = None
|
||||
last_name: str | None = None
|
||||
|
||||
|
||||
class UserCreate(UserBase):
|
||||
telegram_user_id: int
|
||||
|
||||
|
||||
class UserInDb(UserBase):
|
||||
id: int
|
||||
telegram_user_id: int
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class DutyBase(BaseModel):
|
||||
user_id: int
|
||||
start_at: str # UTC, ISO 8601 with Z
|
||||
end_at: str # UTC, ISO 8601 with Z
|
||||
|
||||
|
||||
class DutyCreate(DutyBase):
|
||||
pass
|
||||
|
||||
|
||||
class DutyInDb(DutyBase):
|
||||
id: int
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class DutyWithUser(DutyInDb):
|
||||
"""Duty with full_name and event_type for calendar display.
|
||||
|
||||
event_type: only these values are returned; unknown DB values are mapped to "duty" in the API.
|
||||
"""
|
||||
|
||||
full_name: str
|
||||
event_type: Literal["duty", "unavailable", "vacation"] = "duty"
|
||||
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
|
||||
class CalendarEvent(BaseModel):
|
||||
"""External calendar event (e.g. holiday) for a single day."""
|
||||
|
||||
date: str # YYYY-MM-DD
|
||||
summary: str
|
||||
57
duty_teller/db/session.py
Normal file
57
duty_teller/db/session.py
Normal file
@@ -0,0 +1,57 @@
|
||||
"""SQLAlchemy engine and session factory.
|
||||
|
||||
Engine and session factory are cached globally per process. Only one DATABASE_URL
|
||||
is effectively used for the process lifetime. Using a different URL later (e.g. in
|
||||
tests with in-memory SQLite) would still use the first engine. To use a different
|
||||
URL in tests, set env (e.g. DATABASE_URL) before the first import of this module, or
|
||||
clear _engine and _SessionLocal in test fixtures. Prefer session_scope() for all
|
||||
callers so sessions are always closed and rolled back on error.
|
||||
"""
|
||||
|
||||
from contextlib import contextmanager
|
||||
from typing import Generator
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
|
||||
_engine = None
|
||||
_SessionLocal = None
|
||||
|
||||
|
||||
@contextmanager
|
||||
def session_scope(database_url: str) -> Generator[Session, None, None]:
|
||||
"""Context manager: yields a session, rolls back on exception, closes on exit."""
|
||||
session = get_session(database_url)
|
||||
try:
|
||||
yield session
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
|
||||
def get_engine(database_url: str):
|
||||
global _engine
|
||||
if _engine is None:
|
||||
_engine = create_engine(
|
||||
database_url,
|
||||
connect_args={"check_same_thread": False}
|
||||
if "sqlite" in database_url
|
||||
else {},
|
||||
echo=False,
|
||||
)
|
||||
return _engine
|
||||
|
||||
|
||||
def get_session_factory(database_url: str) -> sessionmaker[Session]:
|
||||
global _SessionLocal
|
||||
if _SessionLocal is None:
|
||||
engine = get_engine(database_url)
|
||||
_SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
return _SessionLocal
|
||||
|
||||
|
||||
def get_session(database_url: str) -> Session:
|
||||
return get_session_factory(database_url)()
|
||||
Reference in New Issue
Block a user