util/trace.py

import os
from google.cloud import trace_v2
from google.oauth2 import service_account
from google.protobuf.timestamp_pb2 import Timestamp
from datetime import datetime
import uuid

class TraceManager:
    def __init__(self):
        self.client = self._initialize_client()
        self.project_name = f"projects/{os.getenv('PROJECT')}"

    def _initialize_client(self):
        with open(os.getenv("GOOGLE_APPLICATION_CREDENTIALS"), "w") as f:
            f.write(os.getenv("gcp_service_json"))
            f.close()
        credentials = service_account.Credentials.from_service_account_file(os.getenv("GOOGLE_APPLICATION_CREDENTIALS"))
        return trace_v2.TraceServiceClient(credentials=credentials)

    def start_span(self, span_display_name="Default Span"):
        span_id = uuid.uuid4().hex[:16]
        trace_id = uuid.uuid4().hex
        span_name = f"{self.project_name}/traces/{trace_id}/spans/{span_id}"
        start_timestamp = self._get_current_timestamp()
        return trace_id, span_id, span_name, start_timestamp

    def close_span(self, trace_id, span_id, start_timestamp, span_display_name="Default Span", status="OK", error_message=None):
        span_name = f"{self.project_name}/traces/{trace_id}/spans/{span_id}"
        end_timestamp = self._get_current_timestamp()
        attributes = self._create_attributes(status, error_message)
        
        span = trace_v2.Span(
            name=span_name,
            span_id=span_id,
            display_name=trace_v2.types.TruncatableString(value=span_display_name),
            start_time=start_timestamp,
            end_time=end_timestamp,
            attributes=attributes
        )
        self.client.batch_write_spans(name=self.project_name, spans=[span])
        return span_name

    def _get_current_timestamp(self):
        now = datetime.utcnow()
        timestamp = Timestamp()
        timestamp.FromDatetime(now)
        return timestamp

    def _create_attributes(self, status, error_message):
        attribute_map = {
            "status": trace_v2.types.AttributeValue(string_value=trace_v2.types.TruncatableString(value=status))
        }
        if error_message:
            attribute_map["error.message"] = trace_v2.types.AttributeValue(string_value=trace_v2.types.TruncatableString(value=error_message))
        return trace_v2.types.Attributes(attribute_map=attribute_map)