from __future__ import annotations from typing import TYPE_CHECKING from django.contrib.auth import get_user_model from django.contrib.postgres.fields import DateTimeRangeField from django.db import models from django.db.models import F, Q, QuerySet from django.utils import timezone from model_utils.managers import InheritanceQuerySetMixin from psycopg.types.range import Range if TYPE_CHECKING: from collections.abc import Iterable from datetime import datetime class Resource(models.Model): name = models.CharField(max_length=256) parent = models.ForeignKey( "Resource", on_delete=models.CASCADE, blank=True, null=True, related_name="children", ) min_length = models.DurationField() max_length = models.DurationField() google_calendar = models.CharField(max_length=1024, unique=True) def __str__(self) -> str: if self.parent: return f"{self.parent} / {self.name}" else: return self.name def get_recursive_parents(self) -> Iterable[Resource]: if self.parent: yield self.parent yield from self.parent.get_recursive_parents() def get_recursive_children(self) -> Iterable[Resource]: for child in self.children.all(): yield child yield from child.get_recursive_children() def get_related(self) -> Iterable[Resource]: yield self yield from self.get_recursive_parents() yield from self.get_recursive_children() class ReservationQuerySet(InheritanceQuerySetMixin, QuerySet): def filter_after(self, start: datetime) -> ReservationQuerySet: """ Selects events that are after the specified datetime, including overlaps but excluding exactly matching endpoints. """ return self.filter(timespan__overlap=Range(start, None)) def filter_before(self, end: datetime) -> ReservationQuerySet: """ Selects events that are before the specified datetime, including overlaps but excluding exactly matching endpoints. """ return self.filter(timespan__overlap=Range(None, end)) def filter_between(self, start: datetime, end: datetime) -> ReservationQuerySet: """ Selects events that are between the specified datetime, including overlaps but excluding exactly matching endpoints. """ return self.filter(timespan__overlap=Range(start, end)) class Reservation(models.Model): resources = models.ManyToManyField(Resource, blank=True) start = models.DateTimeField() end = models.DateTimeField() timespan = models.GeneratedField( expression=models.Func(F("start"), F("end"), function="tstzrange"), output_field=DateTimeRangeField(), db_persist=True, ) google_calendar_event_id = models.CharField( max_length=1024, null=True, blank=True, unique=True ) duration = models.GeneratedField( expression=F("end") - F("start"), output_field=models.DurationField(), db_persist=True, ) objects = ReservationQuerySet.as_manager() class Meta: constraints = [ models.CheckConstraint(check=Q(end__gt=F("start")), name="end_after_start"), ] def __str__(self) -> str: resources = ", ".join(str(resource) for resource in self.resources.all()) return f"{resources}: {self.start} - {self.end}" def __repr__(self) -> str: """Redefined to avoid __str__ and therefore .resources, which was causing recursion issues""" return f"<{self.__class__.__name__}: {self.start} - {self.end}>" def get_title(self) -> str: return "Unknown Reservation" def make_google_calendar_event(self): event = { "summary": "CMSManage Reservation", "start": { "dateTime": self.start.isoformat(), "timeZone": timezone.get_default_timezone_name(), }, "end": { "dateTime": self.end.isoformat(), "timeZone": timezone.get_default_timezone_name(), }, "extendedProperties": {"private": {"cmsmanage": "1"}}, "status": "confirmed", } # use existing id if it exists, otherwise let Google generate it if self.google_calendar_event_id: event["id"] = self.google_calendar_event_id return event class UserReservation(Reservation): user = models.ForeignKey( get_user_model(), on_delete=models.CASCADE, null=True, blank=True ) def __str__(self) -> str: return f"{self.user} | {super().__str__()}" def get_title(self) -> str: return str(self.user) def make_google_calendar_event(self): return super().make_google_calendar_event() | { "summary": str(self.user), } class ExternalReservation(Reservation): """Reservations created by something else in Google Calendar""" title = models.CharField(max_length=1024) def __str__(self) -> str: return f'External "{self.title}" | {super().__str__()}' def get_title(self) -> str: return str(self.title) def make_google_calendar_event(self): """This should never be called, as these are reservations from Google Calendar and shouldn't be synced back""" raise AttributeError( "External Reservations should not be pushed back to Google Calendar" )