Add helper methods

openrelik_common/telemetry.py

@@ -1,11 +1,14 @@
 # OpenRelik OpenTelemetry helper methods
 
+import json
 import os
 
 from opentelemetry import trace
 
 from opentelemetry.exporter.otlp.proto.grpc import trace_exporter as grpc_exporter
 from opentelemetry.exporter.otlp.proto.http import trace_exporter as http_exporter
+from opentelemetry.instrumentation.celery import CeleryInstrumentor
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
 from opentelemetry.sdk.resources import Resource
 from opentelemetry.sdk.trace import TracerProvider
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
@@ -39,3 +42,34 @@ def setup_telemetry(service_name: str):
         raise Exception("Unsupported OTEL tracing mode %s", otel_mode)
 
     trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(trace_exporter))
+
+
+def instrument_celery_app(celery_app):
+    """Helper method to call the OpenTelemetry Python instrumentor on an Celery app object.
+
+    Args:
+        celery_app (celery.app.Celery): the celery app to instrument.
+    """
+    CeleryInstrumentor().instrument(celery_app=celery_app)
+
+
+def instrument_fast_api(fast_api):
+    """Helper method to call the OpenTelemetry Python instrumentor on an FastAPI app object.
+
+    Args:
+        fast_api (fastapi.FastAPI): the FastAPI app to instrument.
+    """
+    FastAPIInstrumentor.instrument_app(api_v1)
+
+
+def add_attribute_to_current_span(name: str, value: object):
+    """This methods tried to get a handle of the OpenTelemetry span in the current context, and add
+    an attribute to it, using the name and value passed as arguments.
+
+    Args:
+        name (str): the name for the attribute.
+        value (object): the value of the attribute. This needs to be a json serializable object.
+    """
+    otel_span = trace.get_current_span()
+    if otel_span != trace.span.INVALID_SPAN:
+        otel_span.set_attribute(name, json.dumps(value))