53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from datetime import datetime, timezone
|
|
import json
|
|
from typing import Any, Dict, Optional
|
|
|
|
from sqlalchemy import func
|
|
|
|
from .models import WorkflowEvent
|
|
|
|
|
|
class EventWriter:
|
|
def __init__(self, run_id: str, workflow_id: str, db_factory) -> None:
|
|
self.run_id = run_id
|
|
self.workflow_id = workflow_id
|
|
self.db_factory = db_factory
|
|
self._lock = asyncio.Lock()
|
|
db = db_factory()
|
|
try:
|
|
self._sequence = int(db.query(func.max(WorkflowEvent.sequence)).filter(WorkflowEvent.workflow_run_id == run_id).scalar() or 0)
|
|
finally:
|
|
db.close()
|
|
|
|
async def emit(self, event_type: str, payload: Dict[str, Any], node_id: Optional[str] = None) -> Dict[str, Any]:
|
|
async with self._lock:
|
|
self._sequence += 1
|
|
timestamp = datetime.now(timezone.utc)
|
|
event = {
|
|
"run_id": self.run_id,
|
|
"workflow_id": self.workflow_id,
|
|
"node_id": node_id,
|
|
"sequence": self._sequence,
|
|
"timestamp": timestamp.isoformat(),
|
|
"type": event_type,
|
|
"payload": payload,
|
|
}
|
|
db = self.db_factory()
|
|
try:
|
|
db.add(WorkflowEvent(
|
|
workflow_run_id=self.run_id,
|
|
workflow_id=self.workflow_id,
|
|
node_id=node_id,
|
|
sequence=self._sequence,
|
|
event_type=event_type,
|
|
timestamp=timestamp.replace(tzinfo=None),
|
|
payload_json=json.dumps(payload, ensure_ascii=False, default=str),
|
|
))
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
return event
|