otel-instrument

#!/usr/bin/env python3
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import logging
import os
import sys
from os import environ, system

sys.path.append(os.path.join(*(os.path.dirname(__file__), "python")))

from opentelemetry.environment_variables import (
    OTEL_PYTHON_DISABLED_INSTRUMENTATIONS,
)
from opentelemetry.instrumentation.aws_lambda import AwsLambdaInstrumentor
from opentelemetry.instrumentation.dependencies import (
    get_dist_dependency_conflicts,
)
from opentelemetry.instrumentation.distro import BaseDistro, DefaultDistro
from pkg_resources import iter_entry_points

# the path to the interpreter and all of the originally intended arguments
args = sys.argv[1:]

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# config default traces exporter if missing
environ.setdefault("OTEL_TRACES_EXPORTER", "otlp_proto_grpc_span")

# Set the service name if not already set
original_resource_attributes = environ.get(
    "OTEL_RESOURCE_ATTRIBUTES", ""
).split(",")
if "service.name=" not in original_resource_attributes:
    environ["OTEL_RESOURCE_ATTRIBUTES"] = ",".join(
        ["service.name=%s" % (environ.get("AWS_LAMBDA_FUNCTION_NAME"))]
        + original_resource_attributes
    )

# TODO: Remove this.
#
# TODO: Waiting on OTel Python support for setting Resource Detectors
# using environment variables. Auto Instrumentation (used by this Lambda
# Instrumentation) sets up the global TracerProvider which is the only
# time Resource Detectors can be configured.
#
# OR
#
# TODO: We forcibly modify the Resource Detectors on the TracerProvider after
# it is initialized? Not preferred.
#
# from opentelemetry.sdk.extension.aws.resource._lambda import (
#   AwsLambdaResourceDetector,
# )
# from opentelemetry.sdk.resources import get_aggregated_resources
# trace.get_tracer_provider.resource = get_aggregated_resources(
#        [
#            trace.get_tracer_provider.resource
#            AwsLambdaResourceDetector(),
#        ]
#    )
lambda_resource_attributes = (
    "cloud.region=%s,cloud.provider=aws,faas.name=%s,faas.version=%s"
    % (
        environ.get("AWS_REGION"),
        environ.get("AWS_LAMBDA_FUNCTION_NAME"),
        environ.get("AWS_LAMBDA_FUNCTION_VERSION"),
    )
)
environ["OTEL_RESOURCE_ATTRIBUTES"] = "%s,%s" % (
    lambda_resource_attributes,
    environ.get("OTEL_RESOURCE_ATTRIBUTES"),
)


def _load_distros() -> BaseDistro:
    for entry_point in iter_entry_points("opentelemetry_distro"):
        try:
            distro = entry_point.load()()
            if not isinstance(distro, BaseDistro):
                logger.debug(
                    "%s is not an OpenTelemetry Distro. Skipping",
                    entry_point.name,
                )
                continue
            logger.debug("Distribution %s will be configured", entry_point.name)
            return distro
        except Exception as exc:  # pylint: disable=broad-except
            logger.debug(
                "Distribution %s configuration failed", entry_point.name
            )
    return DefaultDistro()


def _load_instrumentors(distro):
    package_to_exclude = os.environ.get(
        OTEL_PYTHON_DISABLED_INSTRUMENTATIONS, []
    )
    if isinstance(package_to_exclude, str):
        package_to_exclude = package_to_exclude.split(",")
        # to handle users entering "requests , flask" or "requests, flask" with spaces
        package_to_exclude = [x.strip() for x in package_to_exclude]

    for entry_point in iter_entry_points("opentelemetry_instrumentor"):
        if entry_point.name in package_to_exclude:
            logger.debug(
                "Instrumentation skipped for library %s", entry_point.name
            )
            continue

        try:
            conflict = get_dist_dependency_conflicts(entry_point.dist)
            if conflict:
                logger.debug(
                    "Skipping instrumentation %s: %s",
                    entry_point.name,
                    conflict,
                )
                continue

            # tell instrumentation to not run dep checks again as we already did it above
            distro.load_instrumentor(entry_point, skip_dep_check=True)
            logger.info("Instrumented %s", entry_point.name)
        except Exception as exc:  # pylint: disable=broad-except
            logger.debug("Instrumenting of %s failed", entry_point.name)


def _load_configurators():
    configured = None
    for entry_point in iter_entry_points("opentelemetry_configurator"):
        if configured is not None:
            logger.warning(
                "Configuration of %s not loaded, %s already loaded",
                entry_point.name,
                configured,
            )
            continue
        try:
            entry_point.load()().configure()  # type: ignore
            configured = entry_point.name
        except Exception as exc:  # pylint: disable=broad-except
            logger.debug("Configuration of %s failed", entry_point.name)


def modify_module_name(module_name):
    """Returns a valid modified module to get imported"""
    return ".".join(module_name.split("/"))


class HandlerError(Exception):
    pass


distro = _load_distros()
distro.configure()
_load_configurators()
_load_instrumentors(distro)
# TODO: move to python-contrib
AwsLambdaInstrumentor().instrument(skip_dep_check=True)

# start the runtime with the extra options
system(" ".join(args))