__package__ = 'archivebox.workers' import uuid import json from typing import ClassVar, Type, Iterable, TypedDict from datetime import datetime, timedelta from statemachine.mixins import MachineMixin from django.db import models from django.db.models import QuerySet from django.core import checks from django.utils import timezone from django.utils.functional import classproperty from base_models.models import ABIDModel, ABIDField from machine.models import Process from statemachine import registry, StateMachine, State class DefaultStatusChoices(models.TextChoices): QUEUED = 'queued', 'Queued' STARTED = 'started', 'Started' SEALED = 'sealed', 'Sealed' default_status_field: models.CharField = models.CharField(choices=DefaultStatusChoices.choices, max_length=15, default=DefaultStatusChoices.QUEUED, null=False, blank=False, db_index=True) default_retry_at_field: models.DateTimeField = models.DateTimeField(default=timezone.now, null=True, blank=True, db_index=True) ObjectState = State | str ObjectStateList = Iterable[ObjectState] class BaseModelWithStateMachine(models.Model, MachineMixin): id: models.UUIDField StatusChoices: ClassVar[Type[models.TextChoices]] # status: models.CharField # retry_at: models.DateTimeField state_machine_name: ClassVar[str] state_field_name: ClassVar[str] state_machine_attr: ClassVar[str] = 'sm' bind_events_as_methods: ClassVar[bool] = True active_state: ClassVar[ObjectState] retry_at_field_name: ClassVar[str] class Meta: abstract = True @classmethod def check(cls, sender=None, **kwargs): errors = super().check(**kwargs) found_id_field = False found_status_field = False found_retry_at_field = False for field in cls._meta.get_fields(): if getattr(field, '_is_state_field', False): if cls.state_field_name == field.name: found_status_field = True if getattr(field, 'choices', None) != cls.StatusChoices.choices: errors.append(checks.Error( f'{cls.__name__}.{field.name} must have choices set to {cls.__name__}.StatusChoices.choices', hint=f'{cls.__name__}.{field.name}.choices = {getattr(field, "choices", None)!r}', obj=cls, id='workers.E011', )) if getattr(field, '_is_retry_at_field', False): if cls.retry_at_field_name == field.name: found_retry_at_field = True if field.name == 'id' and getattr(field, 'primary_key', False): found_id_field = True if not found_status_field: errors.append(checks.Error( f'{cls.__name__}.state_field_name must be defined and point to a StatusField()', hint=f'{cls.__name__}.state_field_name = {cls.state_field_name!r} but {cls.__name__}.{cls.state_field_name!r} was not found or does not refer to StatusField', obj=cls, id='workers.E012', )) if not found_retry_at_field: errors.append(checks.Error( f'{cls.__name__}.retry_at_field_name must be defined and point to a RetryAtField()', hint=f'{cls.__name__}.retry_at_field_name = {cls.retry_at_field_name!r} but {cls.__name__}.{cls.retry_at_field_name!r} was not found or does not refer to RetryAtField', obj=cls, id='workers.E013', )) if not found_id_field: errors.append(checks.Error( f'{cls.__name__} must have an id field that is a primary key', hint=f'{cls.__name__}.id = {cls.id!r}', obj=cls, id='workers.E014', )) if not isinstance(cls.state_machine_name, str): errors.append(checks.Error( f'{cls.__name__}.state_machine_name must be a dotted-import path to a StateMachine class', hint=f'{cls.__name__}.state_machine_name = {cls.state_machine_name!r}', obj=cls, id='workers.E015', )) try: cls.StateMachineClass except Exception as err: errors.append(checks.Error( f'{cls.__name__}.state_machine_name must point to a valid StateMachine class, but got {type(err).__name__} {err} when trying to access {cls.__name__}.StateMachineClass', hint=f'{cls.__name__}.state_machine_name = {cls.state_machine_name!r}', obj=cls, id='workers.E016', )) if cls.INITIAL_STATE not in cls.StatusChoices.values: errors.append(checks.Error( f'{cls.__name__}.StateMachineClass.initial_state must be present within {cls.__name__}.StatusChoices', hint=f'{cls.__name__}.StateMachineClass.initial_state = {cls.StateMachineClass.initial_state!r}', obj=cls, id='workers.E017', )) if cls.ACTIVE_STATE not in cls.StatusChoices.values: errors.append(checks.Error( f'{cls.__name__}.active_state must be set to a valid State present within {cls.__name__}.StatusChoices', hint=f'{cls.__name__}.active_state = {cls.active_state!r}', obj=cls, id='workers.E018', )) for state in cls.FINAL_STATES: if state not in cls.StatusChoices.values: errors.append(checks.Error( f'{cls.__name__}.StateMachineClass.final_states must all be present within {cls.__name__}.StatusChoices', hint=f'{cls.__name__}.StateMachineClass.final_states = {cls.StateMachineClass.final_states!r}', obj=cls, id='workers.E019', )) break return errors @staticmethod def _state_to_str(state: ObjectState) -> str: """Convert a statemachine.State, models.TextChoices.choices value, or Enum value to a str""" return str(state.value) if isinstance(state, State) else str(state) @property def RETRY_AT(self) -> datetime: return getattr(self, self.retry_at_field_name) @RETRY_AT.setter def RETRY_AT(self, value: datetime): setattr(self, self.retry_at_field_name, value) @property def STATE(self) -> str: return getattr(self, self.state_field_name) @STATE.setter def STATE(self, value: str): setattr(self, self.state_field_name, value) def bump_retry_at(self, seconds: int = 10): self.RETRY_AT = timezone.now() + timedelta(seconds=seconds) @classproperty def ACTIVE_STATE(cls) -> str: return cls._state_to_str(cls.active_state) @classproperty def INITIAL_STATE(cls) -> str: return cls._state_to_str(cls.StateMachineClass.initial_state) @classproperty def FINAL_STATES(cls) -> list[str]: return [cls._state_to_str(state) for state in cls.StateMachineClass.final_states] @classproperty def FINAL_OR_ACTIVE_STATES(cls) -> list[str]: return [*cls.FINAL_STATES, cls.ACTIVE_STATE] @classmethod def extend_choices(cls, base_choices: Type[models.TextChoices]): """ Decorator to extend the base choices with extra choices, e.g.: class MyModel(ModelWithStateMachine): @ModelWithStateMachine.extend_choices(ModelWithStateMachine.StatusChoices) class StatusChoices(models.TextChoices): SUCCEEDED = 'succeeded' FAILED = 'failed' SKIPPED = 'skipped' """ assert issubclass(base_choices, models.TextChoices), f'@extend_choices(base_choices) must be a TextChoices class, not {base_choices.__name__}' def wrapper(extra_choices: Type[models.TextChoices]) -> Type[models.TextChoices]: joined = {} for item in base_choices.choices: joined[item[0]] = item[1] for item in extra_choices.choices: joined[item[0]] = item[1] return models.TextChoices('StatusChoices', joined) return wrapper @classmethod def StatusField(cls, **kwargs) -> models.CharField: """ Used on subclasses to extend/modify the status field with updated kwargs. e.g.: class MyModel(ModelWithStateMachine): class StatusChoices(ModelWithStateMachine.StatusChoices): QUEUED = 'queued', 'Queued' STARTED = 'started', 'Started' SEALED = 'sealed', 'Sealed' BACKOFF = 'backoff', 'Backoff' FAILED = 'failed', 'Failed' SKIPPED = 'skipped', 'Skipped' status = ModelWithStateMachine.StatusField(choices=StatusChoices.choices, default=StatusChoices.QUEUED) """ default_kwargs = default_status_field.deconstruct()[3] updated_kwargs = {**default_kwargs, **kwargs} field = models.CharField(**updated_kwargs) field._is_state_field = True # type: ignore return field @classmethod def RetryAtField(cls, **kwargs) -> models.DateTimeField: """ Used on subclasses to extend/modify the retry_at field with updated kwargs. e.g.: class MyModel(ModelWithStateMachine): retry_at = ModelWithStateMachine.RetryAtField(editable=False) """ default_kwargs = default_retry_at_field.deconstruct()[3] updated_kwargs = {**default_kwargs, **kwargs} field = models.DateTimeField(**updated_kwargs) field._is_retry_at_field = True # type: ignore return field @classproperty def StateMachineClass(cls) -> Type[StateMachine]: """Get the StateMachine class for the given django Model that inherits from MachineMixin""" model_state_machine_name = getattr(cls, 'state_machine_name', None) if model_state_machine_name: StateMachineCls = registry.get_machine_cls(model_state_machine_name) assert issubclass(StateMachineCls, StateMachine) return StateMachineCls raise NotImplementedError(f'ActorType[{cls.__name__}] must define .state_machine_name: str that points to a valid StateMachine') # @classproperty # def final_q(cls) -> Q: # """Get the filter for objects that are in a final state""" # return Q(**{f'{cls.state_field_name}__in': cls.final_states}) # @classproperty # def active_q(cls) -> Q: # """Get the filter for objects that are actively processing right now""" # return Q(**{cls.state_field_name: cls._state_to_str(cls.active_state)}) # e.g. Q(status='started') # @classproperty # def stalled_q(cls) -> Q: # """Get the filter for objects that are marked active but have timed out""" # return cls.active_q & Q(retry_at__lte=timezone.now()) # e.g. Q(status='started') AND Q() # @classproperty # def future_q(cls) -> Q: # """Get the filter for objects that have a retry_at in the future""" # return Q(retry_at__gt=timezone.now()) # @classproperty # def pending_q(cls) -> Q: # """Get the filter for objects that are ready for processing.""" # return ~(cls.active_q) & ~(cls.final_q) & ~(cls.future_q) # @classmethod # def get_queue(cls) -> QuerySet: # """ # Get the sorted and filtered QuerySet of objects that are ready for processing. # e.g. qs.exclude(status__in=('sealed', 'started'), retry_at__gt=timezone.now()).order_by('retry_at') # """ # return cls.objects.filter(cls.pending_q) class ModelWithStateMachine(BaseModelWithStateMachine): StatusChoices: ClassVar[Type[DefaultStatusChoices]] = DefaultStatusChoices status: models.CharField = BaseModelWithStateMachine.StatusField() retry_at: models.DateTimeField = BaseModelWithStateMachine.RetryAtField() state_machine_name: ClassVar[str] # e.g. 'core.statemachines.ArchiveResultMachine' state_field_name: ClassVar[str] = 'status' state_machine_attr: ClassVar[str] = 'sm' bind_events_as_methods: ClassVar[bool] = True active_state: ClassVar[str] = StatusChoices.STARTED retry_at_field_name: ClassVar[str] = 'retry_at' class Meta: abstract = True class EventDict(TypedDict, total=False): name: str id: str | uuid.UUID path: str content: str status: str retry_at: datetime | None url: str seed_id: str | uuid.UUID crawl_id: str | uuid.UUID snapshot_id: str | uuid.UUID process_id: str | uuid.UUID extractor: str error: str on_success: dict | None on_failure: dict | None class EventManager(models.Manager): pass class EventQuerySet(models.QuerySet): def get_next_unclaimed(self) -> 'Event | None': return self.filter(claimed_at=None).order_by('deliver_at').first() def expired(self, older_than: int=60 * 10) -> QuerySet['Event']: return self.filter(claimed_at__lt=timezone.now() - timedelta(seconds=older_than)) class Event(ABIDModel): abid_prefix = 'evn_' abid_ts_src = 'self.deliver_at' # e.g. 'self.created_at' abid_uri_src = 'self.name' # e.g. 'self.uri' (MUST BE SET) abid_subtype_src = 'self.emitted_by' # e.g. 'self.extractor' abid_rand_src = 'self.id' # e.g. 'self.uuid' or 'self.id' abid_drift_allowed: bool = False # set to True to allow abid_field values to change after a fixed ABID has been issued (NOT RECOMMENDED: means values can drift out of sync from original ABID) read_only_fields = ('id', 'deliver_at', 'name', 'kwargs', 'timeout', 'parent', 'emitted_by', 'on_success', 'on_failure') id = models.UUIDField(primary_key=True, default=uuid.uuid4, null=False, editable=False, unique=True, verbose_name='ID') # disable these fields from inherited models, they're not needed / take up too much room abid = None created_at = None created_by = None created_by_id = None # immutable fields deliver_at = models.DateTimeField(default=timezone.now, null=False, editable=False, unique=True, db_index=True) name = models.CharField(max_length=255, null=False, blank=False, db_index=True) kwargs = models.JSONField(default=dict) timeout = models.IntegerField(null=False, default=60) parent = models.ForeignKey('Event', null=True, on_delete=models.SET_NULL, related_name='child_events') emitted_by = models.ForeignKey(Process, null=False, on_delete=models.PROTECT, related_name='emitted_events') on_success = models.JSONField(null=True) on_failure = models.JSONField(null=True) # mutable fields modified_at = models.DateTimeField(auto_now=True) claimed_proc = models.ForeignKey(Process, null=True, on_delete=models.CASCADE, related_name='claimed_events') claimed_at = models.DateTimeField(null=True) finished_at = models.DateTimeField(null=True) error = models.TextField(null=True) objects: EventManager = EventManager.from_queryset(EventQuerySet)() child_events: models.RelatedManager['Event'] @classmethod def get_next_timestamp(cls): """Get the next monotonically increasing timestamp for the next event.dispatch_at""" latest_event = cls.objects.order_by('-deliver_at').first() ts = timezone.now() if latest_event: assert ts > latest_event.deliver_at, f'Event.deliver_at is not monotonically increasing: {latest_event.deliver_at} > {ts}' return ts @classmethod def dispatch(cls, name: str | EventDict | None = None, event: EventDict | None = None, **event_init_kwargs) -> 'Event': """ Create a new Event and save it to the database. Can be called as either: >>> Event.dispatch(name, {**kwargs}, **event_init_kwargs) # OR >>> Event.dispatch({name, **kwargs}, **event_init_kwargs) """ event_kwargs: EventDict = event or {} if isinstance(name, dict): event_kwargs.update(name) assert isinstance(event_kwargs, dict), 'must be called as Event.dispatch(name, {**kwargs}) or Event.dispatch({name, **kwargs})' event_name: str = name if (isinstance(name, str) and name) else event_kwargs.pop('name') new_event = cls( name=event_name, kwargs=event_kwargs, emitted_by=Process.current(), **event_init_kwargs, ) new_event.save() return new_event def clean(self, *args, **kwargs) -> None: """Fill and validate all the event fields""" # check uuid and deliver_at are set assert self.id, 'Event.id must be set to a valid v4 UUID' if not self.deliver_at: self.deliver_at = self.get_next_timestamp() assert self.deliver_at and (datetime(2024, 12, 8, 12, 0, 0, tzinfo=timezone.utc) < self.deliver_at < datetime(2100, 12, 31, 23, 59, 0, tzinfo=timezone.utc)), ( f'Event.deliver_at must be set to a valid UTC datetime (got Event.deliver_at = {self.deliver_at})') # if name is not set but it's found in the kwargs, move it out of the kwargs to the name field if 'type' in self.kwargs and ((self.name == self.kwargs['type']) or not self.name): self.name = self.kwargs.pop('type') if 'name' in self.kwargs and ((self.name == self.kwargs['name']) or not self.name): self.name = self.kwargs.pop('name') # check name is set and is a valid identifier assert isinstance(self.name, str) and len(self.name) > 3, 'Event.name must be set to a non-empty string' assert self.name.isidentifier(), f'Event.name must be a valid identifier (got Event.name = {self.name})' assert self.name.isupper(), f'Event.name must be in uppercase (got Event.name = {self.name})' # check that kwargs keys and values are valid for key, value in self.kwargs.items(): assert isinstance(key, str), f'Event kwargs keys can only be strings (got Event.kwargs[{key}: {type(key).__name__}])' assert key not in self._meta.get_fields(), f'Event.kwargs cannot contain "{key}" key (Event.kwargs[{key}] conflicts with with reserved attr Event.{key} = {getattr(self, key)})' assert json.dumps(value, sort_keys=True), f'Event can only contain JSON serializable values (got Event.kwargs[{key}]: {type(value).__name__} = {value})' # validate on_success and on_failure are valid event dicts if set if self.on_success: assert isinstance(self.on_success, dict) and self.on_success.get('name', '!invalid').isidentifier(), f'Event.on_success must be a valid event dict (got {self.on_success})' if self.on_failure: assert isinstance(self.on_failure, dict) and self.on_failure.get('name', '!invalid').isidentifier(), f'Event.on_failure must be a valid event dict (got {self.on_failure})' # validate mutable fields like claimed_at, claimed_proc, finished_at are set correctly if self.claimed_at: assert self.claimed_proc, f'Event.claimed_at and Event.claimed_proc must be set together (only found Event.claimed_at = {self.claimed_at})' if self.claimed_proc: assert self.claimed_at, f'Event.claimed_at and Event.claimed_proc must be set together (only found Event.claimed_proc = {self.claimed_proc})' if self.finished_at: assert self.claimed_at, f'If Event.finished_at is set, Event.claimed_at and Event.claimed_proc must also be set (Event.claimed_proc = {self.claimed_proc} and Event.claimed_at = {self.claimed_at})' # validate error is a non-empty string or None if isinstance(self.error, BaseException): self.error = f'{type(self.error).__name__}: {self.error}' if self.error: assert isinstance(self.error, str) and str(self.error).strip(), f'Event.error must be a non-empty string (got Event.error: {type(self.error).__name__} = {self.error})' else: assert self.error is None, f'Event.error must be None or a non-empty string (got Event.error: {type(self.error).__name__} = {self.error})' def save(self, *args, **kwargs): self.clean() return super().save(*args, **kwargs) def reset(self): """Force-update an event to a pending/unclaimed state (without running any of its handlers or callbacks)""" self.claimed_proc = None self.claimed_at = None self.finished_at = None self.error = None self.save() def abort(self): """Force-update an event to a completed/failed state (without running any of its handlers or callbacks)""" self.claimed_proc = Process.current() self.claimed_at = timezone.now() self.finished_at = timezone.now() self.error = 'Aborted' self.save() def __repr__(self) -> str: label = f'[{self.name} {self.kwargs}]' if self.is_finished: label += f' ✅' elif self.claimed_proc: label += f' 🏃' return label def __str__(self) -> str: return repr(self) @property def type(self) -> str: return self.name @property def is_queued(self): return not self.is_claimed and not self.is_finished @property def is_claimed(self): return self.claimed_at is not None @property def is_expired(self): if not self.claimed_at: return False elapsed_time = timezone.now() - self.claimed_at return elapsed_time > timedelta(seconds=self.timeout) @property def is_processing(self): return self.is_claimed and not self.is_finished @property def is_finished(self): return self.finished_at is not None @property def is_failed(self): return self.is_finished and bool(self.error) @property def is_succeeded(self): return self.is_finished and not bool(self.error) def __getattr__(self, key: str): """ Allow access to the event kwargs as attributes e.g. Event(name='CRAWL_CREATE', kwargs={'some_key': 'some_val'}).some_key -> 'some_val' """ return self.kwargs.get(key)