Feat(Escalation policies): A Hackweek project

bin/escalation-policy-seed-data.py

@@ -0,0 +1,391 @@
+#!/usr/bin/env python
+import dataclasses
+import datetime
+import random
+from functools import cached_property
+from typing import Protocol
+
+from sentry.runner import configure
+
+configure()
+
+import click
+
+
+@click.command()
+def seed_policy_data():
+    from sentry.escalation_policies.models.escalation_policy import (
+        EscalationPolicy,
+        EscalationPolicyStep,
+        EscalationPolicyStepRecipient,
+    )
+    from sentry.escalation_policies.models.rotation_schedule import (
+        RotationSchedule,
+        RotationScheduleLayer,
+        RotationScheduleLayerRotationType,
+        RotationScheduleOverride,
+        RotationScheduleUserOrder,
+        ScheduleLayerRestriction,
+    )
+    from sentry.models.organization import Organization
+    from sentry.models.team import Team, TeamStatus
+    from sentry.testutils.factories import Factories  # noqa: S007
+
+    def make_restrictions() -> list[tuple[str, str]]:
+        start = 0
+        result = []
+        for i in range(random.randint(0, 4)):
+            start += random.randint(1, 60 * 5)
+            end = start + random.randint(1, 60 * 5)
+            if end >= 60 * 24:
+                break
+            result.append(
+                (f"{int(start / 60):02d}:{start % 60:02d}", f"{int(end / 60):02d}:{end % 60:02d}")
+            )
+        return result
+
+    class TeamAndUserId(Protocol):
+        team: Team
+        user_id: int
+
+    @dataclasses.dataclass
+    class OrgContext:
+        org_id: int = 1
+        team_count: int = 5
+        user_count: int = 5
+        schedule_count: int = 8
+        policy_count: int = 5
+
+        def ensure_team_and_user_pool(self):
+            for i in range(self.team_count - len(self.teams)):
+                self.teams.append(
+                    Factories.create_team(
+                        organization=self.organization,
+                        name=f"Demo Team {i}",
+                    )
+                )
+            for i in range(self.user_count - len(self.user_ids)):
+                u = Factories.create_user()
+                Factories.create_member(teams=self.teams, user=u, organization=self.organization)
+                self.user_ids.append(u.id)
+            for i in range(self.schedule_count - len(self.schedules)):
+                self.schedules.append(RotationScheduleFactory(org_context=self).build())
+
+            # Fix restrictions by recomputing
+            for schedule in self.schedules:
+                for layer in RotationScheduleLayer.objects.filter(schedule=schedule):
+                    restriction: ScheduleLayerRestriction = {
+                        "Sun": make_restrictions(),
+                        "Mon": make_restrictions(),
+                        "Tue": make_restrictions(),
+                        "Wed": make_restrictions(),
+                        "Thu": make_restrictions(),
+                        "Fri": make_restrictions(),
+                        "Sat": make_restrictions(),
+                    }
+
+                    layer.schedule_layer_restrictions = restriction
+                    layer.save()
+
+        def ensure_policies(self):
+            for i in range(self.policy_count - len(self.policies)):
+                pf = PolicyFactory(self)
+                self.policies.append(pf.build())
+
+        def contextualize_slug_or_user_id(
+            self, model: TeamAndUserId, slug_or_user_id: str | int | None
+        ):
+            if isinstance(slug_or_user_id, str):
+                model.team = next(t for t in self.teams if t.slug == slug_or_user_id)
+            elif isinstance(slug_or_user_id, int):
+                model.user_id = slug_or_user_id
+            else:
+                if random.random() < 0.5:
+                    slug_or_user_id = random.choice(self.teams).slug
+                else:
+                    slug_or_user_id = random.choice(self.user_ids)
+                self.contextualize_slug_or_user_id(model, slug_or_user_id)
+
+        @cached_property
+        def organization(self) -> Organization:
+            return Organization.objects.get(pk=self.org_id)
+
+        @cached_property
+        def policies(self) -> list[EscalationPolicy]:
+            return list(EscalationPolicy.objects.filter(organization=self.organization))
+
+        @cached_property
+        def user_ids(self) -> list[int]:
+            return list(self.organization.member_set.values_list("user_id", flat=True))
+
+        @cached_property
+        def schedules(self) -> list[RotationSchedule]:
+            return list(RotationSchedule.objects.filter(organization=self.organization))
+
+        @cached_property
+        def teams(self) -> list[Team]:
+            return list(
+                Team.objects.filter(
+                    status=TeamStatus.ACTIVE,
+                    organization=self.organization,
+                )
+            )
+
+    @dataclasses.dataclass
+    class RotationScheduleFactory:
+        org_context: OrgContext
+        name: str = ""
+        owner_user_id_or_team_slug: str | int | None = None
+
+        @cached_property
+        def organization(self) -> Organization:
+            return self.org_context.organization
+
+        @cached_property
+        def rotation_schedule(self) -> RotationSchedule:
+            schedule = RotationSchedule(
+                organization=self.organization,
+                name=self.final_name,
+                description=self.final_desc,
+            )
+
+            self.org_context.contextualize_slug_or_user_id(
+                schedule, self.owner_user_id_or_team_slug
+            )
+
+            schedule.save()
+            return schedule
+
+        @cached_property
+        def final_name(self) -> str:
+            c = 1
+            name = self.name or "Rotation Schedule"
+            while True:
+                if not RotationSchedule.objects.filter(
+                    organization=self.organization, name=name
+                ).exists():
+                    break
+                name = name.removesuffix(" " + str(c))
+                c += 1
+                name += " " + str(c)
+            return name
+
+        @cached_property
+        def final_desc(self) -> str:
+            words = [
+                "liberty",
+                "kettle",
+                "courage",
+                "stand",
+                "satisfaction",
+                "compliance",
+                "work",
+                "voyage",
+                "district",
+                "dialogue",
+                "quit",
+                "management",
+                "standard",
+                "remain",
+                "graze",
+                "carpet",
+                "master",
+                "relevance",
+                "machinery",
+                "rubbish",
+                "singer",
+                "presidency",
+                "horn",
+                "activate",
+                "carry",
+                "sphere",
+                "shape",
+                "waiter",
+                "peasant",
+                "string",
+                "result",
+                "rhetoric",
+                "occasion",
+                "thanks",
+                "auction",
+                "dawn",
+                "head",
+                "wind",
+                "cater",
+                "brilliance",
+                "packet",
+                "unfair",
+                "wolf",
+                "loan",
+                "thick",
+                "strikebreaker",
+                "sunrise",
+                "minimum",
+                "resource",
+                "digital",
+            ]
+            return "This is a schedule for " + " ".join(random.choices(words, k=5))
+
+        @cached_property
+        def overrides(self) -> list[RotationScheduleOverride]:
+            start = datetime.datetime.utcnow() - datetime.timedelta(hours=12)
+            results = []
+            for i in range(random.randint(0, 6)):
+                end = start + datetime.timedelta(hours=random.randint(3, 9))
+                user_id = random.choice(self.org_context.user_ids)
+                override = RotationScheduleOverride(
+                    rotation_schedule=self.rotation_schedule,
+                    user_id=user_id,
+                    start_time=start,
+                    end_time=end,
+                )
+                override.save()
+                start = end + datetime.timedelta(hours=random.randint(3, 9))
+                results.append(override)
+            return results
+
+        @cached_property
+        def schedule_layers(self) -> list[RotationScheduleLayer]:
+            results = []
+            for i in range(random.randint(1, 5)):
+                layer = RotationScheduleLayer(
+                    schedule=self.rotation_schedule,
+                    precedence=i,
+                    rotation_type=random.choice(list(RotationScheduleLayerRotationType)),
+                    handoff_time=f"{random.randint(0, 23):02d}:{random.randint(0, 59):02d}",
+                    start_date=(
+                        datetime.datetime.utcnow() - datetime.timedelta(days=random.randint(1, 30))
+                    ).date(),
+                )
+
+                restriction: ScheduleLayerRestriction = {
+                    "Sun": make_restrictions(),
+                    "Mon": make_restrictions(),
+                    "Tue": make_restrictions(),
+                    "Wed": make_restrictions(),
+                    "Thu": make_restrictions(),
+                    "Fri": make_restrictions(),
+                    "Sat": make_restrictions(),
+                }
+
+                layer.schedule_layer_restrictions = restriction
+                layer.save()
+                results.append(layer)
+            return results
+
+        @cached_property
+        def user_orders(self) -> list[RotationScheduleUserOrder]:
+            result = []
+            for layer in self.schedule_layers:
+                user_ids = random.sample(self.org_context.user_ids, random.randint(1, 5))
+                for i, user_id in enumerate(user_ids):
+                    uo = RotationScheduleUserOrder(
+                        schedule_layer=layer,
+                        user_id=user_id,
+                        order=i,
+                    )
+                    uo.save()
+                    result.append(uo)
+            return result
+
+        def build(self):
+            return (
+                self.rotation_schedule,
+                self.schedule_layers,
+                self.user_orders,
+            )[0]
+
+    @dataclasses.dataclass
+    class PolicyFactory:
+        org_context: OrgContext
+        owner_user_id_or_team_slug: str | int | None = None
+        name: str = ""
+        repeat_n_times: int = dataclasses.field(default_factory=lambda: random.randint(1, 5))
+
+        @cached_property
+        def organization(self) -> Organization:
+            return self.org_context.organization
+
+        @cached_property
+        def final_name(self) -> str:
+            c = 1
+            name = self.name or "Escalation Policy"
+            while True:
+                if not EscalationPolicy.objects.filter(
+                    organization=self.organization, name=name
+                ).exists():
+                    break
+                name = name.removesuffix(" " + str(c))
+                c += 1
+                name += " " + str(c)
+            return name
+
+        @cached_property
+        def policy(self) -> EscalationPolicy:
+            policy = EscalationPolicy(
+                organization=self.organization,
+                name=self.final_name,
+                description="",
+                repeat_n_times=self.repeat_n_times,
+            )
+
+            self.org_context.contextualize_slug_or_user_id(policy, self.owner_user_id_or_team_slug)
+
+            policy.save()
+            return policy
+
+        @cached_property
+        def steps(self):
+            result = []
+            for i in range(5):
+                step = EscalationPolicyStep(
+                    policy=self.policy,
+                    step_number=i + 1,
+                    escalate_after_sec=random.randint(1, 5) * 15,
+                )
+                step.save()
+                result.append(step)
+            return result
+
+        @cached_property
+        def recipients(self):
+            result = []
+            for step in self.steps:
+                result.append(
+                    EscalationPolicyStepRecipient(
+                        escalation_policy_step=step,
+                        schedule=random.choice(self.org_context.schedules),
+                    )
+                )
+                result.append(
+                    EscalationPolicyStepRecipient(
+                        escalation_policy_step=step,
+                        team=random.choice(self.org_context.teams),
+                    )
+                )
+                result.append(
+                    EscalationPolicyStepRecipient(
+                        escalation_policy_step=step,
+                        user_id=random.choice(self.org_context.user_ids),
+                    )
+                )
+
+            for recipient in result:
+                recipient.save()
+
+            return result
+
+        def build(self):
+            return (
+                self.policy,
+                self.steps,
+                self.recipients,
+            )[0]
+
+    context = OrgContext()
+    context.ensure_team_and_user_pool()
+    context.ensure_policies()
+
+
+if __name__ == "__main__":
+    seed_policy_data()

migrations_lockfile.txt

@@ -10,6 +10,6 @@ hybridcloud: 0016_add_control_cacheversion
 nodestore: 0002_nodestore_no_dictfield
 remote_subscriptions: 0003_drop_remote_subscription
 replays: 0004_index_together
-sentry: 0747_create_datasecrecywaiver_table
+sentry: 0751_escalation_policies3
 social_auth: 0002_default_auto_field
 uptime: 0006_projectuptimesubscription_name_owner