main.py

import aiohttp
import discord
import json
import logging
import os
import sentry_sdk
from aiohttp import web
from sentry_sdk.integrations.asyncio import AsyncioIntegration
from sentry_sdk.integrations.logging import LoggingIntegration
from sentry_sdk.crons import capture_checkin
from sentry_sdk.crons.consts import MonitorStatus
from time import perf_counter, sleep, time


# BUILD_INFO is generated by the build pipeline (e.g. docker/metadata-action).
# It looks like:
# {"tags":["ghcr.io/watonomous/repo-ingestion:main"],"labels":{"org.opencontainers.image.title":"repo-ingestion","org.opencontainers.image.description":"Simple server to receive file changes and open GitHub pull requests","org.opencontainers.image.url":"https://github.com/WATonomous/repo-ingestion","org.opencontainers.image.source":"https://github.com/WATonomous/repo-ingestion","org.opencontainers.image.version":"main","org.opencontainers.image.created":"2024-01-20T16:10:39.421Z","org.opencontainers.image.revision":"1d55b62b15c78251e0560af9e97927591e260a98","org.opencontainers.image.licenses":""}}
BUILD_INFO=json.loads(os.getenv("DOCKER_METADATA_OUTPUT_JSON", "{}"))
IS_SENTRY_ENABLED = os.getenv("SENTRY_DSN") is not None

# Set up Sentry
if IS_SENTRY_ENABLED:
    build_labels = BUILD_INFO.get("labels", {})
    image_title = build_labels.get("org.opencontainers.image.title", "unknown_image")
    image_version = build_labels.get("org.opencontainers.image.version", "unknown_version")
    image_rev = build_labels.get("org.opencontainers.image.revision", "unknown_rev")

    sentry_config = {
        "dsn": os.environ["SENTRY_DSN"],
        "environment": os.getenv("DEPLOYMENT_ENVIRONMENT", "unknown"),
        "release": os.getenv("SENTRY_RELEASE", f'{image_title}:{image_version}@{image_rev}'),
    }

    print(f"Sentry SDK version: {sentry_sdk.VERSION}")
    print(f"Sentry DSN found. Setting up Sentry with config: {sentry_config}")

    sentry_logging = LoggingIntegration(
        level=logging.INFO,        # Capture info and above as breadcrumbs
        event_level=logging.ERROR  # Send errors as events
    )

    def sentry_traces_sampler(sampling_context):
        # Inherit parent sampling decision
        if sampling_context["parent_sampled"] is not None:
            return sampling_context["parent_sampled"]

        # Don't need to sample health checks
        aiohttp_request = sampling_context.get("aiohttp_request")
        if aiohttp_request is not None and aiohttp_request.path == "/health":
            return 0
        
        # Sample everything else
        return 1

    sentry_sdk.init(
        **sentry_config,
        integrations=[sentry_logging, AsyncioIntegration()],

        # Set traces_sample_rate to 1.0 to capture 100%
        # of transactions for performance monitoring.
        # We recommend adjusting this value in production,
        # traces_sample_rate=1.0,
        traces_sampler=sentry_traces_sampler,

        enable_tracing=True,
        keep_alive=True,
    )
else:
    print("No Sentry DSN found. Skipping Sentry setup.")

logger = logging.getLogger('discord.wato-provisioner')

intents = discord.Intents.none()
intents.members = True

client = discord.Client(intents=intents)
state = {
    "sentry_cron_last_ping_time": 0,
}

@client.event
async def on_member_join(member):
    logger.info(f'{member} has joined the server.')

    # Trigger the provisioner
    logger.info("Triggering the provisioner")
    async with aiohttp.ClientSession() as session:
        async with session.post(
            'https://api.github.com/repos/WATonomous/infra-config/dispatches',
            headers={
                'Content-Type': 'application/json',
                'Authorization': f'Bearer {os.environ["GITHUB_TOKEN"]}',
                'Accept': 'application/vnd.github+json',
                'X-GitHub-Api-Version': "2022-11-28",
            },
            json={'event_type': 'provision-discord'},
        ) as resp:
            if resp.status == 204:
                logger.info('Successfully triggered the Discord provisioner.')
            else:
                logger.error(f'Failed to trigger the Discord provisioner. {resp.status=}')

@client.event
async def on_ready():
    logger.info(f'{client.user} has connected to Discord. Now setting up healthchecks')

    health_endpoint_app = web.Application()
    health_endpoint_app.add_routes([web.get('/health', health_endpoint)])
    health_endpoint_runner = web.AppRunner(health_endpoint_app)
    await health_endpoint_runner.setup()
    healthcheck_site = web.TCPSite(health_endpoint_runner, '0.0.0.0', 8000)
    await healthcheck_site.start()

    logger.info("ready")

async def health_endpoint(_request):
    if client.is_closed():
        success = False
    else:
        success = True
        
    current_time = time()
    # Ping Sentry every 2 minutes
    if IS_SENTRY_ENABLED and current_time - state["sentry_cron_last_ping_time"] >= 120:
        state["sentry_cron_last_ping_time"] = current_time
        capture_checkin(
            monitor_slug='discord-provisioner-bot',
            status=MonitorStatus.OK if success else MonitorStatus.ERROR,
            monitor_config={
                "schedule": { "type": "interval", "value": 1, "unit": "minute" },
                "checkin_margin": 5, # minutes
                "max_runtime": 1, # minutes
                "failure_issue_threshold": 1,
                "recovery_threshold": 2,
            }
        )
        logger.info(f"Pinged Sentry CRON with status {'OK' if success else 'ERROR'}")

    if success:
        return web.Response(text='OK')
    else:
        return web.Response(text='Client is closed!', status=500)

client.run(os.environ['DISCORD_TOKEN'])