sentry/bin/escalation-policy-seed-data.py

#!/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()