app.py

import json
import math
import os
import random

import pika
from flask import Flask, jsonify
from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.flask import FlaskInstrumentor
from opentelemetry.instrumentation.psycopg2 import Psycopg2Instrumentor
from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
from opentelemetry.propagate import inject
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from sqlalchemy import create_engine, text


# OpenTelemetry
resource = Resource.create()
provider = TracerProvider(resource=resource)
otlp_exporter = OTLPSpanExporter()
processor = BatchSpanProcessor(otlp_exporter)
provider.add_span_processor(processor)
trace.set_tracer_provider(provider)
tracer = trace.get_tracer(__name__)

# Database
DB_URL = os.environ["DATABASE_URL"]
engine = create_engine(DB_URL)
SQLAlchemyInstrumentor().instrument(engine=engine)
Psycopg2Instrumentor().instrument()

# RabbitMQ
RABBITMQ_HOST = "rabbitmq"
QUEUE_NAME = "my_queue"
connection = pika.BlockingConnection(
    pika.ConnectionParameters(RABBITMQ_HOST, heartbeat=3600)
)
channel = connection.channel()
channel.queue_declare(queue=QUEUE_NAME, durable=True)

# Flask
app = Flask(__name__)
FlaskInstrumentor().instrument_app(app)


@app.route('/db-check')
def db_check():
    with tracer.start_as_current_span("db-query-execution"):
        with engine.connect() as connection:
            sleep_time = random.random() / 50.0
            result = connection.execute(
                text("SELECT pg_sleep(:t);"), {"t": sleep_time}
            ).scalar()

    return jsonify({
        "status": "success",
        "result": result,
    })


def is_prime(n: int):
    if n < 2:
        return False
    for i in range(2, int(math.sqrt(n)) + 1):
        if n % i == 0:
            return False
    return True


@app.route('/cpu-heavy')
def cpu_heavy():
    limit = 20000
    prime_count = 0
    with tracer.start_as_current_span("cpu-intensive-task") as span:
        for number in range(2, limit):
            if is_prime(number):
                prime_count += 1

        span.set_attribute("prime_count", prime_count)

    return jsonify({
        "limit": limit,
        "prime_count": prime_count,
    })


@app.route('/send-message')
def send_message():
    with tracer.start_as_current_span(
        "publish_message", kind=trace.SpanKind.PRODUCER
    ):
        headers = {}
        inject(headers)

        message = {"sleep": random.random() / 50.0}
        properties = pika.BasicProperties(
            headers=headers,
            delivery_mode=pika.DeliveryMode.Persistent,
        )
        channel.basic_publish(
            exchange="",
            routing_key=QUEUE_NAME,
            body=json.dumps(message),
            properties=properties,
        )

    return jsonify({"status": "OK"})