Nikolay Tatarinov
ef5dbca5df
- Added a new command `/import_duty_schedule` for importing duty schedules via JSON, restricted to admin users. - Introduced a two-step import process: specifying handover time and uploading a JSON file. - Updated the database schema to allow `telegram_user_id` to be nullable for user creation by full name. - Implemented repository functions for user management, including `get_or_create_user_by_full_name` and `delete_duties_in_range`. - Enhanced README documentation with details on the new import command and JSON format requirements. - Added comprehensive tests for the duty schedule parser and integration tests for the import functionality.
112 lines
3.1 KiB
Python
112 lines
3.1 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,
|
|
) -> Duty:
|
|
"""Create a duty. start_at and end_at must be UTC, ISO 8601 with Z (e.g. 2025-01-15T09:00:00Z)."""
|
|
duty = Duty(user_id=user_id, start_at=start_at, end_at=end_at)
|
|
session.add(duty)
|
|
session.commit()
|
|
session.refresh(duty)
|
|
return duty
|