Source code for visualization_toolkit.helpers.dash.monitoring

import logging
import os

from dash import Dash, Input, Output
from flask import Flask, jsonify
from urllib3 import Retry

import requests
from envelop import Environment

from visualization_toolkit.constants import (
    APP_ENVIRONMENT,
    JINJA_ENVIRONMENT,
    LIBRARY_PATH,
    DATADOG_CLIENT_TOKEN,
    DATADOG_APPLICATION_ID,
    DATADOG_SERVICE_NAME,
    DATADOG_ENVIRONMENT,
    ATLAS_LOCATION_ID,
    ATLAS_USER_STORE_ID,
    ATLAS_ENABLE_DATADOG,
    ATLAS_DEBUG,
)

logger = logging.getLogger(__name__)




def setup_sentry(env: Environment = None, env_var: str = "SENTRY_DSN"):
    """
    Intialize sentry tracking of exceptions. Will connect to sentry via the ``SENTRY_DSN`` variable set in the app environment.

    :param env: Enveleop Environment client to use. Defaults to the libraries ``APP_ENVIRONMENT``. Generally not needed to be set.
    :param env_var: Environment variable name that stores the sentry dsn. Defaults to ``SENTRY_DSN``.
    :return:
    """
    env = env or APP_ENVIRONMENT
    SENTRY_DSN = env.get(env_var, None)
    if SENTRY_DSN:
        import sentry_sdk

        sentry_sdk.init(dsn=SENTRY_DSN)





def setup_datadog():
    """
    Initialize datadog APM tracing of the dash app.
    :return:
    """
    if ATLAS_ENABLE_DATADOG:
        from ddtrace import tracer
        from ddtrace import patch_all

        DD_PRIVATE_IP = get_aws_ip()
        tracer.configure(hostname=DD_PRIVATE_IP)

        patch_all()



def setup_datadog_rum(
    app: Dash,
    user_store_id: str = ATLAS_USER_STORE_ID,
    location_id: str = ATLAS_LOCATION_ID,
    debug: bool = ATLAS_DEBUG,
):
    """
    Initialize datadog RUM tracking of the dash app on the browser. Requires adding ``user_store`` and ``location_store`` to the dash app.

    :param app: Dash app to enable RUm on
    :param user_store_id: ID of user store that user information will exist in. Defaults to ``ATLAS_USER_STORE_ID``.
    :param location_id: Environment variable name that stores the sentry dsn. Defaults to ``USE_DATADOG``.
    :param debug: If True, additional logging will be included in the client side callbacks to debug in the browser.
    :return:
    """
    if ATLAS_ENABLE_DATADOG:
        # initialize DD rum client via the browser
        with open(
            os.path.join(LIBRARY_PATH, "templates/datadog/initializeDDRum.js")
        ) as f:
            init_script = JINJA_ENVIRONMENT.from_string(f.read()).render(
                parameters={
                    "debug": debug,
                    "dd_client_token": DATADOG_CLIENT_TOKEN,
                    "dd_application_id": DATADOG_APPLICATION_ID,
                    "dd_service_name": DATADOG_SERVICE_NAME,
                    "dd_environment": DATADOG_ENVIRONMENT,
                }
            )

        app.clientside_callback(
            init_script,
            Input(location_id, "href"),
        )

        # Then register a callback to register the user once DD rum is initialized and user is authenticated
        with open(os.path.join(LIBRARY_PATH, "templates/datadog/ddRumSetUser.js")) as f:
            user_script = JINJA_ENVIRONMENT.from_string(f.read()).render(
                parameters={
                    "debug": debug,
                }
            )

            app.clientside_callback(
                user_script,
                Input(user_store_id, "data"),
            )


def get_aws_ip():
    if os.environ.get("AWS_EXECUTION_ENV") == "AWS_ECS_FARGATE":
        return "127.0.0.1"  # for fargate
    # implement a retry strategy in case the EC2 Instance MetaData Server v1 endpoint isn't available yet
    retry_strategy = Retry(total=3, backoff_factor=0.5)
    session = requests.Session()
    session.mount("http://", requests.adapters.HTTPAdapter(max_retries=retry_strategy))
    r = session.get("http://169.254.169.254/latest/meta-data/local-ipv4", timeout=15)
    r.raise_for_status()
    return r.text




def setup_healthcheck(server: Flask | Dash):
    """
    Adds a custom healthcheck endpoint on the dash app required for deployments to AWS.
    Healthcheck endpoint will return a 200 and is found on ``/status/healthcheck/``

    :param server: Flask server. Generally this is accessed via the ``server`` attribute on an initialized dash app. Dash app can also be passed in and the server is retrieved inside the function.
    :return:
    """
    if isinstance(server, Dash):
        server = server.server

    @server.route("/status/healthcheck/")
    def healthcheck():
        response = {"status": "ok"}
        return jsonify(response)