from django.db import models from django.utils import timezone import uuid class RotaSchedule(models.Model): """Represents a named rota (collection of shifts over a date range). This is the top-level object that groups workers and runs. """ name = models.CharField(max_length=200, unique=True) start_date = models.DateField() end_date = models.DateField() description = models.TextField(blank=True) created_at = models.DateTimeField(auto_now_add=True) # Store the SingleShift definitions as a list of dicts (pydantic -> dict) shifts = models.JSONField(default=list, blank=True) # Additional builder options / constraints (serialisable) options = models.JSONField(default=dict, blank=True) def __str__(self): return self.name def to_rota_builder(self): """Construct a `rota.shifts.RotaBuilder` from this schedule and its shifts/workers. This will: - instantiate a RotaBuilder with `start_date` and computed `weeks_to_rota`; - convert stored `shifts` (list of dicts) to `SingleShift` pydantic models and add them; - convert assigned Django `Worker` objects to pydantic `Worker` and add them. Returns the RotaBuilder instance. """ # Try importing the project's rota generator implementation. Some setups # provide the implementation as `rota_generator.shifts` (script/package), # while others provide it as `rota.shifts` (library folder). Try both # so the Django app works in either layout. RotaBuilder = None SingleShift = None import importlib for mod_name in ("rota_generator.shifts", "rota.shifts"): try: mod = importlib.import_module(mod_name) RotaBuilder = getattr(mod, "RotaBuilder") SingleShift = getattr(mod, "SingleShift") break except Exception: continue if RotaBuilder is None or SingleShift is None: raise RuntimeError( "Unable to import rota generator (tried 'rota_generator.shifts' and 'rota.shifts'); " "ensure the rota package is importable and contains RotaBuilder/SingleShift" ) # compute weeks_to_rota (integer number of weeks) # Prefer an explicit value stored in `self.options['weeks']` (set by # the rota edit form). If absent, fall back to deriving from the # stored end_date (backwards compatible). opts_preview = dict(self.options or {}) weeks_opt = None if "weeks" in opts_preview: try: weeks_opt = int(opts_preview.get("weeks")) except Exception: weeks_opt = None if weeks_opt is not None: weeks_to_rota = max(1, weeks_opt) else: delta = self.end_date - self.start_date weeks_to_rota = max(1, int(delta.days // 7)) # Build kwargs for RotaBuilder from any explicit builder args stored in # `self.options` (kept for backward compatibility), and treat the # remaining keys as constraint options which should map to # `RotaConstraintOptions`. opts = dict(self.options or {}) # Remove 'weeks' from opts before handing to constraint options so # it doesn't cause validation failures in RotaConstraintOptions. if "weeks" in opts: try: opts.pop("weeks") except Exception: pass builder_keys = [ "balance_offset_modifier", "ltft_balance_offset", "use_previous_shifts", "use_shift_balance_extra", "use_bank_holiday_extra", "allow_force_assignment_with_leave_conflict", ] rb_kwargs = {k: opts.pop(k) for k in builder_keys if k in opts} rb = RotaBuilder(start_date=self.start_date, weeks_to_rota=weeks_to_rota, name=self.name, **rb_kwargs) # If the rota package exposes a typed RotaConstraintOptions model, use it # to validate/coerce the remaining `opts` and attach it to the builder. try: RotaConstraintOptions = getattr(mod, "RotaConstraintOptions") if opts: # validate raw JSON into the typed model opts_model = RotaConstraintOptions.model_validate(opts) else: opts_model = RotaConstraintOptions() # attach the typed model to the builder (overrides defaults) setattr(rb, "constraint_options_model", opts_model) except Exception: # If the model isn't available or validation fails, keep builder # defaults and proceed (backwards compatible behaviour). pass # Add shifts (expecting list of dicts compatible with SingleShift) shifts_objs = [] for s in self.shifts or []: try: obj = SingleShift(**s) except Exception as exc: raise ValueError(f"Invalid shift data for rota {self.name}: {exc}") shifts_objs.append(obj) if shifts_objs: rb.add_shifts(*shifts_objs) # Add workers assigned to this rota for w in self.workers.all(): try: pyd_worker = w.to_pydantic() except Exception: # skip workers that cannot be converted continue rb.add_worker(pyd_worker) return rb class Worker(models.Model): """Worker / staff member that can be assigned to one or more rotas.""" name = models.CharField(max_length=200) email = models.EmailField(blank=True) site = models.CharField(max_length=100, blank=True) grade = models.IntegerField(null=True, blank=True) fte = models.FloatField(default=1.0) active = models.BooleanField(default=True) # Store additional worker settings / metadata that map to the # pydantic `Worker` model used by the rota generator. This allows the # web UI to expose richer worker options without changing the core # Django schema for frequently used fields. options = models.JSONField(default=dict, blank=True) rotas = models.ManyToManyField(RotaSchedule, through="Assignment", related_name="workers") # Token used to allow workers (without Django accounts) to access # self-service links (request leave, view their calendar, etc.). This # is a UUID4 and should be treated as a secret. Regenerate to revoke. # During migration we add this field as nullable and non-unique first, # then populate unique values in a data migration and finally make it # non-null and unique. See migrations/ for the migration sequence. request_token = models.UUIDField(default=uuid.uuid4, editable=False, null=True, unique=False) def __str__(self): return self.name def to_pydantic(self): """Return a `rota.workers.Worker` pydantic model populated from this Django model. Only maps a minimal set of fields required by the scheduler; additional pydantic defaults will fill the rest. """ try: from rota_generator.workers import Worker as PydWorker except Exception: raise RuntimeError("Unable to import rota.workers.Worker; ensure project package is importable") # Map fte: Django stores `fte` as a float (fractional e.g. 1.0 == 100%), # while the pydantic Worker expects an integer percentage (100). # Convert conservatively: # - if stored value looks like a fraction (<= 10) treat it as fraction # and multiply by 100 (e.g. 1.0 -> 100). # - otherwise treat it as already a percentage and coerce to int. raw_fte = getattr(self, "fte", None) try: if raw_fte is None: fte_pct = 100 else: # numeric values f = float(raw_fte) if f <= 10: fte_pct = int(round(f * 100)) else: fte_pct = int(round(f)) except Exception: fte_pct = 100 pyd_kwargs = { "name": self.name, "site": self.site or "", "grade": self.grade or 1, "id": self.pk, "fte": fte_pct, } # Merge any additional options stored on the Django model into the # kwargs passed to the pydantic Worker model. The pydantic model will # honour its own defaults and types; storing these as a dict keeps the # Django DB schema simple while exposing full worker settings in the UI. try: opts = getattr(self, "options", None) or {} if isinstance(opts, dict): # shallow merge; keys in `options` override defaults above when present pyd_kwargs.update(opts) # Defensive normalization: the DB may contain `null`/`[]` for fields # that the pydantic model expects to be ints/dicts/lists. Coerce a # few common keys to safe defaults so Worker(**kwargs) doesn't raise # on older or malformed records. int_keys = [ "locum_max_shifts", "locum_max_shifts_per_week", "bank_holiday_extra", "prefer_multi_shift_together", "weekend_shift_target_number", ] dict_keys = [ "previous_shifts", "shift_balance_extra", "force_assign_with", "max_shifts_per_week_by_shift_name", "assign_as_block_preferences", "shift_fte_overrides", "exact_shifts", ] list_keys = [ "oop", "not_available_to_work", "pref_not_to_work", "work_requests", "locum_availability", "allowed_multi_shift_sets", "hard_day_dependencies", "hard_day_exclusions", "forced_assignments", "forced_assignments_by_date", ] for k in int_keys: v = pyd_kwargs.get(k, None) try: if v is None: pyd_kwargs[k] = 0 else: pyd_kwargs[k] = int(v) except Exception: pyd_kwargs[k] = 0 for k in dict_keys: v = pyd_kwargs.get(k, None) if not isinstance(v, dict): # If a list was stored accidentally (common migration/serialization # mistakes) coerce to an empty dict rather than letting pydantic # fail. pyd_kwargs[k] = {} for k in list_keys: v = pyd_kwargs.get(k, None) if v is None: pyd_kwargs[k] = [] elif isinstance(v, list): # leave as-is pass elif isinstance(v, (set, tuple)): pyd_kwargs[k] = list(v) else: pyd_kwargs[k] = [] set_keys = ["groups"] for k in set_keys: v = pyd_kwargs.get(k, None) try: if v is None: pyd_kwargs[k] = set() else: pyd_kwargs[k] = set(v) except Exception: pyd_kwargs[k] = set() # allowed_multi_shift_sets is expected to be a list of sets; if the # DB contains lists-of-lists convert inner lists to sets. ams = pyd_kwargs.get("allowed_multi_shift_sets") if isinstance(ams, list): new = [] for item in ams: if isinstance(item, list): new.append(set(item)) elif isinstance(item, (set, tuple)): new.append(set(item)) else: # leave unknown types to pydantic to either accept or coerce new.append(item) pyd_kwargs["allowed_multi_shift_sets"] = new except Exception: # ignore any issues reading options and proceed with core fields pass return PydWorker(**pyd_kwargs) class Assignment(models.Model): """Join model assigning a worker to a rota (can hold role / metadata).""" # Enforce that a Worker may only be assigned to a single RotaSchedule by # using a OneToOneField here. This ensures at the DB level that each # Worker has at most one Assignment (and therefore at most one rota). worker = models.OneToOneField(Worker, on_delete=models.CASCADE) rota = models.ForeignKey(RotaSchedule, on_delete=models.CASCADE) role = models.CharField(max_length=100, blank=True) created_at = models.DateTimeField(auto_now_add=True) class Meta: # With OneToOneField on worker, uniqueness for (worker, rota) is # implied by the worker uniqueness. Keep ordering for convenience. ordering = ("-created_at",) def __str__(self): return f"{self.worker} -> {self.rota}" class Leave(models.Model): """Leave record for a worker. Stored as inclusive start/end dates.""" worker = models.ForeignKey(Worker, on_delete=models.CASCADE, related_name="leaves") start_date = models.DateField() end_date = models.DateField() reason = models.CharField(max_length=255, blank=True) created_at = models.DateTimeField(auto_now_add=True) class Meta: ordering = ("-start_date",) def __str__(self): return f"{self.worker}: {self.start_date} → {self.end_date}" class RotaRun(models.Model): """Represents an execution of the scheduler for a RotaSchedule. The run stores status and (optionally) JSON results. A background runner or management command can update this record as the job progresses. """ STATUS_PENDING = "pending" STATUS_RUNNING = "running" STATUS_SUCCESS = "success" STATUS_FAILED = "failed" STATUS_CHOICES = [ (STATUS_PENDING, "Pending"), (STATUS_RUNNING, "Running"), (STATUS_SUCCESS, "Success"), (STATUS_FAILED, "Failed"), ] rota = models.ForeignKey(RotaSchedule, on_delete=models.CASCADE, related_name="runs") created_at = models.DateTimeField(auto_now_add=True) started_at = models.DateTimeField(null=True, blank=True) finished_at = models.DateTimeField(null=True, blank=True) status = models.CharField(max_length=20, choices=STATUS_CHOICES, default=STATUS_PENDING) result = models.JSONField(null=True, blank=True) log = models.TextField(blank=True) # Optional HTML export produced by a run (stored inline for convenience). # May be large; keep nullable to avoid breaking existing runs. export_html = models.TextField(null=True, blank=True) def mark_running(self): self.status = self.STATUS_RUNNING self.started_at = timezone.now() self.save(update_fields=["status", "started_at"]) def mark_finished(self, result=None): self.status = self.STATUS_SUCCESS self.finished_at = timezone.now() if result is not None: self.result = result self.save(update_fields=["status", "finished_at", "result"]) def mark_failed(self, message=""): self.status = self.STATUS_FAILED self.finished_at = timezone.now() self.log = (self.log or "") + "\n" + message self.save(update_fields=["status", "finished_at", "log"]) def __str__(self): return f"RotaRun({self.rota.name} - {self.status} @ {self.created_at.isoformat()})"