032: Phase 5 (T036-T039) + Phase 6 (T040-T044) completed
T036: superset_compilation_adapter fully async T037: fileio.py async wrappers (aiofiles+run_blocking) T038-T039: tests for plugins fileio concurrency T040: providers async (aiosmtplib, httpx.AsyncClient) T041: dispatch_report parallel via asyncio.gather T042: TaskManager ThreadPoolExecutor->asyncio.create_task T043: EventBus asyncio.Queue(maxsize=10000) T044: lifecycle async context manager Remaining: T045 git/_base.py, T046-T048 tests, T005-T020-T021 tests
This commit is contained in:
@@ -1,110 +1,123 @@
|
||||
# #region EventBusModule [C:5] [TYPE Module] [SEMANTICS event,bus,log,buffer,subscriber,flush]
|
||||
# @BRIEF Task-scoped log buffering, persistence flushes, and subscriber fan-out for real-time
|
||||
# WebSocket observers — extracted from the monolithic TaskManager.
|
||||
# WebSocket observers — fully async with asyncio.Queue fan-out.
|
||||
# @LAYER Core
|
||||
# @RELATION DEPENDS_ON -> [LogEntry]
|
||||
# @RELATION DEPENDS_ON -> [TaskLogPersistenceService]
|
||||
# @RELATION DEPENDS_ON -> [should_log_task_level]
|
||||
# @RATIONALE Extracted from TaskManager to satisfy INV_7. EventBus owns log buffering, batch
|
||||
# persistence flushing, and pub/sub subscription management — a cohesive unit with
|
||||
# independent threading and I/O patterns.
|
||||
# @REJECTED Keeping log/event concerns inside TaskManager was rejected — the background flusher
|
||||
# thread, buffer locks, and subscriber queues are architecturally distinct from task
|
||||
# lifecycle state machines and the in-memory task registry.
|
||||
# @RATIONALE Extracted from TaskManager to satisfy INV_7. Migrated from threading to asyncio
|
||||
# for T043 — replaces queue.Queue + call_soon_threadsafe with asyncio.Queue
|
||||
# and async subscriber generator pattern.
|
||||
# @REJECTED Keeping threading-based EventBus was rejected — call_soon_threadsafe is fragile
|
||||
# under high load and cannot be tested without a running event loop.
|
||||
# @PRE TaskLogPersistenceService is initialized.
|
||||
# @POST Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers.
|
||||
# @SIDE_EFFECT Writes task logs to persistence, mutates in-memory buffers, pushes log entries
|
||||
# into subscriber queues, runs background flusher thread.
|
||||
# into subscriber queues via async put.
|
||||
# @INVARIANT Buffered logs are retried on persistence failure; every subscriber receives only
|
||||
# task-scoped events.
|
||||
# @DATA_CONTRACT Input: LogEntry -> Output: persisted log + subscriber notification
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from typing import Any
|
||||
|
||||
from ..cot_logger import seed_trace_id
|
||||
from ..logger import belief_scope, logger, should_log_task_level
|
||||
from .models import LogEntry, LogFilter, LogStats
|
||||
from .persistence import TaskLogPersistenceService
|
||||
from src.core.cot_logger import seed_trace_id
|
||||
from src.core.logger import belief_scope, logger, should_log_task_level
|
||||
from src.core.task_manager.models import LogEntry, LogFilter, LogStats
|
||||
from src.core.task_manager.persistence import TaskLogPersistenceService
|
||||
|
||||
QUEUE_MAXSIZE = 10000
|
||||
|
||||
|
||||
# #region EventBus [C:5] [TYPE Class] [SEMANTICS event,bus,log,buffer,subscriber,flush]
|
||||
# @BRIEF Coordinates task-scoped log buffering, persistence flushes, and subscriber fan-out
|
||||
# for real-time observers.
|
||||
# for real-time observers. Fully async — uses asyncio.Queue for fan-out delivery.
|
||||
# @RELATION DEPENDS_ON -> [LogEntry]
|
||||
# @RELATION DEPENDS_ON -> [TaskLogPersistenceService]
|
||||
# @RELATION DEPENDS_ON -> [should_log_task_level]
|
||||
# @PRE asyncio loop is available for thread-safe queue operations.
|
||||
# @PRE TaskLogPersistenceService is initialized.
|
||||
# @POST Accepted task logs are buffered, optionally persisted, and dispatched to active subscribers.
|
||||
# @INVARIANT Buffered logs are retried on persistence failure.
|
||||
# @RATIONALE Extracted as standalone pub/sub event bus to isolate threading (background flusher
|
||||
# thread, buffer locks) and subscription management from task lifecycle and registry.
|
||||
# @REJECTED Keeping log buffering inside TaskManager was rejected — the background flusher
|
||||
# thread and subscriber queues are architecturally distinct from task state machines
|
||||
# and should be independently stoppable/testable.
|
||||
# @RATIONALE Fully async event bus with asyncio.Queue(maxsize=10000) for each subscriber.
|
||||
# Broadcast methods use await queue.put() instead of call_soon_threadsafe.
|
||||
# @REJECTED Keeping threading-based EventBus with call_soon_threadsafe — fragile under load,
|
||||
# hard to test, and incompatible with fully async task execution.
|
||||
class EventBus:
|
||||
"""Task-scoped log buffering, persistence flushes, and subscriber fan-out."""
|
||||
"""Task-scoped log buffering, persistence flushes, and subscriber fan-out — fully async."""
|
||||
|
||||
LOG_FLUSH_INTERVAL = 2.0
|
||||
|
||||
# #region __init__ [TYPE Function]
|
||||
# @BRIEF Initialize empty log buffer, subscriber registry, status subscribers, and start background flusher.
|
||||
def __init__(self, log_persistence_service: TaskLogPersistenceService, loop: asyncio.AbstractEventLoop):
|
||||
# @BRIEF Initialize log buffer, subscriber registries, and async flusher.
|
||||
# @POST EventBus is ready to accept logs. Subscriber dicts are empty.
|
||||
def __init__(self, log_persistence_service: TaskLogPersistenceService):
|
||||
self.log_persistence_service = log_persistence_service
|
||||
self.loop = loop
|
||||
self.subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
# Per-task status subscribers (same pattern as log subscribers)
|
||||
self._status_subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
# Global task event subscribers (all tasks)
|
||||
self._task_event_subscribers: list[asyncio.Queue] = []
|
||||
# Maintenance event subscribers
|
||||
self._maintenance_subscribers: list[asyncio.Queue] = []
|
||||
self._log_buffer: dict[str, list[LogEntry]] = {}
|
||||
self._log_buffer_lock = threading.Lock()
|
||||
self._flusher_stop_event = threading.Event()
|
||||
self._flusher_thread = threading.Thread(
|
||||
target=self._flusher_loop, daemon=True
|
||||
)
|
||||
self._flusher_thread.start()
|
||||
self._flusher_stop_event = asyncio.Event()
|
||||
self._flusher_task: asyncio.Task | None = None
|
||||
# #endregion __init__
|
||||
|
||||
# #region start [TYPE Function]
|
||||
# @BRIEF Start the background async flusher task.
|
||||
# @POST Flusher task is created and scheduled in the event loop.
|
||||
def start(self) -> None:
|
||||
if self._flusher_task is None or self._flusher_task.done():
|
||||
self._flusher_task = asyncio.create_task(self.async_flusher_loop())
|
||||
# #endregion start
|
||||
|
||||
# #region stop [TYPE Function]
|
||||
# @BRIEF Signal the flusher thread to stop and wait for it.
|
||||
def stop(self) -> None:
|
||||
# @BRIEF Signal the flusher to stop and wait for it.
|
||||
# @POST Flusher task is cancelled and awaited.
|
||||
async def stop(self) -> None:
|
||||
self._flusher_stop_event.set()
|
||||
self._flusher_thread.join(timeout=2)
|
||||
if self._flusher_task and not self._flusher_task.done():
|
||||
self._flusher_task.cancel()
|
||||
from contextlib import suppress
|
||||
with suppress(asyncio.CancelledError, Exception):
|
||||
await self._flusher_task
|
||||
self._flusher_task = None
|
||||
# #endregion stop
|
||||
|
||||
# #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,thread]
|
||||
# @BRIEF Background thread that periodically flushes log buffer to database.
|
||||
def _flusher_loop(self):
|
||||
# #region async_flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,async]
|
||||
# @BRIEF Background async task that periodically flushes log buffer to database.
|
||||
async def async_flusher_loop(self):
|
||||
seed_trace_id()
|
||||
while not self._flusher_stop_event.is_set():
|
||||
self._flush_logs()
|
||||
self._flusher_stop_event.wait(self.LOG_FLUSH_INTERVAL)
|
||||
# #endregion _flusher_loop
|
||||
try:
|
||||
while not self._flusher_stop_event.is_set():
|
||||
await self._flush_logs()
|
||||
try:
|
||||
await asyncio.wait_for(
|
||||
self._flusher_stop_event.wait(),
|
||||
timeout=self.LOG_FLUSH_INTERVAL,
|
||||
)
|
||||
break # Event was set
|
||||
except TimeoutError:
|
||||
continue # Timeout is normal — flush again
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
# #endregion async_flusher_loop
|
||||
|
||||
# #region _flush_logs [C:3] [TYPE Function] [SEMANTICS flush,batch,persistence]
|
||||
# @BRIEF Flush all buffered logs to the database.
|
||||
# @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.add_logs]
|
||||
def _flush_logs(self):
|
||||
async def _flush_logs(self):
|
||||
seed_trace_id()
|
||||
with self._log_buffer_lock:
|
||||
task_ids = list(self._log_buffer.keys())
|
||||
task_ids = list(self._log_buffer.keys())
|
||||
for task_id in task_ids:
|
||||
with self._log_buffer_lock:
|
||||
logs = self._log_buffer.pop(task_id, [])
|
||||
logs = self._log_buffer.pop(task_id, [])
|
||||
if logs:
|
||||
try:
|
||||
self.log_persistence_service.add_logs(task_id, logs)
|
||||
logger.debug(f"Flushed {len(logs)} logs for task {task_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to flush logs for task {task_id}: {e}")
|
||||
with self._log_buffer_lock:
|
||||
if task_id not in self._log_buffer:
|
||||
self._log_buffer[task_id] = []
|
||||
self._log_buffer[task_id].extend(logs)
|
||||
if task_id not in self._log_buffer:
|
||||
self._log_buffer[task_id] = []
|
||||
self._log_buffer[task_id].extend(logs)
|
||||
# #endregion _flush_logs
|
||||
|
||||
# #region flush_task_logs [C:3] [TYPE Function] [SEMANTICS flush,single,persistence]
|
||||
@@ -112,10 +125,9 @@ class EventBus:
|
||||
# @PRE task_id exists.
|
||||
# @POST Task's buffered logs are written to database.
|
||||
# @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.add_logs]
|
||||
def flush_task_logs(self, task_id: str):
|
||||
async def flush_task_logs(self, task_id: str):
|
||||
with belief_scope("EventBus.flush_task_logs"):
|
||||
with self._log_buffer_lock:
|
||||
logs = self._log_buffer.pop(task_id, [])
|
||||
logs = self._log_buffer.pop(task_id, [])
|
||||
if logs:
|
||||
try:
|
||||
self.log_persistence_service.add_logs(task_id, logs)
|
||||
@@ -124,11 +136,11 @@ class EventBus:
|
||||
# #endregion flush_task_logs
|
||||
|
||||
# #region add_log [C:3] [TYPE Function] [SEMANTICS log,buffer,notify,subscriber]
|
||||
# @BRIEF Adds a log entry to a task buffer and notifies subscribers.
|
||||
# @BRIEF Adds a log entry to a task buffer and notifies subscribers via async put.
|
||||
# @PRE Task exists.
|
||||
# @POST Log added to buffer and pushed to queues (if level meets task_log_level filter).
|
||||
# @POST Log added to buffer and pushed to subscriber queues (if level passes filter).
|
||||
# @RELATION CALLS -> [should_log_task_level]
|
||||
def add_log(
|
||||
async def add_log(
|
||||
self,
|
||||
task_id: str,
|
||||
level: str,
|
||||
@@ -148,25 +160,31 @@ class EventBus:
|
||||
metadata=metadata,
|
||||
context=context,
|
||||
)
|
||||
# Add to in-memory logs (for backward compatibility with legacy JSON field)
|
||||
# Add to in-memory logs (backward compatibility)
|
||||
if task_logs_list is not None:
|
||||
task_logs_list.append(log_entry)
|
||||
# Add to buffer for batch persistence
|
||||
with self._log_buffer_lock:
|
||||
if task_id not in self._log_buffer:
|
||||
self._log_buffer[task_id] = []
|
||||
self._log_buffer[task_id].append(log_entry)
|
||||
# Notify subscribers (for real-time WebSocket updates)
|
||||
if task_id not in self._log_buffer:
|
||||
self._log_buffer[task_id] = []
|
||||
self._log_buffer[task_id].append(log_entry)
|
||||
# Notify subscribers via async put (fan-out)
|
||||
if task_id in self.subscribers:
|
||||
for queue in self.subscribers[task_id]:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, log_entry)
|
||||
try:
|
||||
await queue.put(log_entry)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
f"Subscriber queue full for task {task_id}, dropping log entry"
|
||||
)
|
||||
# #endregion add_log
|
||||
|
||||
# ── Log Subscribers ──
|
||||
|
||||
# #region subscribe_logs [C:2] [TYPE Function] [SEMANTICS subscribe,queue,realtime]
|
||||
# @BRIEF Subscribes to real-time logs for a task.
|
||||
# @POST Returns an asyncio.Queue for log entries.
|
||||
# @POST Returns an asyncio.Queue (maxsize=10000) for log entries.
|
||||
async def subscribe_logs(self, task_id: str) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
|
||||
if task_id not in self.subscribers:
|
||||
self.subscribers[task_id] = []
|
||||
self.subscribers[task_id].append(queue)
|
||||
@@ -188,9 +206,9 @@ class EventBus:
|
||||
|
||||
# #region subscribe_status [C:2] [TYPE Function] [SEMANTICS subscribe,status,realtime]
|
||||
# @BRIEF Subscribes to real-time status updates for a specific task.
|
||||
# @POST Returns an asyncio.Queue for Task status dicts.
|
||||
# @POST Returns an asyncio.Queue (maxsize=10000) for Task status dicts.
|
||||
async def subscribe_status(self, task_id: str) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
|
||||
if task_id not in self._status_subscribers:
|
||||
self._status_subscribers[task_id] = []
|
||||
self._status_subscribers[task_id].append(queue)
|
||||
@@ -198,7 +216,7 @@ class EventBus:
|
||||
# #endregion subscribe_status
|
||||
|
||||
# #region unsubscribe_status [C:2] [TYPE Function] [SEMANTICS unsubscribe,status,remove]
|
||||
# @BRIEF Unsubscribes from real-time status updates for a task.
|
||||
# @BRIEF Unsubscribes from status updates for a task.
|
||||
# @POST Queue removed from status subscribers.
|
||||
def unsubscribe_status(self, task_id: str, queue: asyncio.Queue):
|
||||
if task_id in self._status_subscribers:
|
||||
@@ -210,32 +228,35 @@ class EventBus:
|
||||
|
||||
# #region broadcast_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,notify]
|
||||
# @BRIEF Broadcast a task status update to all per-task and global subscribers.
|
||||
# @POST Status event pushed to all matching queues.
|
||||
def broadcast_status(self, task_id: str, task_dict: dict):
|
||||
from ..logger import logger as _lg
|
||||
# @POST Status event pushed to all matching queues via async put.
|
||||
async def broadcast_status(self, task_id: str, task_dict: dict):
|
||||
event = {"type": "task_status", "task_id": task_id, "task": task_dict}
|
||||
# Per-task subscribers
|
||||
if task_id in self._status_subscribers:
|
||||
for queue in self._status_subscribers[task_id]:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception as _e:
|
||||
_lg.warning("Failed to send status to per-task subscriber", extra={"task_id": task_id, "error": str(_e)})
|
||||
await queue.put(event)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
f"Status subscriber queue full for task {task_id}, dropping event"
|
||||
)
|
||||
# Global task event subscribers
|
||||
for queue in self._task_event_subscribers:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception as _e:
|
||||
_lg.warning("Failed to send status to global subscriber", extra={"task_id": task_id, "error": str(_e)})
|
||||
await queue.put(event)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
"Global task event subscriber queue full, dropping event"
|
||||
)
|
||||
# #endregion broadcast_status
|
||||
|
||||
# ── Maintenance Event Subscribers ──
|
||||
|
||||
# #region subscribe_maintenance_events [C:2] [TYPE Function] [SEMANTICS subscribe,maintenance,events]
|
||||
# @BRIEF Subscribes to maintenance event updates (created/ended/banner changes).
|
||||
# @POST Returns an asyncio.Queue for maintenance event dicts.
|
||||
# @POST Returns an asyncio.Queue (maxsize=10000) for maintenance event dicts.
|
||||
async def subscribe_maintenance_events(self) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
|
||||
self._maintenance_subscribers.append(queue)
|
||||
return queue
|
||||
# #endregion subscribe_maintenance_events
|
||||
@@ -250,23 +271,24 @@ class EventBus:
|
||||
|
||||
# #region broadcast_maintenance_event [C:2] [TYPE Function] [SEMANTICS broadcast,maintenance,notify]
|
||||
# @BRIEF Broadcast a maintenance event to all subscribers.
|
||||
# @POST Event pushed to all maintenance subscriber queues.
|
||||
def broadcast_maintenance_event(self, event: dict):
|
||||
from ..logger import logger as _lg
|
||||
# @POST Event pushed to all maintenance subscriber queues via async put.
|
||||
async def broadcast_maintenance_event(self, event: dict):
|
||||
for queue in self._maintenance_subscribers:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception as _e:
|
||||
_lg.warning("Failed to send maintenance event to subscriber", extra={"event_type": event.get("type"), "error": str(_e)})
|
||||
await queue.put(event)
|
||||
except asyncio.QueueFull:
|
||||
logger.warning(
|
||||
"Maintenance subscriber queue full, dropping event"
|
||||
)
|
||||
# #endregion broadcast_maintenance_event
|
||||
|
||||
# ── Global Task Event Subscribers ──
|
||||
|
||||
# #region subscribe_task_events [C:2] [TYPE Function] [SEMANTICS subscribe,global,events]
|
||||
# @BRIEF Subscribes to ALL task events (status changes) globally.
|
||||
# @POST Returns an asyncio.Queue for task event dicts.
|
||||
# @POST Returns an asyncio.Queue (maxsize=10000) for task event dicts.
|
||||
async def subscribe_task_events(self) -> asyncio.Queue:
|
||||
queue = asyncio.Queue()
|
||||
queue: asyncio.Queue = asyncio.Queue(maxsize=QUEUE_MAXSIZE)
|
||||
self._task_event_subscribers.append(queue)
|
||||
return queue
|
||||
# #endregion subscribe_task_events
|
||||
@@ -279,13 +301,15 @@ class EventBus:
|
||||
self._task_event_subscribers.remove(queue)
|
||||
# #endregion unsubscribe_task_events
|
||||
|
||||
# ── Log retrieval methods (sync, delegates to persistence) ──
|
||||
|
||||
# #region get_task_logs [C:3] [TYPE Function] [SEMANTICS logs,read,persistence,backfill]
|
||||
# @BRIEF Retrieves logs for a specific task (from memory for running, persistence for completed).
|
||||
# @RELATION CALLS -> [EXT:method:TaskLogPersistenceService.get_logs]
|
||||
def get_task_logs(
|
||||
self, task_id: str, log_filter: LogFilter | None = None, task_status=None, task_logs: list | None = None
|
||||
) -> list[LogEntry]:
|
||||
from .models import TaskStatus
|
||||
from src.core.task_manager.models import TaskStatus
|
||||
# For completed tasks, fetch from persistence
|
||||
if task_status in [TaskStatus.SUCCESS, TaskStatus.FAILED]:
|
||||
if log_filter is None:
|
||||
|
||||
@@ -23,16 +23,15 @@
|
||||
# a new wait future is created.
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from datetime import UTC, datetime
|
||||
import inspect
|
||||
from typing import Any
|
||||
|
||||
from ..cot_logger import seed_trace_id
|
||||
from ..logger import belief_scope, logger
|
||||
from .context import TaskContext
|
||||
from .graph import TaskGraph
|
||||
from .models import Task, TaskStatus
|
||||
from src.core.cot_logger import seed_trace_id
|
||||
from src.core.logger import belief_scope, logger
|
||||
from src.core.task_manager.context import TaskContext
|
||||
from src.core.task_manager.graph import TaskGraph
|
||||
from src.core.task_manager.models import Task, TaskStatus
|
||||
|
||||
def _task_to_dict(task: Task) -> dict:
|
||||
"""Serialize a Task to a dict for WebSocket broadcast."""
|
||||
@@ -78,15 +77,11 @@ class JobLifecycle:
|
||||
graph: TaskGraph,
|
||||
event_bus: Any,
|
||||
persistence_service: Any,
|
||||
executor: ThreadPoolExecutor,
|
||||
loop: asyncio.AbstractEventLoop,
|
||||
):
|
||||
self.plugin_loader = plugin_loader
|
||||
self.graph = graph
|
||||
self.event_bus = event_bus
|
||||
self.persistence_service = persistence_service
|
||||
self.executor = executor
|
||||
self.loop = loop
|
||||
self._dataset_subscribers: dict[str, list[asyncio.Queue]] = {}
|
||||
# #endregion __init__
|
||||
|
||||
@@ -107,16 +102,13 @@ class JobLifecycle:
|
||||
if not isinstance(params, dict):
|
||||
logger.error("Task parameters must be a dictionary.")
|
||||
raise ValueError("Task parameters must be a dictionary.")
|
||||
logger.reason("Creating task node and scheduling execution")
|
||||
logger.reason("Creating task node")
|
||||
task = Task(plugin_id=plugin_id, params=params, user_id=user_id)
|
||||
self.graph.add_task(task)
|
||||
self.persistence_service.persist_task(task)
|
||||
logger.info(f"Task {task.id} created and scheduled for execution")
|
||||
self.loop.create_task(
|
||||
self._run_task(task.id, add_log_callback=add_log_callback)
|
||||
)
|
||||
logger.info(f"Task {task.id} created (execution scheduled by TaskManager)")
|
||||
logger.reflect(
|
||||
"Task creation persisted and execution scheduled",
|
||||
"Task creation persisted",
|
||||
extra={"task_id": task.id, "plugin_id": plugin_id},
|
||||
)
|
||||
return task
|
||||
@@ -145,9 +137,9 @@ class JobLifecycle:
|
||||
task.status = TaskStatus.RUNNING
|
||||
task.started_at = datetime.now(UTC)
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
await self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
await add_log_callback(
|
||||
task_id, "INFO",
|
||||
f"Task started for plugin '{plugin.name}'",
|
||||
source="system",
|
||||
@@ -168,22 +160,21 @@ class JobLifecycle:
|
||||
if asyncio.iscoroutinefunction(plugin.execute):
|
||||
task.result = await plugin.execute(params, context=context)
|
||||
else:
|
||||
task.result = await self.loop.run_in_executor(
|
||||
self.executor,
|
||||
lambda: plugin.execute(params, context=context),
|
||||
task.result = await asyncio.to_thread(
|
||||
plugin.execute, params, context=context,
|
||||
)
|
||||
else:
|
||||
if asyncio.iscoroutinefunction(plugin.execute):
|
||||
task.result = await plugin.execute(params)
|
||||
else:
|
||||
task.result = await self.loop.run_in_executor(
|
||||
self.executor, plugin.execute, params
|
||||
task.result = await asyncio.to_thread(
|
||||
plugin.execute, params
|
||||
)
|
||||
logger.info(f"Task {task_id} completed successfully")
|
||||
task.status = TaskStatus.SUCCESS
|
||||
self._broadcast_task_status(task)
|
||||
await self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
await add_log_callback(
|
||||
task_id, "INFO",
|
||||
f"Task completed successfully for plugin '{plugin.name}'",
|
||||
source="system",
|
||||
@@ -191,9 +182,9 @@ class JobLifecycle:
|
||||
except Exception as e:
|
||||
logger.error(f"Task {task_id} failed: {e}")
|
||||
task.status = TaskStatus.FAILED
|
||||
self._broadcast_task_status(task)
|
||||
await self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
await add_log_callback(
|
||||
task_id, "ERROR",
|
||||
f"Task failed: {e}",
|
||||
source="system",
|
||||
@@ -201,7 +192,7 @@ class JobLifecycle:
|
||||
)
|
||||
finally:
|
||||
task.finished_at = datetime.now(UTC)
|
||||
self.event_bus.flush_task_logs(task_id)
|
||||
await self.event_bus.flush_task_logs(task_id)
|
||||
self.persistence_service.persist_task(task)
|
||||
logger.info(
|
||||
f"Task {task_id} execution finished with status: {task.status}"
|
||||
@@ -215,7 +206,7 @@ class JobLifecycle:
|
||||
dataset_ids = task.params.get("dataset_ids") or [task.params.get("dataset_id")]
|
||||
env_id = task.params.get("env") or task.params.get("environment_id", "")
|
||||
if dataset_ids and env_id:
|
||||
self._broadcast_dataset_updated(env_id, dataset_ids)
|
||||
await self._broadcast_dataset_updated(env_id, dataset_ids)
|
||||
# #endregion _run_task
|
||||
|
||||
# #region resolve_task [C:3] [TYPE Function] [SEMANTICS task,resolve,mapping,resume]
|
||||
@@ -233,7 +224,7 @@ class JobLifecycle:
|
||||
task.params.update(resolution_params)
|
||||
task.status = TaskStatus.RUNNING
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
await self._broadcast_task_status(task)
|
||||
self.graph.resolve_future(task_id, True)
|
||||
# #endregion resolve_task
|
||||
|
||||
@@ -248,8 +239,9 @@ class JobLifecycle:
|
||||
return
|
||||
task.status = TaskStatus.AWAITING_MAPPING
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
future = self.loop.create_future()
|
||||
await self._broadcast_task_status(task)
|
||||
loop = asyncio.get_running_loop()
|
||||
future = loop.create_future()
|
||||
self.graph.create_future(task_id, future)
|
||||
try:
|
||||
await future
|
||||
@@ -266,7 +258,8 @@ class JobLifecycle:
|
||||
task = self.graph.get_task(task_id)
|
||||
if not task:
|
||||
return
|
||||
future = self.loop.create_future()
|
||||
loop = asyncio.get_running_loop()
|
||||
future = loop.create_future()
|
||||
self.graph.create_future(task_id, future)
|
||||
try:
|
||||
await future
|
||||
@@ -279,7 +272,7 @@ class JobLifecycle:
|
||||
# @PRE Task exists and is in RUNNING state.
|
||||
# @POST Task status changed to AWAITING_INPUT, input_request set, persisted.
|
||||
# @RAISES ValueError if task not found or not RUNNING.
|
||||
def await_input(
|
||||
async def await_input(
|
||||
self, task_id: str, input_request: dict[str, Any],
|
||||
add_log_callback=None,
|
||||
) -> None:
|
||||
@@ -295,9 +288,9 @@ class JobLifecycle:
|
||||
task.input_required = True
|
||||
task.input_request = input_request
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
await self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
await add_log_callback(
|
||||
task_id, "INFO",
|
||||
"Task paused for user input",
|
||||
metadata={"input_request": input_request},
|
||||
@@ -309,7 +302,7 @@ class JobLifecycle:
|
||||
# @PRE Task exists and is in AWAITING_INPUT state.
|
||||
# @POST Task status changed to RUNNING, passwords injected, task resumed.
|
||||
# @RAISES ValueError if task not found, not awaiting input, or passwords invalid.
|
||||
def resume_task_with_password(
|
||||
async def resume_task_with_password(
|
||||
self, task_id: str, passwords: dict[str, str],
|
||||
add_log_callback=None,
|
||||
) -> None:
|
||||
@@ -330,9 +323,9 @@ class JobLifecycle:
|
||||
task.input_request = None
|
||||
task.status = TaskStatus.RUNNING
|
||||
self.persistence_service.persist_task(task)
|
||||
self._broadcast_task_status(task)
|
||||
await self._broadcast_task_status(task)
|
||||
if add_log_callback:
|
||||
add_log_callback(
|
||||
await add_log_callback(
|
||||
task_id, "INFO",
|
||||
"Task resumed with passwords",
|
||||
metadata={"databases": list(passwords.keys())},
|
||||
@@ -363,19 +356,19 @@ class JobLifecycle:
|
||||
# #region _broadcast_task_status [C:2] [TYPE Function] [SEMANTICS broadcast,status,websocket]
|
||||
# @BRIEF Broadcast task status update to all subscribers via EventBus.
|
||||
# @POST EventBus.broadcast_status called with serialized task dict.
|
||||
def _broadcast_task_status(self, task: Task):
|
||||
self.event_bus.broadcast_status(task.id, _task_to_dict(task))
|
||||
async def _broadcast_task_status(self, task: Task):
|
||||
await self.event_bus.broadcast_status(task.id, _task_to_dict(task))
|
||||
# #endregion _broadcast_task_status
|
||||
|
||||
# #region _broadcast_dataset_updated [C:3] [TYPE Function] [SEMANTICS dataset,broadcast,event,websocket]
|
||||
# @BRIEF Broadcast dataset.updated event to all subscribers for a given environment.
|
||||
def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]):
|
||||
async def _broadcast_dataset_updated(self, env_id: str, dataset_ids: list[int]):
|
||||
event = {"type": "dataset.updated", "payload": {"env_id": env_id, "dataset_ids": dataset_ids}}
|
||||
if env_id in self._dataset_subscribers:
|
||||
for queue in self._dataset_subscribers[env_id]:
|
||||
try:
|
||||
self.loop.call_soon_threadsafe(queue.put_nowait, event)
|
||||
except Exception:
|
||||
await queue.put(event)
|
||||
except asyncio.QueueFull:
|
||||
pass
|
||||
# #endregion _broadcast_dataset_updated
|
||||
# #endregion JobLifecycle
|
||||
|
||||
@@ -33,15 +33,14 @@
|
||||
# with async execution, and created an unmaintainable god class.
|
||||
|
||||
import asyncio
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from typing import Any
|
||||
|
||||
from ..logger import belief_scope, logger
|
||||
from .event_bus import EventBus
|
||||
from .graph import TaskGraph
|
||||
from .lifecycle import JobLifecycle
|
||||
from .models import LogFilter, LogStats, Task, TaskStatus
|
||||
from .persistence import TaskLogPersistenceService, TaskPersistenceService
|
||||
from src.core.logger import belief_scope, logger
|
||||
from src.core.task_manager.event_bus import EventBus
|
||||
from src.core.task_manager.graph import TaskGraph
|
||||
from src.core.task_manager.lifecycle import JobLifecycle
|
||||
from src.core.task_manager.models import LogFilter, LogStats, Task, TaskStatus
|
||||
from src.core.task_manager.persistence import TaskLogPersistenceService, TaskPersistenceService
|
||||
|
||||
|
||||
# #region TaskManager [C:5] [TYPE Class] [SEMANTICS task, manager, lifecycle, execution, state]
|
||||
@@ -81,18 +80,14 @@ class TaskManager:
|
||||
with belief_scope("TaskManager.__init__"):
|
||||
logger.reason("Initializing task manager runtime services")
|
||||
self.plugin_loader = plugin_loader
|
||||
self.executor = ThreadPoolExecutor(max_workers=5)
|
||||
# Track running asyncio tasks for management/cancellation
|
||||
self._async_tasks: dict[str, asyncio.Task] = {}
|
||||
self.persistence_service = TaskPersistenceService()
|
||||
self.log_persistence_service = TaskLogPersistenceService()
|
||||
|
||||
try:
|
||||
self.loop = asyncio.get_running_loop()
|
||||
except RuntimeError:
|
||||
self.loop = asyncio.get_event_loop()
|
||||
|
||||
# Create sub-services
|
||||
self.graph = TaskGraph(self.persistence_service)
|
||||
self.event_bus = EventBus(self.log_persistence_service, self.loop)
|
||||
self.event_bus = EventBus(self.log_persistence_service)
|
||||
|
||||
# Create add_log callback that lifecycle and plugins use
|
||||
self._add_log = self._make_add_log_callback()
|
||||
@@ -102,8 +97,6 @@ class TaskManager:
|
||||
graph=self.graph,
|
||||
event_bus=self.event_bus,
|
||||
persistence_service=self.persistence_service,
|
||||
executor=self.executor,
|
||||
loop=self.loop,
|
||||
)
|
||||
|
||||
# Load persisted tasks on startup
|
||||
@@ -112,11 +105,6 @@ class TaskManager:
|
||||
# Backward-compatible property aliases for tests
|
||||
self.tasks = self.graph.tasks
|
||||
self.task_futures = self.graph.task_futures
|
||||
self.subscribers = self.event_bus.subscribers
|
||||
self._log_buffer = self.event_bus._log_buffer
|
||||
self._log_buffer_lock = self.event_bus._log_buffer_lock
|
||||
self._flusher_stop_event = self.event_bus._flusher_stop_event
|
||||
self._flusher_thread = self.event_bus._flusher_thread
|
||||
|
||||
logger.reflect(
|
||||
"Task manager runtime initialized",
|
||||
@@ -124,14 +112,14 @@ class TaskManager:
|
||||
)
|
||||
# #endregion __init__
|
||||
|
||||
# #region _make_add_log_callback [C:2] [TYPE Function]
|
||||
# @BRIEF Create a closure for adding logs that looks up the task and delegates to EventBus.
|
||||
# #region _make_add_log_callback [C:3] [TYPE Function]
|
||||
# @BRIEF Create an async closure for adding logs that looks up the task and delegates to EventBus.
|
||||
def _make_add_log_callback(self):
|
||||
def _add_log(task_id, level, message, source="system", metadata=None, context=None):
|
||||
async def _add_log(task_id, level, message, source="system", metadata=None, context=None):
|
||||
task = self.graph.get_task(task_id)
|
||||
if not task:
|
||||
return
|
||||
self.event_bus.add_log(
|
||||
await self.event_bus.add_log(
|
||||
task_id, level, message, source, metadata, context,
|
||||
task_logs_list=task.logs,
|
||||
)
|
||||
@@ -140,22 +128,22 @@ class TaskManager:
|
||||
|
||||
# ── Deprecated legacy aliases (delegate to EventBus) ──
|
||||
|
||||
# #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,thread]
|
||||
# @BRIEF Legacy alias delegating to EventBus._flusher_loop.
|
||||
def _flusher_loop(self):
|
||||
self.event_bus._flusher_loop()
|
||||
# #region _flusher_loop [C:3] [TYPE Function] [SEMANTICS flush,background,async]
|
||||
# @BRIEF Legacy alias delegating to EventBus.async_flusher_loop.
|
||||
async def _flusher_loop(self):
|
||||
await self.event_bus.async_flusher_loop()
|
||||
# #endregion _flusher_loop
|
||||
|
||||
# #region _flush_logs [C:3] [TYPE Function] [SEMANTICS flush,batch,persistence]
|
||||
# @BRIEF Legacy alias delegating to EventBus._flush_logs.
|
||||
def _flush_logs(self):
|
||||
self.event_bus._flush_logs()
|
||||
async def _flush_logs(self):
|
||||
await self.event_bus._flush_logs()
|
||||
# #endregion _flush_logs
|
||||
|
||||
# #region _flush_task_logs [C:3] [TYPE Function] [SEMANTICS flush,single,persistence]
|
||||
# @BRIEF Legacy alias delegating to EventBus.flush_task_logs.
|
||||
def _flush_task_logs(self, task_id: str):
|
||||
self.event_bus.flush_task_logs(task_id)
|
||||
async def _flush_task_logs(self, task_id: str):
|
||||
await self.event_bus.flush_task_logs(task_id)
|
||||
# #endregion _flush_task_logs
|
||||
|
||||
# ── Task CRUD delegates to TaskGraph ──
|
||||
@@ -294,18 +282,45 @@ class TaskManager:
|
||||
async def create_task(
|
||||
self, plugin_id: str, params: dict[str, Any], user_id: str | None = None
|
||||
) -> Task:
|
||||
return await self.lifecycle.create_task(
|
||||
task = await self.lifecycle.create_task(
|
||||
plugin_id, params, user_id,
|
||||
add_log_callback=self._add_log,
|
||||
)
|
||||
# Schedule execution and track the asyncio task
|
||||
async_task = asyncio.create_task(
|
||||
self.lifecycle._run_task(task.id, add_log_callback=self._add_log)
|
||||
)
|
||||
self._async_tasks[task.id] = async_task
|
||||
return task
|
||||
# #endregion create_task
|
||||
|
||||
# #region _run_task [TYPE Function] [C:4]
|
||||
# @BRIEF Internal method to execute a task with TaskContext support (delegates to lifecycle).
|
||||
# Tracks the asyncio.Task for management/cancellation.
|
||||
async def _run_task(self, task_id: str):
|
||||
await self.lifecycle._run_task(task_id, add_log_callback=self._add_log)
|
||||
async_task = asyncio.create_task(
|
||||
self.lifecycle._run_task(task_id, add_log_callback=self._add_log)
|
||||
)
|
||||
self._async_tasks[task_id] = async_task
|
||||
try:
|
||||
await async_task
|
||||
finally:
|
||||
self._async_tasks.pop(task_id, None)
|
||||
# #endregion _run_task
|
||||
|
||||
# #region cancel_task [TYPE Function] [C:3]
|
||||
# @BRIEF Cancel a running task by ID.
|
||||
# @PRE Task must be currently tracked as running.
|
||||
# @POST Task is cancelled and removed from tracking dict.
|
||||
async def cancel_task(self, task_id: str) -> bool:
|
||||
async_task = self._async_tasks.get(task_id)
|
||||
if async_task is None or async_task.done():
|
||||
return False
|
||||
async_task.cancel()
|
||||
self._async_tasks.pop(task_id, None)
|
||||
return True
|
||||
# #endregion cancel_task
|
||||
|
||||
# #region resolve_task [TYPE Function] [C:3]
|
||||
# @BRIEF Resumes a task that is awaiting mapping.
|
||||
async def resolve_task(self, task_id: str, resolution_params: dict[str, Any]):
|
||||
@@ -326,8 +341,8 @@ class TaskManager:
|
||||
|
||||
# #region await_input [TYPE Function] [C:3]
|
||||
# @BRIEF Transition a task to AWAITING_INPUT state with input request.
|
||||
def await_input(self, task_id: str, input_request: dict[str, Any]) -> None:
|
||||
self.lifecycle.await_input(
|
||||
async def await_input(self, task_id: str, input_request: dict[str, Any]) -> None:
|
||||
await self.lifecycle.await_input(
|
||||
task_id, input_request,
|
||||
add_log_callback=self._add_log,
|
||||
)
|
||||
@@ -335,10 +350,10 @@ class TaskManager:
|
||||
|
||||
# #region resume_task_with_password [TYPE Function] [C:3]
|
||||
# @BRIEF Resume a task that is awaiting input with provided passwords.
|
||||
def resume_task_with_password(
|
||||
async def resume_task_with_password(
|
||||
self, task_id: str, passwords: dict[str, str]
|
||||
) -> None:
|
||||
self.lifecycle.resume_task_with_password(
|
||||
await self.lifecycle.resume_task_with_password(
|
||||
task_id, passwords,
|
||||
add_log_callback=self._add_log,
|
||||
)
|
||||
@@ -360,8 +375,8 @@ class TaskManager:
|
||||
|
||||
# #region broadcast_maintenance_event [TYPE Function] [C:2]
|
||||
# @BRIEF Broadcast a maintenance event to all subscribers.
|
||||
def broadcast_maintenance_event(self, event: dict):
|
||||
self.event_bus.broadcast_maintenance_event(event)
|
||||
async def broadcast_maintenance_event(self, event: dict):
|
||||
await self.event_bus.broadcast_maintenance_event(event)
|
||||
# #endregion broadcast_maintenance_event
|
||||
|
||||
# ── Dataset event delegates to JobLifecycle ──
|
||||
|
||||
@@ -16,7 +16,7 @@ from typing import Any, LiteralString
|
||||
import zipfile
|
||||
import zlib
|
||||
|
||||
from ..logger import belief_scope, logger as app_logger
|
||||
from src.core.logger import belief_scope, logger as app_logger
|
||||
|
||||
|
||||
# #region InvalidZipFormatError [TYPE Class]
|
||||
@@ -415,4 +415,107 @@ def consolidate_archive_folders(root_directory: Path) -> None:
|
||||
except Exception as e:
|
||||
app_logger.error("[consolidate_archive_folders][Failure] Failed to remove source directory %s: %s", source_dir, e)
|
||||
# #endregion consolidate_archive_folders
|
||||
# #region AsyncFileIOWrappers [C:3] [TYPE Module] [SEMANTICS async, fileio, aiofiles, run_blocking]
|
||||
# @BRIEF Async wrappers for file I/O operations using run_blocking executor and aiofiles.
|
||||
# @LAYER Infrastructure
|
||||
# @RELATION DEPENDS_ON -> [BlockingExecutorsModule]
|
||||
# @RELATION DEPENDS_ON -> [aiofiles]
|
||||
# @PRE init_executors has been called (executors ready).
|
||||
# @POST Async wrappers delegate blocking I/O to thread pool or aiofiles.
|
||||
import aiofiles # noqa: E402
|
||||
|
||||
from .executors import run_blocking # noqa: E402
|
||||
|
||||
|
||||
# #region aread_dashboard_from_disk [C:3] [TYPE Function] [SEMANTICS async, file, read]
|
||||
# @BRIEF Async read dashboard file from disk using aiofiles for large files or run_blocking fallback.
|
||||
# @PRE file_path must point to an existing file.
|
||||
# @POST Returns file bytes content.
|
||||
# @SIDE_EFFECT Reads file from disk (non-blocking via aiofiles or thread pool).
|
||||
async def aread_dashboard_from_disk(file_path: str) -> bytes:
|
||||
with belief_scope(f"Async read dashboard from {file_path}"):
|
||||
path = Path(file_path)
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"Dashboard file not found: {file_path}")
|
||||
if path.stat().st_size > 10 * 1024 * 1024: # > 10MB → use aiofiles
|
||||
async with aiofiles.open(file_path, "rb") as f:
|
||||
content = await f.read()
|
||||
else:
|
||||
content = await run_blocking(kind="file", fn=path.read_bytes)
|
||||
if not content:
|
||||
app_logger.warning(
|
||||
"[aread_dashboard_from_disk][Warning] File is empty: %s", file_path
|
||||
)
|
||||
return content
|
||||
# #endregion aread_dashboard_from_disk
|
||||
|
||||
|
||||
# #region asave_and_unpack_dashboard [C:3] [TYPE Function] [SEMANTICS async, zip, save, unpack]
|
||||
# @BRIEF Async save ZIP content to disk and optionally unpack it.
|
||||
# @PRE zip_content must be valid ZIP archive bytes.
|
||||
# @POST ZIP file saved to target_dir; if unpack=True, contents extracted.
|
||||
# @SIDE_EFFECT Writes ZIP to disk; may unpack archive.
|
||||
async def asave_and_unpack_dashboard(
|
||||
zip_content: bytes, target_dir: str, unpack: bool = False, original_filename: str | None = None
|
||||
) -> tuple[Path, Path | None]:
|
||||
with belief_scope(f"Async save and unpack dashboard in {target_dir}"):
|
||||
app_logger.info(
|
||||
"[asave_and_unpack_dashboard][Enter] Processing dashboard. Unpack: %s", unpack
|
||||
)
|
||||
output_path = Path(target_dir)
|
||||
await run_blocking(kind="file", fn=output_path.mkdir, parents=True, exist_ok=True)
|
||||
zip_name = (
|
||||
sanitize_filename(original_filename)
|
||||
if original_filename
|
||||
else f"dashboard_export_{datetime.now().strftime('%Y%m%d_%H%M%S')}.zip"
|
||||
)
|
||||
zip_path = output_path / zip_name
|
||||
async with aiofiles.open(zip_path, "wb") as f:
|
||||
await f.write(zip_content)
|
||||
app_logger.info(
|
||||
"[asave_and_unpack_dashboard][State] Dashboard saved to: %s", zip_path
|
||||
)
|
||||
if unpack:
|
||||
await run_blocking(
|
||||
"file",
|
||||
_unzip_to_directory,
|
||||
str(zip_path),
|
||||
str(output_path),
|
||||
)
|
||||
app_logger.info(
|
||||
"[asave_and_unpack_dashboard][State] Dashboard unpacked to: %s", output_path
|
||||
)
|
||||
return zip_path, output_path
|
||||
return zip_path, None
|
||||
# #endregion asave_and_unpack_dashboard
|
||||
|
||||
|
||||
# #region _unzip_to_directory [C:1] [TYPE Function] [SEMANTICS zip, helper]
|
||||
# @BRIEF (Helper) Extract ZIP file to directory. Used by asave_and_unpack_dashboard in thread pool.
|
||||
def _unzip_to_directory(zip_path: str, output_path: str) -> None:
|
||||
with zipfile.ZipFile(zip_path, "r") as zip_ref:
|
||||
zip_ref.extractall(output_path)
|
||||
# #endregion _unzip_to_directory
|
||||
|
||||
|
||||
# #region acleanup_directory [C:2] [TYPE Function] [SEMANTICS async, cleanup, directory]
|
||||
# @BRIEF Async remove directory tree using run_blocking executor.
|
||||
# @PRE path is a valid filesystem path.
|
||||
# @POST Directory and all contents removed.
|
||||
# @SIDE_EFFECT Deletes files and directories from disk.
|
||||
async def acleanup_directory(path: str) -> None:
|
||||
with belief_scope(f"Async cleanup directory: {path}"):
|
||||
await run_blocking("file", shutil.rmtree, path)
|
||||
# #endregion acleanup_directory
|
||||
|
||||
|
||||
# #region aread_bytes [C:2] [TYPE Function] [SEMANTICS async, read, bytes]
|
||||
# @BRIEF Async read file bytes from a Path using run_blocking executor.
|
||||
# @PRE path must point to an existing file.
|
||||
# @POST Returns file bytes content.
|
||||
async def aread_bytes(path: Path) -> bytes:
|
||||
with belief_scope(f"Async read bytes: {path}"):
|
||||
return await run_blocking(kind="file", fn=path.read_bytes)
|
||||
# #endregion aread_bytes
|
||||
# #endregion AsyncFileIOWrappers
|
||||
# #endregion FileIO
|
||||
|
||||
@@ -71,7 +71,7 @@ class SupersetCompilationAdapter:
|
||||
# @POST returns a ready or failed preview artifact backed only by Superset-originated SQL or diagnostics.
|
||||
# @SIDE_EFFECT performs upstream preview requests.
|
||||
# @DATA_CONTRACT Input[PreviewCompilationPayload] -> Output[CompiledPreview]
|
||||
def compile_preview(self, payload: PreviewCompilationPayload) -> CompiledPreview:
|
||||
async def compile_preview(self, payload: PreviewCompilationPayload) -> CompiledPreview:
|
||||
with belief_scope("SupersetCompilationAdapter.compile_preview"):
|
||||
if payload.dataset_id <= 0:
|
||||
logger.explore(
|
||||
@@ -92,7 +92,7 @@ class SupersetCompilationAdapter:
|
||||
},
|
||||
)
|
||||
try:
|
||||
preview_result = self._request_superset_preview(payload)
|
||||
preview_result = await self._request_superset_preview(payload)
|
||||
except Exception as exc:
|
||||
logger.explore(
|
||||
"Superset preview compilation failed with explicit upstream error",
|
||||
@@ -167,7 +167,7 @@ class SupersetCompilationAdapter:
|
||||
# @POST returns one canonical SQL Lab session reference from Superset.
|
||||
# @SIDE_EFFECT performs upstream SQL Lab execution/session creation.
|
||||
# @DATA_CONTRACT Input[SqlLabLaunchPayload] -> Output[str]
|
||||
def create_sql_lab_session(self, payload: SqlLabLaunchPayload) -> str:
|
||||
async def create_sql_lab_session(self, payload: SqlLabLaunchPayload) -> str:
|
||||
with belief_scope("SupersetCompilationAdapter.create_sql_lab_session"):
|
||||
compiled_sql = str(payload.compiled_sql or "").strip()
|
||||
if not compiled_sql:
|
||||
@@ -187,7 +187,7 @@ class SupersetCompilationAdapter:
|
||||
"preview_id": payload.preview_id,
|
||||
},
|
||||
)
|
||||
result = self._request_sql_lab_session(payload)
|
||||
result = await self._request_sql_lab_session(payload)
|
||||
sql_lab_session_ref = str(
|
||||
result.get("sql_lab_session_ref")
|
||||
or result.get("query_id")
|
||||
@@ -223,7 +223,7 @@ class SupersetCompilationAdapter:
|
||||
# @POST returns one normalized upstream compilation response including the chosen strategy metadata.
|
||||
# @SIDE_EFFECT issues one or more Superset preview requests through the client fallback chain.
|
||||
# @DATA_CONTRACT Input[PreviewCompilationPayload] -> Output[Dict[str,Any]]
|
||||
def _request_superset_preview(
|
||||
async def _request_superset_preview(
|
||||
self, payload: PreviewCompilationPayload
|
||||
) -> dict[str, Any]:
|
||||
direct_compile_preview = getattr(self.client, "compile_preview", None)
|
||||
@@ -238,9 +238,9 @@ class SupersetCompilationAdapter:
|
||||
"session_id": payload.session_id,
|
||||
},
|
||||
)
|
||||
response = direct_compile_preview(payload)
|
||||
response = await direct_compile_preview(payload)
|
||||
except TypeError:
|
||||
response = direct_compile_preview(
|
||||
response = await direct_compile_preview(
|
||||
payload.dataset_id,
|
||||
template_params=payload.template_params,
|
||||
effective_filters=payload.effective_filters,
|
||||
@@ -274,7 +274,7 @@ class SupersetCompilationAdapter:
|
||||
"template_param_count": len(payload.template_params),
|
||||
},
|
||||
)
|
||||
response = direct_compile_dataset_preview(
|
||||
response = await direct_compile_dataset_preview(
|
||||
dataset_id=payload.dataset_id,
|
||||
template_params=payload.template_params,
|
||||
effective_filters=payload.effective_filters,
|
||||
@@ -306,7 +306,7 @@ class SupersetCompilationAdapter:
|
||||
},
|
||||
)
|
||||
if self._supports_client_method("compile_dataset_preview"):
|
||||
response = self.client.compile_dataset_preview(
|
||||
response = await self.client.compile_dataset_preview(
|
||||
dataset_id=payload.dataset_id,
|
||||
template_params=payload.template_params,
|
||||
effective_filters=payload.effective_filters,
|
||||
@@ -323,7 +323,7 @@ class SupersetCompilationAdapter:
|
||||
f"/dataset/{payload.dataset_id}/sql",
|
||||
):
|
||||
try:
|
||||
response = self.client.network.request(
|
||||
response = await self.client.network.request(
|
||||
method="POST",
|
||||
endpoint=endpoint,
|
||||
data=self._dump_json(
|
||||
@@ -361,8 +361,8 @@ class SupersetCompilationAdapter:
|
||||
# @POST returns the first successful SQL Lab execution response from Superset.
|
||||
# @SIDE_EFFECT issues Superset dataset lookup and SQL Lab execution requests.
|
||||
# @DATA_CONTRACT Input[SqlLabLaunchPayload] -> Output[Dict[str,Any]]
|
||||
def _request_sql_lab_session(self, payload: SqlLabLaunchPayload) -> dict[str, Any]:
|
||||
dataset_raw = self.client.get_dataset(payload.dataset_id)
|
||||
async def _request_sql_lab_session(self, payload: SqlLabLaunchPayload) -> dict[str, Any]:
|
||||
dataset_raw = await self.client.get_dataset(payload.dataset_id)
|
||||
dataset_record = (
|
||||
dataset_raw.get("result", dataset_raw)
|
||||
if isinstance(dataset_raw, dict)
|
||||
@@ -391,7 +391,7 @@ class SupersetCompilationAdapter:
|
||||
errors: list[str] = []
|
||||
for candidate in candidate_calls:
|
||||
try:
|
||||
response = self.client.network.request(
|
||||
response = await self.client.network.request(
|
||||
method=candidate["http_method"],
|
||||
endpoint=candidate["target"],
|
||||
data=self._dump_json(request_payload),
|
||||
|
||||
Reference in New Issue
Block a user