File size: 2,377 Bytes
3d35409 942388a 3d35409 942388a 3d35409 942388a 3d35409 942388a 3d35409 942388a 3d35409 942388a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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) |