proc-rota/djangorota/rota/models.py

from django.db import models
from django.utils import timezone


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:
			from rota_generator.shifts import RotaBuilder, SingleShift
		except Exception:
			raise RuntimeError("Unable to import rota_generator.shifts; ensure the rota package is importable")

		# compute weeks_to_rota (integer number of weeks)
		delta = self.end_date - self.start_date
		weeks_to_rota = max(1, int(delta.days // 7))

		# Allow passing options directly if they match RotaBuilder kwargs
		rb_kwargs = dict(self.options or {})
		rb = RotaBuilder(start_date=self.start_date, weeks_to_rota=weeks_to_rota, name=self.name, **rb_kwargs)

		# 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)

	rotas = models.ManyToManyField(RotaSchedule, through="Assignment", related_name="workers")

	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 percentage (100) while pydantic Worker expects int
		pyd_kwargs = {
			"name": self.name,
			"site": self.site or "",
			"grade": self.grade or 1,
			"id": self.pk,
			"fte": int(self.fte),
		}

		return PydWorker(**pyd_kwargs)


class Assignment(models.Model):
	"""Join model assigning a worker to a rota (can hold role / metadata)."""

	worker = models.ForeignKey(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:
		unique_together = ("worker", "rota")

	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)

	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()})"