oncall-engine/tools/migrators/lib/splunk/resources/schedules.py

import datetime
import typing
from dataclasses import dataclass
from typing import Optional
from uuid import uuid4

from lib.constants import ONCALL_SHIFT_WEB_SOURCE
from lib.oncall import types as oncall_types
from lib.oncall.api_client import OnCallAPIClient
from lib.splunk import types
from lib.utils import dt_to_oncall_datetime, duration_to_frequency_and_interval

TIME_ZONE = "UTC"
"""
Note: The Splunk schedule rotations do return a `timezone` attribute, but I don't think
we need to worry about this as the all of the timestamps that we touch are in UTC.
"""


def generate_splunk_schedule_name(
    schedule: types.SplunkScheduleWithTeamAndRotations,
) -> str:
    return f"{schedule['policy']['name']} schedule"


def match_schedule(
    schedule: types.SplunkScheduleWithTeamAndRotations,
    oncall_schedules: list[oncall_types.OnCallSchedule],
    user_id_map: dict[str, str],
) -> None:
    schedule_name = generate_splunk_schedule_name(schedule)
    schedule["name"] = schedule_name

    oncall_schedule = None

    for candidate in oncall_schedules:
        if schedule_name.lower().strip() == candidate["name"].lower().strip():
            oncall_schedule = candidate

    _, errors = Schedule.from_dict(schedule).to_oncall_schedule(user_id_map)

    schedule["migration_errors"] = errors
    schedule["oncall_schedule"] = oncall_schedule


def migrate_schedule(
    schedule: types.SplunkScheduleWithTeamAndRotations,
    user_id_map: dict[str, str],
) -> None:
    if schedule["oncall_schedule"]:
        OnCallAPIClient.delete("schedules/{}".format(schedule["oncall_schedule"]["id"]))

    schedule["oncall_schedule"] = Schedule.from_dict(schedule).migrate(user_id_map)


def _splunk_datetime_to_dt(text: str) -> datetime.datetime:
    """
    Convert a Splunk datetime string to a datetime object.
    """
    return datetime.datetime.strptime(text, "%Y-%m-%dT%H:%M:%SZ")


@dataclass
class Schedule:
    """
    Utility class for converting a Splunk schedule to an OnCall schedule.
    """

    name: str
    rotation_shifts: list["RotationShift"]
    overrides: list["Override"]

    @classmethod
    def from_dict(
        cls, schedule: types.SplunkScheduleWithTeamAndRotations
    ) -> "Schedule":
        """
        Create a Schedule object from a Splunk API response for a schedule.
        """
        rotation_shifts = []
        num_oncall_shift_layers = len(schedule["rotations"])

        for idx, rotation in enumerate(schedule["rotations"]):
            for shift in rotation["shifts"]:
                rotation_shifts.append(
                    RotationShift.from_dict(shift, num_oncall_shift_layers - idx)
                )

        return cls(
            name=generate_splunk_schedule_name(schedule),
            rotation_shifts=rotation_shifts,
            overrides=[
                Override.from_dict(override) for override in schedule["overrides"]
            ],
        )

    def to_oncall_schedule(
        self, user_id_map: dict[str, str]
    ) -> tuple[Optional[dict], list[str]]:
        """
        Convert a Schedule object to an OnCall schedule.
        Note that it also returns shifts, but these are not created at the same time as the schedule (see migrate method for more info).
        """
        shifts = []
        errors = []
        for rotation_shift in self.rotation_shifts:
            # Check if all users in the rotation exist in OnCall
            missing_user_ids = [
                user_id
                for user_id in rotation_shift.user_ids
                if user_id_map.get(user_id) is None
            ]
            if missing_user_ids:
                errors.append(
                    f"{rotation_shift.name}: Users with IDs {missing_user_ids} not found. The user(s) don't seem to exist in Grafana."
                )
                continue

            shifts.append(rotation_shift.to_oncall_shift(user_id_map))

        for override in self.overrides:
            user_id = override.user_id

            if user_id_map.get(user_id) is None:
                errors.append(
                    f"Override: User with ID '{user_id}' not found. The user doesn't seem to exist in Grafana."
                )
                continue

            shifts.append(override.to_oncall_shift(user_id_map))

        if errors:
            return None, errors

        return {
            "name": self.name,
            "type": "web",
            "team_id": None,
            "time_zone": TIME_ZONE,
            "shifts": shifts,
        }, []

    def migrate(self, user_id_map: dict[str, str]) -> dict:
        """
        Create an OnCall schedule and its shifts.
        First create the shifts, then create a schedule with shift IDs provided.
        """

        schedule, errors = self.to_oncall_schedule(user_id_map)
        assert not errors, "Unexpected errors: {}".format(errors)

        # Create shifts in OnCall
        shift_ids = [
            OnCallAPIClient.create("on_call_shifts", shift)["id"]
            for shift in schedule["shifts"]
        ]

        # Create schedule in OnCall with shift IDs provided
        schedule["shifts"] = shift_ids
        new_schedule = OnCallAPIClient.create("schedules", schedule)

        return new_schedule


@dataclass
class RotationShift:
    """
    Utility class for converting a Splunk schedule rotation layer to OnCall shifts.
    """

    name: str
    level: int

    shift_type: typing.Literal["std", "pho", "cstm"]
    start: datetime.datetime
    duration: datetime.timedelta
    mask: types.SplunkRotationShiftMask
    mask2: typing.Optional[types.SplunkRotationShiftMask]
    mask3: typing.Optional[types.SplunkRotationShiftMask]

    user_ids: list[str]

    MONDAY = "m"
    TUESDAY = "t"
    WEDNESDAY = "w"
    THURSDAY = "th"
    FRIDAY = "f"
    SATURDAY = "sa"
    SUNDAY = "su"

    SPLUNK_TO_ONCALL_DAY_MASK_MAP = {
        SUNDAY: "SU",
        MONDAY: "MO",
        TUESDAY: "TU",
        WEDNESDAY: "WE",
        THURSDAY: "TH",
        FRIDAY: "FR",
        SATURDAY: "SA",
    }

    @classmethod
    def from_dict(
        cls, rotation_shift: types.SplunkRotationShift, level: int
    ) -> "RotationShift":
        """
        Create a RotationShift object from a Splunk API response for a rotation.
        Converts Splunk datetime strings to datetime objects for easier manipulation.
        """
        return cls(
            name=rotation_shift["label"],
            level=level,
            shift_type=rotation_shift["shifttype"],
            start=_splunk_datetime_to_dt(rotation_shift["start"]),
            duration=datetime.timedelta(days=rotation_shift["duration"]),
            mask=rotation_shift["mask"],
            mask2=rotation_shift.get("mask2"),
            mask3=rotation_shift.get("mask3"),
            user_ids=[u["username"] for u in rotation_shift["shiftMembers"]],
        )

    def _construct_datetime_from_date_and_mask_time(
        self,
        date: datetime.date,
        mask: types.SplunkRotationShiftMask,
        mask_key: typing.Literal["start", "end"],
    ) -> datetime.datetime:
        mask_time = mask["time"][0][mask_key]
        return datetime.datetime.combine(
            date,
            datetime.time(hour=mask_time["hour"], minute=mask_time["minute"]),
        )

    def _calculate_partial_day_duration_from_mask(self) -> datetime.timedelta:
        """
        Calculate the duration of the shift based on the mask.
        """
        today = datetime.date.today()

        start_dt = self._construct_datetime_from_date_and_mask_time(
            today, self.mask, "start"
        )
        end_dt = self._construct_datetime_from_date_and_mask_time(
            today, self.mask, "end"
        )

        return end_dt - start_dt

    def _calculate_by_days_from_partial_day_shift_mask(self) -> list[str]:
        """
        Calculate the days of the week the shift occurs based on the mask.
        """
        return [
            self.SPLUNK_TO_ONCALL_DAY_MASK_MAP[day]
            for day, is_active in self.mask["day"].items()
            if is_active
        ]

    def _next_day_of_week(
        self, starting_date: datetime.date, day_of_week: str
    ) -> datetime.date:
        # Define a mapping of day abbreviations to their corresponding datetime weekday values
        SPLUNK_DAY_ABBREVIATION_TO_DATETIME_WEEKDAY_IDX_MAP = {
            self.MONDAY: 0,
            self.TUESDAY: 1,
            self.WEDNESDAY: 2,
            self.THURSDAY: 3,
            self.FRIDAY: 4,
            self.SATURDAY: 5,
            self.SUNDAY: 6,
        }

        # Calculate the difference between starting_date's weekday and the desired weekday
        days_until_next_day = (
            SPLUNK_DAY_ABBREVIATION_TO_DATETIME_WEEKDAY_IDX_MAP[day_of_week]
            - starting_date.weekday()
            + 7
        ) % 7

        # Calculate the date of the next desired day of the week
        return starting_date + datetime.timedelta(days=days_until_next_day)

    def _get_sole_active_day_from_mask(
        self, mask: types.SplunkRotationShiftMask
    ) -> str:
        """
        making a big assumption here, but it looks like for multi-day shifts, mask and mask3
        only have one active day each
        """
        return [day for day, is_active in mask["day"].items() if is_active][0]

    def _calculate_multi_day_duration_from_masks(self) -> datetime.timedelta:
        start_mask = self.mask
        end_mask = self.mask3

        today = datetime.date.today()
        shift_start_date = self._next_day_of_week(
            today, self._get_sole_active_day_from_mask(start_mask)
        )
        shift_end_date = self._next_day_of_week(
            shift_start_date, self._get_sole_active_day_from_mask(end_mask)
        )

        shift_start_dt = self._construct_datetime_from_date_and_mask_time(
            shift_start_date, start_mask, "start"
        )
        shift_end_dt = self._construct_datetime_from_date_and_mask_time(
            shift_end_date, end_mask, "end"
        )
        return shift_end_dt - shift_start_dt

    def to_oncall_shift(self, user_id_map: dict[str, str]) -> typing.Dict:
        frequency, interval = duration_to_frequency_and_interval(self.duration)
        start = dt_to_oncall_datetime(self.start)

        duration: datetime.timedelta
        extra_kwargs = {}

        if self.shift_type == "std":
            duration = self.duration
        elif self.shift_type == "pho":
            duration = self._calculate_partial_day_duration_from_mask()
            extra_kwargs[
                "by_day"
            ] = self._calculate_by_days_from_partial_day_shift_mask()
        elif self.shift_type == "cstm":
            num_days = self.duration.days

            if num_days != 7:
                # NOTE: we don't currently support multi-day Splunk shifts with a "hand-off" greater than one week
                # https://raintank-corp.slack.com/archives/C04JCU51NF8/p1714581046981109?thread_ts=1714580582.883559&cid=C04JCU51NF8
                raise ValueError(
                    f"Multi-day shifts with a duration greater than 7 days are not supported: {num_days} days"
                )

            duration = self._calculate_multi_day_duration_from_masks()
        else:
            raise ValueError(f"Unknown shift type: {self.shift_type}")

        return {
            "name": self.name,
            "team_id": None,
            "level": self.level,
            "type": "rolling_users",
            "rotation_start": start,
            "start": start,
            "until": None,
            "duration": int(duration.total_seconds()),
            "frequency": frequency,
            "interval": interval,
            "rolling_users": [[user_id_map[user_id]] for user_id in self.user_ids],
            "start_rotation_from_user_index": 0,
            "week_start": "MO",
            "time_zone": TIME_ZONE,
            "source": ONCALL_SHIFT_WEB_SOURCE,
            **extra_kwargs,
        }


@dataclass
class Override:
    start: datetime.datetime
    end: datetime.datetime
    user_id: str

    @classmethod
    def from_dict(cls, override: types.SplunkScheduleOverride) -> "Override":
        # convert start and end to datetime objects in UTC
        return cls(
            start=datetime.datetime.fromisoformat(override["start"]).astimezone(
                datetime.timezone.utc
            ),
            end=datetime.datetime.fromisoformat(override["end"]).astimezone(
                datetime.timezone.utc
            ),
            user_id=override["overrideOnCallUser"]["username"],
        )

    def to_oncall_shift(self, user_id_map: dict[str, str]) -> dict:
        start = dt_to_oncall_datetime(self.start)
        duration = int((self.end - self.start).total_seconds())
        user_id = user_id_map[self.user_id]

        return {
            "name": uuid4().hex,
            "team_id": None,
            "type": "override",
            "time_zone": TIME_ZONE,
            "start": start,
            "duration": duration,
            "rotation_start": start,
            "users": [user_id],
            "source": ONCALL_SHIFT_WEB_SOURCE,
        }