first commit

ui/runner.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections import deque
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Deque, Dict, List, Optional, Tuple
+
+import sync_state_machine.domain  # noqa: F401
+from sync_state_machine.pipeline import (
+    FullSyncPipeline,
+    create_pipeline_from_config,
+)
+from sync_state_machine.config import (
+    PipelineRunConfig,
+    JsonlDataSourceConfig,
+    ApiDataSourceConfig,
+    apply_config_overrides,
+    resolve_log_file_path,
+)
+from sync_state_machine.pipeline.run_logger import clear_pipeline_logger_handlers, init_pipeline_logger
+
+from .config_models import PipelineUiConfig, default_ui_config
+
+
+@dataclass
+class RunnerStatus:
+    running: bool = False
+    started_at: Optional[str] = None
+    ended_at: Optional[str] = None
+    last_error: str = ""
+    last_stats: Optional[Dict[str, Any]] = None
+    current_log_file: str = ""
+
+
+class _UiLogHandler(logging.Handler):
+    def __init__(self, callback):
+        super().__init__(level=logging.INFO)
+        self._callback = callback
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            self._callback(self.format(record) + "\n")
+        except Exception:
+            return
+
+
+class PipelineRunnerService:
+    def __init__(self) -> None:
+        self._status = RunnerStatus(last_stats={})
+        self._config = default_ui_config()
+        self._task: Optional[asyncio.Task] = None
+        self._log_chunks: Deque[str] = deque(maxlen=5000)
+        self._lock = asyncio.Lock()
+
+    def get_default_config(self) -> PipelineUiConfig:
+        return default_ui_config()
+
+    def get_config(self) -> PipelineUiConfig:
+        return self._config
+
+    def set_config(self, cfg: PipelineUiConfig) -> PipelineUiConfig:
+        self._config = cfg
+        return self._config
+
+    def get_status(self) -> RunnerStatus:
+        return self._status
+
+    def get_logs(self, since: int = 0) -> Tuple[List[str], int]:
+        chunks = list(self._log_chunks)
+        if since < 0:
+            since = 0
+        if since >= len(chunks):
+            return [], len(chunks)
+        return chunks[since:], len(chunks)
+
+    async def start(self, cfg: Optional[PipelineUiConfig] = None) -> RunnerStatus:
+        async with self._lock:
+            if self._task and not self._task.done():
+                raise RuntimeError("Pipeline is already running")
+
+            if cfg is not None:
+                self._config = cfg
+
+            self._log_chunks.clear()
+            self._status = RunnerStatus(
+                running=True,
+                started_at=datetime.now().isoformat(timespec="seconds"),
+                ended_at=None,
+                last_error="",
+                last_stats={},
+                current_log_file="",
+            )
+            self._task = asyncio.create_task(self._run(self._config))
+            return self._status
+
+    async def stop(self) -> RunnerStatus:
+        async with self._lock:
+            task = self._task
+            if not task or task.done():
+                return self._status
+            task.cancel()
+
+        try:
+            await task
+        except asyncio.CancelledError:
+            pass
+        return self._status
+
+    def _append_log(self, data: str) -> int:
+        self._log_chunks.append(data)
+        return len(data)
+
+    def _append_log_callback(self, data: str) -> None:
+        self._append_log(data)
+
+    def write(self, data: str) -> int:
+        return self._append_log(data)
+
+    def flush(self) -> None:
+        return None
+
+    async def _run(self, cfg: PipelineUiConfig) -> None:
+        pipeline = None
+        ui_handler = None
+        app_logger = None
+
+        log_cfg = cfg.logging.model_copy(deep=True)
+        log_path = resolve_log_file_path(log_cfg)
+        self._status.current_log_file = str(log_path or "")
+
+        try:
+            app_logger = init_pipeline_logger(
+                log_file_path=log_path,
+                mirror_console=True,
+                replace_handlers=True,
+            )
+            ui_handler = _UiLogHandler(self._append_log_callback)
+            ui_handler.setFormatter(logging.Formatter("%(message)s"))
+            app_logger.addHandler(ui_handler)
+
+            if cfg.mode == "jsonl_to_api":
+                pipeline = await self._create_jsonl_to_api_pipeline(cfg)
+            else:
+                pipeline = await self._create_jsonl_to_jsonl_pipeline(cfg)
+
+            stats = await pipeline.run()
+            self._status.last_stats = stats
+            self._status.last_error = ""
+        except asyncio.CancelledError:
+            self._status.last_error = "cancelled"
+            raise
+        except Exception as exc:
+            self._status.last_error = f"{type(exc).__name__}: {exc}"
+        finally:
+            try:
+                if pipeline is not None:
+                    await pipeline.close()
+            finally:
+                if app_logger is not None and ui_handler is not None:
+                    app_logger.removeHandler(ui_handler)
+                    try:
+                        ui_handler.close()
+                    except Exception:
+                        pass
+                clear_pipeline_logger_handlers(app_logger)
+                self._status.running = False
+                self._status.ended_at = datetime.now().isoformat(timespec="seconds")
+
+    async def _create_jsonl_to_api_pipeline(self, cfg: PipelineUiConfig) -> FullSyncPipeline:
+        local_ds = cfg.local_datasource.model_copy(deep=True)
+        remote_ds = cfg.remote_datasource.model_copy(deep=True)
+        local_ds = JsonlDataSourceConfig.model_validate(local_ds.model_dump())
+        remote_ds = ApiDataSourceConfig.model_validate(remote_ds.model_dump())
+
+        run_cfg = PipelineRunConfig(
+            node_types=list(cfg.node_types),
+            target_project_ids=list(cfg.target_project_ids),
+            local_datasource=local_ds,
+            remote_datasource=remote_ds,
+            strategies=[],
+            persist=cfg.persist.model_copy(deep=True),
+            logging=cfg.logging.model_copy(deep=True),
+        )
+        if cfg.strategies:
+            run_cfg = apply_config_overrides(run_cfg, {"strategies": dict(cfg.strategies)})
+        return await create_pipeline_from_config(run_cfg)
+
+    async def _create_jsonl_to_jsonl_pipeline(self, cfg: PipelineUiConfig) -> FullSyncPipeline:
+        local_ds = cfg.local_datasource.model_copy(deep=True)
+        remote_ds = cfg.remote_datasource.model_copy(deep=True)
+        local_ds = JsonlDataSourceConfig.model_validate(local_ds.model_dump())
+        remote_ds = JsonlDataSourceConfig.model_validate(remote_ds.model_dump())
+
+        run_cfg = PipelineRunConfig(
+            node_types=list(cfg.node_types),
+            target_project_ids=list(cfg.target_project_ids),
+            local_datasource=local_ds,
+            remote_datasource=remote_ds,
+            strategies=[],
+            persist=cfg.persist.model_copy(deep=True),
+            logging=cfg.logging.model_copy(deep=True),
+        )
+        if cfg.strategies:
+            run_cfg = apply_config_overrides(run_cfg, {"strategies": dict(cfg.strategies)})
+        return await create_pipeline_from_config(run_cfg)