e4728ea69f
This PR adds support for migrating data from OpsGenie to Grafana IRM. Closes https://github.com/grafana/irm/issues/1179
679 lines
24 KiB
Python
679 lines
24 KiB
Python
import datetime
|
|
import re
|
|
import typing
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
from typing import Optional
|
|
from uuid import uuid4
|
|
|
|
from lib.common.report import TAB
|
|
from lib.constants import ONCALL_SHIFT_WEB_SOURCE
|
|
from lib.oncall.api_client import OnCallAPIClient
|
|
from lib.pagerduty.config import (
|
|
PAGERDUTY_FILTER_SCHEDULE_REGEX,
|
|
PAGERDUTY_FILTER_TEAM,
|
|
PAGERDUTY_FILTER_USERS,
|
|
SCHEDULE_MIGRATION_MODE,
|
|
SCHEDULE_MIGRATION_MODE_ICAL,
|
|
SCHEDULE_MIGRATION_MODE_WEB,
|
|
VERBOSE_LOGGING,
|
|
)
|
|
from lib.utils import dt_to_oncall_datetime, duration_to_frequency_and_interval
|
|
|
|
|
|
def filter_schedules(
|
|
schedules: typing.List[typing.Dict[str, typing.Any]]
|
|
) -> typing.List[typing.Dict[str, typing.Any]]:
|
|
"""
|
|
Filter schedules based on configured filters.
|
|
|
|
If multiple filters are specified, a schedule only needs to match one of them
|
|
to be included (OR operation between filters).
|
|
"""
|
|
if not any(
|
|
[PAGERDUTY_FILTER_TEAM, PAGERDUTY_FILTER_USERS, PAGERDUTY_FILTER_SCHEDULE_REGEX]
|
|
):
|
|
return schedules # No filters specified, return all
|
|
|
|
filtered_schedules = []
|
|
filtered_out = 0
|
|
filtered_reasons = {}
|
|
|
|
for schedule in schedules:
|
|
matches_any_filter = False
|
|
reasons = []
|
|
|
|
# Filter by team
|
|
if PAGERDUTY_FILTER_TEAM:
|
|
teams = schedule.get("teams", [])
|
|
if any(team["summary"] == PAGERDUTY_FILTER_TEAM for team in teams):
|
|
matches_any_filter = True
|
|
else:
|
|
reasons.append(
|
|
f"No teams found for team filter: {PAGERDUTY_FILTER_TEAM}"
|
|
)
|
|
|
|
# Filter by users
|
|
if PAGERDUTY_FILTER_USERS:
|
|
schedule_users = set()
|
|
for layer in schedule.get("schedule_layers", []):
|
|
for user in layer.get("users", []):
|
|
schedule_users.add(user["user"]["id"])
|
|
|
|
if any(user_id in schedule_users for user_id in PAGERDUTY_FILTER_USERS):
|
|
matches_any_filter = True
|
|
else:
|
|
reasons.append(
|
|
f"No users found for user filter: {','.join(PAGERDUTY_FILTER_USERS)}"
|
|
)
|
|
|
|
# Filter by name regex
|
|
if PAGERDUTY_FILTER_SCHEDULE_REGEX:
|
|
if re.match(PAGERDUTY_FILTER_SCHEDULE_REGEX, schedule["name"]):
|
|
matches_any_filter = True
|
|
else:
|
|
reasons.append(
|
|
f"Schedule regex filter: {PAGERDUTY_FILTER_SCHEDULE_REGEX}"
|
|
)
|
|
|
|
if matches_any_filter:
|
|
filtered_schedules.append(schedule)
|
|
else:
|
|
filtered_out += 1
|
|
filtered_reasons[schedule["id"]] = reasons
|
|
|
|
if filtered_out > 0:
|
|
summary = f"Filtered out {filtered_out} schedules"
|
|
print(summary)
|
|
|
|
# Only print detailed reasons in verbose mode
|
|
if VERBOSE_LOGGING:
|
|
for schedule_id, reasons in filtered_reasons.items():
|
|
print(f"{TAB}Schedule {schedule_id}: {', '.join(reasons)}")
|
|
|
|
return filtered_schedules
|
|
|
|
|
|
def match_schedule(
|
|
schedule: dict, oncall_schedules: list[dict], user_id_map: dict[str, str]
|
|
) -> None:
|
|
oncall_schedule = None
|
|
for candidate in oncall_schedules:
|
|
if schedule["name"].lower().strip() == candidate["name"].lower().strip():
|
|
oncall_schedule = candidate
|
|
|
|
schedule["migration_errors"] = []
|
|
if SCHEDULE_MIGRATION_MODE == SCHEDULE_MIGRATION_MODE_WEB:
|
|
_, errors = Schedule.from_dict(schedule).to_oncall_schedule(user_id_map)
|
|
schedule["migration_errors"] = errors
|
|
|
|
schedule["oncall_schedule"] = oncall_schedule
|
|
|
|
|
|
def migrate_schedule(schedule: dict, user_id_map: dict[str, str]) -> None:
|
|
if schedule["oncall_schedule"]:
|
|
OnCallAPIClient.delete("schedules/{}".format(schedule["oncall_schedule"]["id"]))
|
|
|
|
if SCHEDULE_MIGRATION_MODE == SCHEDULE_MIGRATION_MODE_WEB:
|
|
# Migrate shifts
|
|
oncall_schedule = Schedule.from_dict(schedule).migrate(user_id_map)
|
|
elif SCHEDULE_MIGRATION_MODE == SCHEDULE_MIGRATION_MODE_ICAL:
|
|
# Migrate using ICal URL
|
|
payload = {
|
|
"name": schedule["name"],
|
|
"type": "ical",
|
|
"ical_url_primary": schedule["http_cal_url"],
|
|
"team_id": None,
|
|
}
|
|
oncall_schedule = OnCallAPIClient.create("schedules", payload)
|
|
else:
|
|
raise ValueError("Invalid schedule migration mode")
|
|
|
|
schedule["oncall_schedule"] = oncall_schedule
|
|
|
|
|
|
def _pd_datetime_to_dt(text: str) -> datetime.datetime:
|
|
"""
|
|
Convert a PagerDuty datetime string to a datetime object.
|
|
"""
|
|
dt = datetime.datetime.strptime(text, "%Y-%m-%dT%H:%M:%SZ")
|
|
return dt.replace(tzinfo=datetime.timezone.utc)
|
|
|
|
|
|
@dataclass
|
|
class Schedule:
|
|
"""
|
|
Utility class for converting a PagerDuty schedule to an OnCall schedule.
|
|
A PagerDuty schedule has multiple layers, each with a rotation of users.
|
|
"""
|
|
|
|
name: str
|
|
time_zone: str
|
|
layers: list["Layer"]
|
|
overrides: list["Override"]
|
|
|
|
@classmethod
|
|
def from_dict(cls, schedule: dict) -> "Schedule":
|
|
"""
|
|
Create a Schedule object from a PagerDuty API response for a schedule.
|
|
"""
|
|
|
|
layers = []
|
|
# PagerDuty API returns layers in reverse order (e.g. Layer 3, Layer 2, Layer 1)
|
|
for level, layer_dict in enumerate(
|
|
reversed(schedule["schedule_layers"]), start=1
|
|
):
|
|
layer = Layer.from_dict(layer_dict, level)
|
|
|
|
# skip any layers that have already ended
|
|
if layer.end and layer.end < datetime.datetime.now(datetime.timezone.utc):
|
|
continue
|
|
|
|
layers.append(layer)
|
|
|
|
overrides = []
|
|
for override in schedule["overrides"]:
|
|
overrides.append(Override.from_dict(override))
|
|
|
|
return cls(
|
|
name=schedule["name"],
|
|
time_zone=schedule["time_zone"],
|
|
layers=layers,
|
|
overrides=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 layer in self.layers:
|
|
# Check if all users in the layer exist in PD
|
|
deactivated_user_ids = [
|
|
user_id for user_id in layer.user_ids if user_id not in user_id_map
|
|
]
|
|
if deactivated_user_ids:
|
|
errors.append(
|
|
f"{layer.name}: Users with IDs {deactivated_user_ids} not found. The users probably have been deactivated in PagerDuty."
|
|
)
|
|
continue
|
|
|
|
# A single PagerDuty layer can result in multiple OnCall shifts
|
|
layer_shifts, error = layer.to_oncall_shifts(user_id_map)
|
|
|
|
if layer_shifts:
|
|
shifts += layer_shifts
|
|
|
|
if error:
|
|
error_text = f"{layer.name}: {error}"
|
|
|
|
# If a layer has a single user, it's likely can be easily tweaked in PD to make it possible to migrate
|
|
if len(set(layer.user_ids)) == 1:
|
|
error_text += " Layer has a single user, consider simplifying the rotation in PD."
|
|
|
|
errors.append(error_text)
|
|
|
|
for override in self.overrides:
|
|
if override.user_id not in user_id_map:
|
|
errors.append(
|
|
f"Override: User with ID '{override.user_id}' not found. The user probably has been deactivated in PagerDuty."
|
|
)
|
|
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": self.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 = []
|
|
for shift in schedule["shifts"]:
|
|
created_shift = OnCallAPIClient.create("on_call_shifts", shift)
|
|
shift_ids.append(created_shift["id"])
|
|
|
|
# Create schedule in OnCall with shift IDs provided
|
|
schedule["shifts"] = shift_ids
|
|
new_schedule = OnCallAPIClient.create("schedules", schedule)
|
|
|
|
return new_schedule
|
|
|
|
|
|
@dataclass
|
|
class Layer:
|
|
"""
|
|
Utility class for converting a PagerDuty schedule layer to OnCall shifts.
|
|
"""
|
|
|
|
name: str
|
|
level: int
|
|
|
|
rotation_virtual_start: datetime.datetime
|
|
rotation_turn_length: datetime.timedelta
|
|
|
|
start: datetime.datetime
|
|
end: Optional[datetime.datetime]
|
|
|
|
user_ids: list[str]
|
|
restrictions: list["Restriction"]
|
|
|
|
@classmethod
|
|
def from_dict(cls, layer: dict, level: int) -> "Layer":
|
|
"""
|
|
Create a Layer object from a PagerDuty API response for a schedule layer.
|
|
Converts PagerDuty datetime strings to datetime objects for easier manipulation.
|
|
"""
|
|
return cls(
|
|
name=layer["name"],
|
|
level=level,
|
|
rotation_virtual_start=_pd_datetime_to_dt(layer["rotation_virtual_start"]),
|
|
rotation_turn_length=datetime.timedelta(
|
|
seconds=layer["rotation_turn_length_seconds"]
|
|
),
|
|
start=_pd_datetime_to_dt(layer["start"]),
|
|
end=_pd_datetime_to_dt(layer["end"]) if layer["end"] else None,
|
|
user_ids=[u["user"]["id"] for u in layer["users"]],
|
|
restrictions=[Restriction.from_dict(r) for r in layer["restrictions"]],
|
|
)
|
|
|
|
def to_oncall_shifts(
|
|
self, user_id_map: dict[str, str]
|
|
) -> tuple[Optional[list[dict]], Optional[str]]:
|
|
frequency, interval = duration_to_frequency_and_interval(
|
|
self.rotation_turn_length
|
|
)
|
|
rolling_users = []
|
|
for user_id in self.user_ids:
|
|
oncall_user_id = user_id_map[user_id]
|
|
rolling_users.append([oncall_user_id])
|
|
|
|
if not self.restrictions:
|
|
return [
|
|
{
|
|
"name": uuid4().hex,
|
|
"level": self.level,
|
|
"type": "rolling_users",
|
|
"rotation_start": dt_to_oncall_datetime(self.start),
|
|
"until": dt_to_oncall_datetime(self.end) if self.end else None,
|
|
"start": dt_to_oncall_datetime(self.rotation_virtual_start),
|
|
"duration": int(self.rotation_turn_length.total_seconds()),
|
|
"frequency": frequency,
|
|
"interval": interval,
|
|
"rolling_users": rolling_users,
|
|
"start_rotation_from_user_index": 0,
|
|
"week_start": "MO",
|
|
"time_zone": "UTC",
|
|
"source": ONCALL_SHIFT_WEB_SOURCE,
|
|
}
|
|
], None
|
|
|
|
restrictions_type = self.restrictions[0].type
|
|
|
|
if (frequency, restrictions_type) in (
|
|
("daily", Restriction.Type.DAILY),
|
|
("weekly", Restriction.Type.WEEKLY),
|
|
):
|
|
# TODO: some of this can use by_day?
|
|
shifts, _ = self._generate_shifts(unique=False)
|
|
shifts = [(s[0], s[1], "MO", None) for s in shifts]
|
|
|
|
elif frequency == "weekly" and restrictions_type == Restriction.Type.DAILY:
|
|
shifts, is_split = self._generate_shifts(unique=True)
|
|
|
|
if is_split:
|
|
return (
|
|
None,
|
|
f"Cannot migrate {interval}-weekly rotation with daily restrictions that are split by handoff.",
|
|
)
|
|
|
|
# repeat ["MO", "TU", "WE", "TH", "FR", "SA", "SU"] shift for the number of weeks in the rotation
|
|
shifts_for_multiple_weeks = []
|
|
for shift in shifts:
|
|
for week in range(interval):
|
|
start = shift[0] + datetime.timedelta(weeks=week)
|
|
end = shift[1] + datetime.timedelta(weeks=week)
|
|
shifts_for_multiple_weeks.append((start, end))
|
|
|
|
shifts = [
|
|
(
|
|
shift[0],
|
|
shift[1],
|
|
["MO", "TU", "WE", "TH", "FR", "SA", "SU"][
|
|
shift[0].date().weekday()
|
|
],
|
|
["MO", "TU", "WE", "TH", "FR", "SA", "SU"],
|
|
)
|
|
for shift in shifts_for_multiple_weeks
|
|
]
|
|
|
|
elif (
|
|
frequency == "daily"
|
|
and restrictions_type == Restriction.Type.WEEKLY
|
|
and interval == 1
|
|
):
|
|
# the only case when it's possible to migrate a daily rotation with weekly restrictions
|
|
# is when the restrictions start at the same time as the shift start
|
|
# and the restrictions are a multiple of 24 hours
|
|
restrictions = Restriction.merge_restrictions(self.restrictions)
|
|
for restriction in restrictions:
|
|
if (
|
|
not restriction.start_time_of_day
|
|
== self.rotation_virtual_start.time()
|
|
):
|
|
return (
|
|
None,
|
|
f"Cannot migrate {interval}-daily rotation with weekly restrictions that start at a different time than the shift start.",
|
|
)
|
|
if not restriction.duration % datetime.timedelta(
|
|
days=1
|
|
) == datetime.timedelta(0):
|
|
return (
|
|
None,
|
|
f"Cannot migrate {interval}-daily rotation with weekly restrictions that have durations that are not a multiple of a 24 hours.",
|
|
)
|
|
|
|
# get the first restriction and use its start time as the start of the shift
|
|
restriction, shift_start = Restriction.current_or_next_restriction(
|
|
restrictions, self.rotation_virtual_start
|
|
)
|
|
shift_end = shift_start + datetime.timedelta(days=1)
|
|
|
|
# determine which days of the week are covered
|
|
by_day = set()
|
|
for restriction in restrictions:
|
|
days = restriction.duration // datetime.timedelta(days=1)
|
|
|
|
for day in range(days):
|
|
weekday = ["MO", "TU", "WE", "TH", "FR", "SA", "SU"][
|
|
(restriction.start_day_of_week + day) % 7
|
|
]
|
|
by_day.add(weekday)
|
|
|
|
# sort by_day so that the order is always the same
|
|
by_day = sorted(
|
|
list(by_day),
|
|
key=lambda d: ["MO", "TU", "WE", "TH", "FR", "SA", "SU"].index(d),
|
|
)
|
|
|
|
shifts = [(shift_start, shift_end, "MO", by_day)]
|
|
|
|
else:
|
|
restrictions_type_verbal = (
|
|
"daily" if restrictions_type == Restriction.Type.DAILY else "weekly"
|
|
)
|
|
return (
|
|
None,
|
|
f"Cannot migrate {interval}-{frequency} rotation with {restrictions_type_verbal} restrictions.",
|
|
)
|
|
|
|
payloads = []
|
|
for shift in shifts:
|
|
payload = {
|
|
"name": uuid4().hex,
|
|
"level": self.level,
|
|
"type": "rolling_users",
|
|
"rotation_start": dt_to_oncall_datetime(self.start),
|
|
"until": dt_to_oncall_datetime(self.end) if self.end else None,
|
|
"start": dt_to_oncall_datetime(shift[0]),
|
|
"duration": int((shift[1] - shift[0]).total_seconds()),
|
|
"frequency": frequency,
|
|
"interval": interval,
|
|
"by_day": shift[3],
|
|
"rolling_users": rolling_users,
|
|
"start_rotation_from_user_index": 0,
|
|
"week_start": shift[2],
|
|
"time_zone": "UTC",
|
|
"source": ONCALL_SHIFT_WEB_SOURCE,
|
|
}
|
|
payloads.append(payload)
|
|
return payloads, None
|
|
|
|
def _generate_shifts(
|
|
self, unique: bool = True
|
|
) -> tuple[list[tuple[datetime.datetime, datetime.datetime]], bool]:
|
|
"""
|
|
Returns a list of (start, end) tuples representing the shifts in this layer.
|
|
Note that these are not the actual shifts for OnCall API but rather a list of
|
|
unique intervals generated by traversing the restrictions.
|
|
|
|
Also returns a boolean indicating whether there are restrictions split by on-call handoff.
|
|
"""
|
|
|
|
start = self.rotation_virtual_start
|
|
end = self.rotation_virtual_start + self.rotation_turn_length
|
|
|
|
# Convert restrictions to weekly restrictions, then merge overlapping ones
|
|
restrictions = []
|
|
for restriction in self.restrictions:
|
|
restrictions += restriction.to_weekly_restrictions()
|
|
restrictions = Restriction.merge_restrictions(restrictions)
|
|
|
|
is_split = False
|
|
current = start
|
|
shifts = []
|
|
unique_shift_times = []
|
|
|
|
while current < end:
|
|
restriction, restriction_start = Restriction.current_or_next_restriction(
|
|
restrictions, current
|
|
)
|
|
restriction_end = restriction_start + restriction.duration
|
|
|
|
shift_start = max(current, restriction_start)
|
|
shift_end = min(restriction_end, end)
|
|
|
|
# If the next restriction starts after the end of the rotation or shift is empty, we're done
|
|
if restriction_start > end or shift_start == shift_end:
|
|
break
|
|
|
|
# Check if restriction is split by handoff
|
|
if (shift_start == start and shift_start > restriction_start) or (
|
|
shift_end == end and shift_end < restriction_end
|
|
):
|
|
is_split = True
|
|
|
|
shift = (shift_start, shift_end)
|
|
|
|
# check that we haven't already added this shift
|
|
if (
|
|
not unique
|
|
or (shift[0].time(), shift[1].time()) not in unique_shift_times
|
|
):
|
|
shifts.append(shift)
|
|
unique_shift_times.append((shift[0].time(), shift[1].time()))
|
|
|
|
current = shift_end
|
|
|
|
return shifts, is_split
|
|
|
|
|
|
@dataclass
|
|
class Restriction:
|
|
"""
|
|
Utility class for representing a restriction on a rotation in PagerDuty.
|
|
"""
|
|
|
|
class Type(Enum):
|
|
DAILY = "daily_restriction"
|
|
WEEKLY = "weekly_restriction"
|
|
|
|
type: Type
|
|
start_time_of_day: datetime.time
|
|
duration: datetime.timedelta
|
|
start_day_of_week: Optional[int] # this is only present for weekly restrictions
|
|
|
|
@classmethod
|
|
def from_dict(cls, restriction: dict) -> "Restriction":
|
|
"""
|
|
Create a Restriction object from PagerDuty's API representation.
|
|
Converts PagerDuty datetime strings to datetime objects for easier manipulation as well.
|
|
"""
|
|
|
|
# PagerDuty's API uses 1-indexed days of the week, converting to 0-indexed for ease of use
|
|
start_day_of_week = restriction.get("start_day_of_week")
|
|
if start_day_of_week is not None:
|
|
start_day_of_week -= 1
|
|
|
|
return cls(
|
|
type=Restriction.Type(restriction["type"]),
|
|
start_time_of_day=datetime.time.fromisoformat(
|
|
restriction["start_time_of_day"]
|
|
),
|
|
duration=datetime.timedelta(seconds=restriction["duration_seconds"]),
|
|
start_day_of_week=start_day_of_week,
|
|
)
|
|
|
|
def to_weekly_restrictions(self) -> list["Restriction"]:
|
|
"""
|
|
Convert a daily restriction to a list of weekly restrictions.
|
|
Daily restriction is basically 7 weekly restrictions with the same start time and duration,
|
|
but different days of the week (e.g. 9am-5pm daily restriction is the same as 9am-5pm on Monday, 9am-5pm on Tuesday, etc.)
|
|
|
|
Converting to weekly restrictions makes it easier to work with restrictions and only care about weekly ones.
|
|
"""
|
|
|
|
if self.type == Restriction.Type.WEEKLY:
|
|
return [self]
|
|
else:
|
|
return [
|
|
Restriction(
|
|
type=Restriction.Type.WEEKLY,
|
|
start_time_of_day=self.start_time_of_day,
|
|
duration=self.duration,
|
|
start_day_of_week=day,
|
|
)
|
|
for day in range(7)
|
|
]
|
|
|
|
@staticmethod
|
|
def merge_restrictions(restrictions: list["Restriction"]) -> list["Restriction"]:
|
|
"""
|
|
Merge a list of weekly restrictions into a list of the fewest possible weekly restrictions that cover the same time period.
|
|
Example: (9am - 5pm restriction on Monday, 10am - 4pm restriction on Monday) -> (9am - 5pm restriction on Monday).
|
|
|
|
Only works on weekly restrictions, as daily restrictions are converted to weekly restrictions first.
|
|
"""
|
|
|
|
assert all(r.type == Restriction.Type.WEEKLY for r in restrictions)
|
|
|
|
restrictions = sorted(
|
|
restrictions, key=lambda r: (r.start_day_of_week, r.start_time_of_day)
|
|
)
|
|
merged = []
|
|
|
|
for restriction in restrictions:
|
|
restriction_start = datetime.datetime.combine(
|
|
datetime.date.min
|
|
+ datetime.timedelta(days=restriction.start_day_of_week),
|
|
restriction.start_time_of_day,
|
|
)
|
|
restriction_end = restriction_start + restriction.duration
|
|
|
|
if not merged:
|
|
merged.append(restriction)
|
|
continue
|
|
|
|
last = merged[-1]
|
|
last_start = datetime.datetime.combine(
|
|
datetime.date.min + datetime.timedelta(days=last.start_day_of_week),
|
|
last.start_time_of_day,
|
|
)
|
|
last_end = last_start + last.duration
|
|
|
|
if last_end < restriction_start:
|
|
merged.append(restriction)
|
|
else:
|
|
restriction.start_day_of_week = last_start.weekday()
|
|
restriction.start_time_of_day = last_start.time()
|
|
restriction.duration = max(last_end, restriction_end) - last_start
|
|
merged = merged[:-1] + [restriction]
|
|
|
|
return merged
|
|
|
|
@staticmethod
|
|
def current_or_next_restriction(
|
|
restrictions: list["Restriction"], dt: datetime.datetime
|
|
) -> tuple["Restriction", datetime.datetime]:
|
|
"""
|
|
Get the current or next restriction for a given datetime.
|
|
This is useful for finding all the restrictions that apply to a given rotation shift.
|
|
"""
|
|
assert all(r.type == Restriction.Type.WEEKLY for r in restrictions)
|
|
|
|
for weeks in (-1, 0, 1): # check last week, this week, and next week
|
|
for restriction in restrictions:
|
|
restriction_date = (
|
|
dt.date()
|
|
- datetime.timedelta(days=dt.weekday())
|
|
+ datetime.timedelta(days=restriction.start_day_of_week)
|
|
+ datetime.timedelta(weeks=weeks)
|
|
)
|
|
restriction_start = datetime.datetime.combine(
|
|
restriction_date,
|
|
restriction.start_time_of_day,
|
|
tzinfo=datetime.timezone.utc,
|
|
)
|
|
restriction_end = restriction_start + restriction.duration
|
|
|
|
if restriction_end > dt:
|
|
return restriction, restriction_start
|
|
|
|
# there should always be a restriction
|
|
raise ValueError("No restriction found for given datetime")
|
|
|
|
|
|
@dataclass
|
|
class Override:
|
|
start: datetime.datetime
|
|
end: datetime.datetime
|
|
user_id: str
|
|
|
|
@classmethod
|
|
def from_dict(cls, override: dict) -> "Override":
|
|
# convert start and end to datetime objects in UTC
|
|
start = datetime.datetime.fromisoformat(override["start"]).astimezone(
|
|
datetime.timezone.utc
|
|
)
|
|
end = datetime.datetime.fromisoformat(override["end"]).astimezone(
|
|
datetime.timezone.utc
|
|
)
|
|
|
|
return cls(start=start, end=end, user_id=override["user"]["id"])
|
|
|
|
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": "UTC",
|
|
"start": start,
|
|
"duration": duration,
|
|
"rotation_start": start,
|
|
"users": [user_id],
|
|
"source": ONCALL_SHIFT_WEB_SOURCE,
|
|
}
|