oncall-engine/dev/scripts/generate-fake-data/main.py

import argparse
import asyncio
import math
import random
import typing
import uuid
from datetime import datetime

import aiohttp
from faker import Faker
from tqdm.asyncio import tqdm

fake = Faker()

TEAMS_USERS_COMMAND = "generate_teams_and_users"
SCHEDULES_ONCALL_SHIFTS_COMMAND = "generate_schedules_and_oncall_shifts"

GRAFANA_API_URL = None
ONCALL_API_URL = None
ONCALL_API_TOKEN = None


class OnCallApiUser(typing.TypedDict):
    id: str


class OnCallApiOnCallShift(typing.TypedDict):
    id: str


class OnCallApiListUsersResponse(typing.TypedDict):
    results: typing.List[OnCallApiUser]


class GrafanaAPIUser(typing.TypedDict):
    id: int


def _generate_unique_email() -> str:
    user = fake.profile()
    return f'{uuid.uuid4()}-{user["mail"]}'


async def _grafana_api_request(
    http_session: aiohttp.ClientSession, method: str, url: str, **request_kwargs
) -> typing.Awaitable[typing.Dict]:
    resp = await http_session.request(
        method, f"{GRAFANA_API_URL}{url}", **request_kwargs
    )
    return await resp.json()


async def _oncall_api_request(
    http_session: aiohttp.ClientSession, method: str, url: str, **request_kwargs
) -> typing.Awaitable[typing.Dict]:
    resp = await http_session.request(
        method,
        f"{ONCALL_API_URL}{url}",
        headers={"Authorization": ONCALL_API_TOKEN},
        **request_kwargs,
    )
    return await resp.json()


def generate_team(
    http_session: aiohttp.ClientSession, org_id: int
) -> typing.Callable[[], typing.Awaitable[typing.Dict]]:
    """
    https://grafana.com/docs/grafana/latest/developers/http_api/team/#add-team
    """

    def _generate_team() -> typing.Awaitable[typing.Dict]:
        return _grafana_api_request(
            http_session,
            "POST",
            "/api/teams",
            json={
                "name": str(uuid.uuid4()),
                "email": _generate_unique_email(),
                "orgId": org_id,
            },
        )

    return _generate_team


def generate_user(
    http_session: aiohttp.ClientSession, org_id: int
) -> typing.Callable[[], typing.Awaitable[typing.Dict]]:
    """
    https://grafana.com/docs/grafana/latest/developers/http_api/admin/#global-users
    """

    async def _generate_user() -> typing.Awaitable[typing.Dict]:
        user = fake.profile()

        # create the user in grafana
        grafana_user: GrafanaAPIUser = await _grafana_api_request(
            http_session,
            "POST",
            "/api/admin/users",
            json={
                "name": user["name"],
                "email": _generate_unique_email(),
                "login": str(uuid.uuid4()),
                "password": fake.password(length=20),
                "OrgId": org_id,
            },
        )

        # update the user's basic role in grafana to Admin
        # https://grafana.com/docs/grafana/latest/developers/http_api/org/#updates-the-given-user
        await _grafana_api_request(
            http_session,
            "PATCH",
            f'/api/org/users/{grafana_user["id"]}',
            json={"role": "Admin"},
        )

        return grafana_user

    return _generate_user


def generate_schedule(
    http_session: aiohttp.ClientSession, oncall_shift_ids: typing.List[str]
) -> typing.Callable[[], typing.Awaitable[typing.Dict]]:
    def _generate_schedule() -> typing.Awaitable[typing.Dict]:
        # Create a schedule
        # https://grafana.com/docs/oncall/latest/oncall-api-reference/schedules/#create-a-schedule
        return _oncall_api_request(
            http_session,
            "POST",
            "/api/v1/schedules",
            json={
                "name": f"Schedule {uuid.uuid4()}",
                "type": "calendar",
                "time_zone": "UTC",
                "shifts": oncall_shift_ids,
            },
        )

    return _generate_schedule


def _bulk_generate_data(
    iterations: int,
    data_generator_func: typing.Callable[[], typing.Awaitable[typing.Dict]],
) -> typing.Awaitable[typing.List[typing.Dict]]:
    return tqdm.gather(
        *[asyncio.ensure_future(data_generator_func()) for _ in range(iterations)]
    )


async def _generate_grafana_teams_and_users(
    args: argparse.Namespace, http_session: aiohttp.ClientSession
) -> typing.Awaitable[None]:
    global GRAFANA_API_URL
    GRAFANA_API_URL = args.grafana_api_url

    org_id = args.grafana_org_id

    print("Generating team(s)")
    await _bulk_generate_data(args.teams, generate_team(http_session, org_id))

    print("Generating user(s)")
    await _bulk_generate_data(args.users, generate_user(http_session, org_id))

    print(
        f"""
    Grafana teams and users generated
    Now head to the OnCall plugin and manually visit the plugin to trigger a sync. This will sync grafana
    teams/users to OnCall. Once completed, you can run the {SCHEDULES_ONCALL_SHIFTS_COMMAND} command.
    """
    )


async def _generate_oncall_schedules_and_oncall_shifts(
    args: argparse.Namespace, http_session: aiohttp.ClientSession
) -> typing.Awaitable[None]:
    global ONCALL_API_URL, ONCALL_API_TOKEN
    ONCALL_API_URL = args.oncall_api_url
    ONCALL_API_TOKEN = args.oncall_api_token

    today = datetime.now()

    print("Fetching users from OnCall API")

    # Fetch users from the OnCall API
    users: OnCallApiListUsersResponse = await _oncall_api_request(
        http_session, "GET", "/api/v1/users"
    )
    user_ids: typing.List[str] = [u["id"] for u in users["results"]]
    num_users = len(user_ids)

    print(f"Fetched {num_users} user(s) from the OnCall API")

    async def _create_oncall_shift(shift_start_time: str) -> typing.Awaitable[str]:
        """
        Creates an eight hour shift.

        `shift_start_time` - ex. 09:00:00, 15:00:00

        https://grafana.com/docs/oncall/latest/oncall-api-reference/on_call_shifts/#create-an-oncall-shift
        """
        new_shift: OnCallApiOnCallShift = await _oncall_api_request(
            http_session,
            "POST",
            "/api/v1/on_call_shifts",
            json={
                "name": f"On call shift{uuid.uuid4()}",
                "type": "rolling_users",
                "start": today.strftime(f"%Y-%m-%dT{shift_start_time}"),
                "time_zone": "UTC",
                "duration": 60 * 60 * 8,  # 8 hours
                "frequency": "daily",
                "week_start": "MO",
                "rolling_users": [
                    [u] for u in random.choices(user_ids, k=math.floor(num_users / 2))
                ],
                "start_rotation_from_user_index": 0,
                "team_id": None,
            },
        )

        oncall_shift_id = new_shift["id"]
        print(f"Generated OnCall shift w/ ID {oncall_shift_id}")
        return oncall_shift_id

    print("Creating three 8h on-call shifts")
    morning_shift_id = await _create_oncall_shift("00:00:00")
    afternoon_shift_id = await _create_oncall_shift("08:00:00")
    evening_shift_id = await _create_oncall_shift("16:00:00")

    print("Generating schedules(s)")
    await _bulk_generate_data(
        args.schedules,
        generate_schedule(
            http_session, [morning_shift_id, afternoon_shift_id, evening_shift_id]
        ),
    )


async def main() -> typing.Awaitable[None]:
    parser = argparse.ArgumentParser(
        description="Set of commands to help generate fake data in a Grafana OnCall setup."
    )
    subparsers = parser.add_subparsers(help="sub-command help")

    grafana_command_parser = subparsers.add_parser(
        TEAMS_USERS_COMMAND,
        description="Command to generate teams and users in Grafana",
    )
    grafana_command_parser.set_defaults(func=_generate_grafana_teams_and_users)
    grafana_command_parser.add_argument(
        "--grafana-api-url",
        help="Grafana API URL. This should include the basic authentication username/password in the URL. ex. http://oncall:oncall@localhost:3000",
        default="http://oncall:oncall@localhost:3000",
    )
    grafana_command_parser.add_argument(
        "--grafana-org-id",
        help="Org ID, in Grafana, of the org that you would like to generate data for",
        type=int,
        default=1,
    )
    grafana_command_parser.add_argument(
        "-t", "--teams", help="Number of teams to generate", default=10, type=int
    )
    grafana_command_parser.add_argument(
        "-u", "--users", help="Number of users to generate", default=1_000, type=int
    )

    oncall_command_parser = subparsers.add_parser(
        SCHEDULES_ONCALL_SHIFTS_COMMAND,
        description="Command to generate schedules and on-call shifts in OnCall",
    )
    oncall_command_parser.set_defaults(
        func=_generate_oncall_schedules_and_oncall_shifts
    )
    oncall_command_parser.add_argument(
        "--oncall-api-url",
        help="OnCall API URL",
        default="http://localhost:8080",
    )
    oncall_command_parser.add_argument(
        "--oncall-api-token", help="OnCall API token", required=True
    )
    oncall_command_parser.add_argument(
        "-s",
        "--schedules",
        help="Number of schedules to generate",
        default=100,
        type=int,
    )

    args = parser.parse_args()

    async with aiohttp.ClientSession(
        connector=aiohttp.TCPConnector(limit=5)
    ) as session:
        await args.func(args, session)


if __name__ == "__main__":
    asyncio.run(main())