authentik.tasks.schedules.common

 1import pickle  # nosec
 2from collections.abc import Iterable
 3from dataclasses import dataclass, field
 4from typing import TYPE_CHECKING, Any
 5from uuid import UUID
 6
 7from dramatiq.actor import Actor
 8from psqlextra.types import ConflictAction
 9
10if TYPE_CHECKING:
11    from authentik.tasks.schedules.models import Schedule
12
13
14@dataclass
15class ScheduleSpec:
16    actor: Actor
17    crontab: str
18    paused: bool = False
19    identifier: str | UUID | None = None
20    uid: str | None = None
21
22    args: Iterable[Any] = field(default_factory=tuple)
23    kwargs: dict[str, Any] = field(default_factory=dict)
24    options: dict[str, Any] = field(default_factory=dict)
25
26    rel_obj: Any | None = None
27
28    send_on_save: bool = False
29
30    send_on_startup: bool = False
31
32    def get_args(self) -> bytes:
33        return pickle.dumps(self.args)
34
35    def get_kwargs(self) -> bytes:
36        return pickle.dumps(self.kwargs)
37
38    def get_options(self) -> bytes:
39        options = self.options
40        if self.uid is not None:
41            options["uid"] = self.uid
42        return pickle.dumps(options)
43
44    def update_or_create(self) -> Schedule:
45        from django.contrib.contenttypes.models import ContentType
46
47        from authentik.tasks.schedules.models import Schedule
48
49        update_values = {
50            "_uid": self.uid,
51            "paused": self.paused,
52            "args": self.get_args(),
53            "kwargs": self.get_kwargs(),
54            "options": self.get_options(),
55        }
56        if self.rel_obj is not None:
57            update_values["rel_obj_content_type"] = ContentType.objects.get_for_model(self.rel_obj)
58            update_values["rel_obj_id"] = str(self.rel_obj.pk)
59        create_values = {
60            **update_values,
61            "crontab": self.crontab,
62        }
63
64        schedule = Schedule.objects.on_conflict(
65            ["actor_name", "identifier"],
66            ConflictAction.UPDATE,
67            update_values=update_values,
68        ).insert_and_get(
69            actor_name=self.actor.actor_name,
70            identifier=str(self.identifier),
71            **create_values,
72        )
73
74        return schedule