Update util/trace.py

util/trace.py

@@ -1,37 +1,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
 
-with open(os.getenv("GOOGLE_APPLICATION_CREDENTIALS"), "w") as f:
-    f.write(os.getenv("gcp_service_json"))
-    f.close()
+class TraceManager:
+    def __init__(self):
+        self.client = self._initialize_client()
+        self.project_name = f"projects/{os.getenv('PROJECT')}"
 
-def initialize_trace_client():
-    credentials = service_account.Credentials.from_service_account_file(os.getenv("GOOGLE_APPLICATION_CREDENTIALS"))
-    client = trace_v2.TraceServiceClient(credentials=credentials)
-    return client, 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
 
-from google.protobuf.timestamp_pb2 import Timestamp
-from datetime import datetime
+    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_span(client, parent_span_id=None, span_display_name="Default Span"):
-    trace_id = uuid.uuid4().hex
-    span_id = uuid.uuid4().hex[:16]
-    span_name = f"{os.getenv("PROJECT")}/traces/{trace_id}/spans/{span_id}"
-    now = datetime.utcnow()
-    timestamp = Timestamp()
-    timestamp.FromDatetime(now)
-    
-    span = trace_v2.Span(
-        name=span_name,
-        span_id=span_id,
-        parent_span_id=parent_span_id,
-        display_name=trace_v2.types.TruncatableString(value=span_display_name),
-        start_time=timestamp,
-        end_time=timestamp
-    )
-    
-    client.batch_write_spans(name=os.getenv("PROJECT"), spans=[span])
-    return span_name  # 返り値としてスパンの名前を返す
+    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)