Sentry Automated Project Setup - Includes Alerting for Slack and Pagerduty

sentry_setup.py

"""
Inspiration: https://gist.github.com/nikolaik/85e19b89223686b9ab560822fb63bc01

Quickstart:
    1. Setup a virtualenv
    2. pip install requests pydantic rich click python-dotenv
    3. python sentry_setup_v2.py

"""
import json
import os
import pprint
import sys
from collections import namedtuple
from typing import Optional, Union, Any

import click
import requests
from dotenv import load_dotenv
from pydantic import BaseModel, RootModel, Field
from requests.auth import AuthBase
from rich.console import Console
from rich.json import JSON

load_dotenv()

RateLimitConfig = namedtuple("RateLimitConfig", ["window_in_minutes", "count"])

SENTRY_HOST = "https://sentry.io"
SENTRY_ORG_SLUG = '<your org>'
SENTRY_AUTH_TOKEN = os.environ.get("SENTRY_AUTH_TOKEN")

SENTRY_DEFAULT_CLIENT_KEY_NAME = "Default"
SENTRY_PROD_CLIENT_KEY_NAME = "Production"
SENTY_PROD_RATE_LIMIT_CONFIG = RateLimitConfig(900, 10)  # comes to around 30k events per month
SENTRY_NONPROD_CLIENT_KEY_NAME = "NonProduction"
SENTY_NONPROD_RATE_LIMIT_CONFIG = RateLimitConfig(900, 3)  # comes to around 10k events per month

SLACK_ACCOUNT_NAME = "<slack integration account name>"
PAGERDUTY_ACCOUNT_NAME = "<pagerduty integration account name>"

if not SENTRY_AUTH_TOKEN:
    raise sys.exit("Sentry auth token is not in env")

console = Console(color_system="standard", soft_wrap=True)
grey_print_models = lambda x: console.print(type(x).__name__, JSON(x.model_dump_json(), highlight=False), style="white")


class Project(BaseModel):
    id: str
    name: str
    slug: str


class Team(BaseModel):
    id: str
    name: str
    slug: str


class RateLimit(BaseModel):
    window: int
    count: int


class ClientKeyDSN(BaseModel):
    public: str


class ClientKey(BaseModel):
    id: str
    name: str
    rate_limit: RateLimit | None = Field(alias="rateLimit")
    dsn: ClientKeyDSN | None


class AlertCondition(BaseModel):
    id: str


class AlertConfiguration(BaseModel):
    id: str


class SlackAlertAction(BaseModel):
    id: str
    channel: str
    tags: str
    workspace: int


class PagerDutyAlertAction(BaseModel):
    id: str
    account: int
    service: int


class Alert(BaseModel):
    id: str | None
    name: str
    conditions: list[AlertCondition] | None
    filters: list[dict]
    actions: list[Union[SlackAlertAction, PagerDutyAlertAction]] | None
    environment: str | None
    frequency: int | None
    action_match: str = Field(alias="actionMatch")


class IterableRootModel(RootModel):
    root: list[Any]

    def __iter__(self):
        return iter(self.root)

    def __getitem__(self, item):
        return self.root[item]

    def __add__(self, other):
        return self.__class__(self.root + other.root)

    def by_name(self, name: str):
        for p in self.root:
            if p.name == name:
                return p
        return None


class Projects(IterableRootModel):
    root: list[Project]


class Teams(IterableRootModel):
    root: list[Team]


class ClientKeys(IterableRootModel):
    root: list[ClientKey]


class Alerts(IterableRootModel):
    root: list[Alert]


class SetupException(Exception):

    def __init__(self, message):
        self.message = message
        super().__init__()


class SentryClient:
    ORG_URL = f"{SENTRY_HOST}/api/0/organizations/{SENTRY_ORG_SLUG}"
    TEAMS_URL = f"{SENTRY_HOST}/api/0/teams/{SENTRY_ORG_SLUG}"
    PROJECTS_URL = f"{SENTRY_HOST}/api/0/projects/{SENTRY_ORG_SLUG}"

    class BearerAuth(AuthBase):
        def __init__(self, token):
            self.token = token

        def __call__(self, r):
            r.headers["authorization"] = "Bearer " + self.token
            return r

    def __init__(self):
        self._client = None

    @property
    def client(self):
        if self._client is None:
            s = requests.Session()
            s.auth = self.BearerAuth(SENTRY_AUTH_TOKEN)
            self._client = s
        return self._client

    def list_teams(self, params=None):
        """
        https://docs.sentry.io/api/teams/list-an-organizations-teams/
        """
        url = self.ORG_URL + "/teams/"
        resp = self.client.get(url, params=params)
        teams = Teams.model_validate_json(resp.content)
        has_next_results = resp.links.get("next", {}).get("results", False)
        if has_next_results and has_next_results != "false":
            params = {"cursor": resp.links["next"]["cursor"]}
            return teams + self.list_teams(params)
        return teams

    def list_projects(self, params=None):
        """
        https://docs.sentry.io/api/organizations/list-an-organizations-projects
        """
        url = self.ORG_URL + "/projects/"
        resp = self.client.get(url, params=params)
        projects = Projects.model_validate_json(resp.content)
        has_next_results = resp.links.get("next", {}).get("results", False)
        if has_next_results and has_next_results != "false":
            params = {"cursor": resp.links["next"]["cursor"]}
            return projects + self.list_projects(params)
        return projects

    def create_team(self, team_name: str) -> Team:
        url = self.ORG_URL + "/teams/"
        params = {
            "name": team_name
        }
        resp = self.client.post(url, data=params)
        if resp.status_code != 201:
            raise SetupException(f"client unable to create team: {resp.text}")
        team = Team.model_validate_json(resp.content)
        return team

    def list_client_keys(self, project: Project, params=None) -> ClientKeys:
        url = self.PROJECTS_URL + f"/{project.slug}/keys/"
        resp = self.client.get(url, params=params)
        if resp.status_code != 200:
            raise SetupException(f"client unable to list client keys: {resp.text}")
        keys = ClientKeys.model_validate_json(resp.content)
        has_next_results = resp.links.get("next", {}).get("results", False)
        if has_next_results and has_next_results != "false":
            params = {"cursor": resp.links["next"]["cursor"]}
            return keys + self.list_client_keys(project, params)
        return keys

    def create_project(self, project_name: str, platform: Optional[str], team: Team) -> Project:
        url = self.TEAMS_URL + f"/{team.slug}/projects/"
        params = {
            "name": project_name,
            "default_rules": False,
        }
        if platform:
            params["platform"] = platform

        resp = self.client.post(url, data=params)
        if resp.status_code != 201:
            raise SetupException(f"client unable to create project: {resp.text}")
        project = Project.model_validate_json(resp.content)
        return project

    def create_client_key(self, project: Project, key: str) -> ClientKey:
        print(f"creating client key {key}...")

        url = self.PROJECTS_URL + f"/{project.slug}/keys/"
        params = {
            "name": key,
            # cant give rateLimit here while creating itself
        }
        resp = self.client.post(url, data=params)
        if resp.status_code != 201:
            raise SetupException(f"client unable to create client key: {resp.text}")
        key = ClientKey.model_validate_json(resp.content)
        return key

    def update_client_key(self, project: Project, key: ClientKey) -> ClientKey:
        url = self.PROJECTS_URL + f"/{project.slug}/keys/{key.id}/"
        params = {
            "rateLimit": {
                "window": key.rate_limit.window,
                "count": key.rate_limit.count,
            }
        }
        resp = self.client.put(url, json=params)
        if resp.status_code != 200:
            raise SetupException(f"client unable to update client key: {resp.text}")
        key = ClientKey.model_validate_json(resp.content)
        return key

    def delete_client_key(self, project: Project, client_key: ClientKey):
        print(f"deleting client key: {client_key.name}...")
        print(json.dumps(client_key.model_dump()))

        url = self.PROJECTS_URL + f"/{project.slug}/keys/{client_key.id}/"
        resp = self.client.delete(url)
        if resp.status_code != 204:
            raise SetupException(f"client unable to delete client key: {resp.text}")

    def get_alerts_configuration(self, project: Project) -> dict:
        # todo: paginate ?  maybe not since links is not present in header
        url = self.PROJECTS_URL + f"/{project.slug}/rules/configuration/"
        resp = self.client.get(url)
        if resp.status_code != 200:
            raise SetupException(f"client unable to list alerts: {resp.text}")
        return resp.json()

    def list_alerts(self, project: Project, params=None) -> Alerts:
        url = self.PROJECTS_URL + f"/{project.slug}/rules/"
        resp = self.client.get(url, params=params)
        if resp.status_code != 200:
            raise SetupException(f"client unable to list alerts: {resp.text}")
        alerts = Alerts.model_validate_json(resp.content)

        has_next_results = resp.links.get("next", {}).get("results", False)
        if has_next_results and has_next_results != "false":
            params = {"cursor": resp.links["next"]["cursor"]}
            return alerts + self.list_alerts(project, params)

        return alerts

    def create_alert(self, project: Project, alert: Alert) -> Alert:
        """
        https://forum.sentry.io/t/feature-request-manage-projects-alerts-via-api/5667
        https://forum.sentry.io/t/duplicate-alert-rules-across-projects/3416/8
        """
        print(f"creating alert {alert.name}...")
        pprint.pprint(alert)

        url = self.PROJECTS_URL + f"/{project.slug}/rules/"
        data = alert.model_dump(by_alias=True)
        resp = self.client.post(url, json=data)
        if resp.status_code != 200:
            raise SetupException(f"client unable to create alert: {resp.text}")
        created_alert = Alert.model_validate_json(resp.content)
        return created_alert


def default_new_issue_alert(
        project: Project,
        alert_config: dict,
        alert_slack_channel_name: str,
        alert_pd_service_name: str
) -> Optional[Alert]:
    if not (alert_slack_channel_name or alert_pd_service_name):
        return None

    condition_new_issue = AlertCondition(
        id="sentry.rules.conditions.first_seen_event.FirstSeenEventCondition",
    )

    actions = []
    if alert_slack_channel_name:
        slack_alert_id = "sentry.integrations.slack.notify_action.SlackNotifyServiceAction"
        slack_workspace_id = None

        for cfg in alert_config["actions"]:
            if cfg.get("id") == slack_alert_id:

                integration_choices = cfg.get("formFields", {}).get("workspace", {}).get("choices", [])
                for ch in integration_choices:
                    if ch[1] == SLACK_ACCOUNT_NAME:
                        slack_workspace_id = ch[0]

        if not slack_workspace_id:
            raise SetupException(f"Unable to find slack workspace for {SLACK_ACCOUNT_NAME}")

        slack_alert = SlackAlertAction(
            id=slack_alert_id,
            channel=alert_slack_channel_name,
            tags="environment,url",
            workspace=slack_workspace_id
        )
        actions.append(slack_alert)

    if alert_pd_service_name:
        pd_alert_id = "sentry.integrations.pagerduty.notify_action.PagerDutyNotifyServiceAction"
        pd_account_id = None
        pd_service_id = None
        for cfg in alert_config["actions"]:
            if cfg.get("id") == pd_alert_id:

                account_choices = cfg.get("formFields", {}).get("account", {}).get("choices")
                for ch in account_choices:
                    if ch[1] == PAGERDUTY_ACCOUNT_NAME:
                        pd_account_id = ch[0]

                service_choices = cfg.get("formFields", {}).get("service", {}).get("choices")
                for ch in service_choices:
                    if ch[1] == alert_pd_service_name:
                        pd_service_id = ch[0]

        if not pd_account_id:
            raise SetupException(
                f"Unable to find pagerduty account for {PAGERDUTY_ACCOUNT_NAME}")

        if not pd_service_id:
            raise SetupException(
                f"Unable to find pagerduty service for {alert_pd_service_name}")

        pd_alert = PagerDutyAlertAction(
            id="sentry.integrations.pagerduty.notify_action.PagerDutyNotifyServiceAction",
            account=pd_account_id,
            service=pd_service_id,
        )
        actions.append(pd_alert)

    ps = "".join([x.capitalize() for x in project.slug.split("-")])
    alert_name = f"{ps}-NewIssueProductionAlert"
    return Alert(
        id=None,
        name=alert_name,
        filters=[],
        conditions=[condition_new_issue],
        actions=actions,
        frequency=1440,
        environment="production",
        actionMatch="any"
    )


def _setup(project_name, team_name, project_platform, alert_slack_channel_name, alert_pd_service_name):
    if alert_slack_channel_name and not alert_slack_channel_name.startswith("#"):
        raise SetupException("slack channel name should start with #")

    print("Initializing...\n")
    sc = SentryClient()

    projects = sc.list_projects()
    teams = sc.list_teams()

    project = projects.by_name(project_name)
    team = teams.by_name(team_name)

    if project and not team:
        raise SetupException("project already exists and may be part of another team")

    if not team:
        print(f"creating {team_name}...")
        team = sc.create_team(team_name)
        print(team)
    else:
        print(f"{team_name} team exists. Not creating")

    if not project:
        print(f"creating {project_name}...")
        project = sc.create_project(project_name, project_platform, team)
        print(project)
    else:
        print(f"{team_name} project exists. Not creating")

    # create client keys
    client_keys = sc.list_client_keys(project)

    default_client_key = client_keys.by_name(SENTRY_DEFAULT_CLIENT_KEY_NAME)
    if default_client_key:
        sc.delete_client_key(project, default_client_key)

    production_client_key = client_keys.by_name(SENTRY_PROD_CLIENT_KEY_NAME)
    if not production_client_key:
        production_client_key = sc.create_client_key(project, SENTRY_PROD_CLIENT_KEY_NAME)
        print(f"ratelimiting {production_client_key.name}...")
        production_client_key.rate_limit = RateLimit(
            window=SENTY_PROD_RATE_LIMIT_CONFIG.window_in_minutes,
            count=SENTY_PROD_RATE_LIMIT_CONFIG.count
        )
        production_client_key = sc.update_client_key(project, production_client_key)
    else:
        print(f"{production_client_key.name} client key already exists. Not creating")

    non_production_client_key = client_keys.by_name(SENTRY_NONPROD_CLIENT_KEY_NAME)
    if not non_production_client_key:
        non_production_client_key = sc.create_client_key(project, SENTRY_NONPROD_CLIENT_KEY_NAME)
        print(f"ratelimiting {non_production_client_key.name}...")
        non_production_client_key.rate_limit = RateLimit(
            window=SENTY_NONPROD_RATE_LIMIT_CONFIG.window_in_minutes,
            count=SENTY_NONPROD_RATE_LIMIT_CONFIG.count
        )
        non_production_client_key = sc.update_client_key(project, non_production_client_key)
    else:
        print(f"{non_production_client_key.name} client key already exists. Not creating.")

    # create alerts
    alerts = sc.list_alerts(project)
    alerts_config = sc.get_alerts_configuration(project)

    try:
        da = default_new_issue_alert(project, alerts_config, alert_slack_channel_name, alert_pd_service_name)
    except Exception as e:
        # print alerts config and raise incase of any exception
        print(alerts_config)
        raise e

    if da:
        if not alerts.by_name(da.name):
            da = sc.create_alert(project, da)
        else:
            print(f"{da.name} already exists. Not creating")
    else:
        print(f"Not creating any alerts since slack channel or pd name is given")

    print("\n", "-" * 20, "\n")

    grey_print_models(project)
    grey_print_models(production_client_key)
    grey_print_models(non_production_client_key)
    if da:
        grey_print_models(da)

    print("\n", "-" * 20, "\n")
    console.print(project.name, style="bold green")
    console.print(f"{production_client_key.name} Client Key: {production_client_key.dsn.public}", style="bold red")
    console.print(f"{non_production_client_key.name} Client Key: {non_production_client_key.dsn.public}",
                  style="bold blue")


@click.command()
@click.option('--project_name', prompt='Project Name (prefer small case hyphenated)', required=True, type=str)
@click.option('--team_name', prompt='Team Name (prefer small case hyphenated)', required=True, type=str)
@click.option('--project_platform', prompt='Platform (go, FastAPI). (optional)', default="")
@click.option('--alert_slack_channel_name', prompt='Alert Slack Channel Name (with #). (optional)', default="")
@click.option('--alert_pd_service_name', prompt='Alert PD Service Name. (optional)', default="")
@click.confirmation_option(prompt='Are you sure you want to proceed setting up project in sentry with above?')
def setup(project_name, team_name, project_platform, alert_slack_channel_name, alert_pd_service_name):
    try:
        return _setup(project_name, team_name, project_platform, alert_slack_channel_name, alert_pd_service_name)
    except SetupException as e:
        console.print(e.message, style="bold red")


# def test_setup_prompt():
#     runner = CliRunner()
#     result = runner.invoke(setup, input="\n".join([
#         'test-project-1',
#         "test--team-1",
#         "go",
#         "#test",
#         "test",
#         "y"
#     ]))
#     assert not result.exception
#
#
# def test_setup():
#     runner = CliRunner()
#     result = runner.invoke(
#         setup,
#         [
#             '--project_name', 'test-project-1',
#             '--team_name', "test-team-1"
#         ]
#     )
#     assert not result.exception


if __name__ == "__main__":
    console.print("Ctrl + c to quit", style="dim blue")
    setup()