Files
duty-teller/db/repository.py
Nikolay Tatarinov 7a963eccd1 Add event type handling for duties in the system
- Introduced a new `event_type` column in the `duties` table to categorize duties as 'duty', 'unavailable', or 'vacation'.
- Updated the duty schedule import functionality to parse and store event types from the JSON input.
- Enhanced the API response to include event types for each duty, improving the calendar display logic.
- Modified the web application to visually differentiate between duty types in the calendar and duty list.
- Updated tests to cover new event type functionality and ensure correct parsing and storage of duties.
- Revised README documentation to reflect changes in duty event types and their representation in the system.
2026-02-17 23:01:07 +03:00

119 lines
3.2 KiB
Python

"""Repository: get_or_create_user, get_duties, insert_duty."""
from datetime import datetime, timedelta
from sqlalchemy.orm import Session
from db.models import User, Duty
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."""
# start_at < to_date + 1 day so duties starting on to_date are included (start_at is ISO with T)
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.
from_date/to_date are YYYY-MM-DD (first/last day of month in client's local calendar).
Duty.start_at and end_at are stored in UTC (ISO 8601 with Z); lexicographic comparison
with date strings yields correct overlap.
"""
q = (
session.query(Duty, User.full_name)
.join(User, Duty.user_id == User.id)
.filter(Duty.start_at <= to_date, 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.
event_type: 'duty' | 'unavailable' | 'vacation'."""
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