Flowfile 0.5.6__py3-none-any.whl → 0.6.1__py3-none-any.whl

flowfile_worker/routes.py

@@ -1,9 +1,8 @@
 import os
 import uuid
-from base64 import encodebytes
 
 import polars as pl
-from fastapi import APIRouter, BackgroundTasks, HTTPException, Response
+from fastapi import APIRouter, BackgroundTasks, HTTPException, Request, Response
 
 from flowfile_worker import CACHE_DIR, PROCESS_MEMORY_USAGE, models, status_dict, status_dict_lock
 from flowfile_worker.configs import logger
@@ -23,32 +22,45 @@ def create_and_get_default_cache_dir(flowfile_flow_id: int) -> str:
 
 
 @router.post("/submit_query/")
-def submit_query(polars_script: models.PolarsScript, background_tasks: BackgroundTasks) -> models.Status:
-    logger.info(f"Processing query with operation: {polars_script.operation_type}")
-
+async def submit_query(request: Request, background_tasks: BackgroundTasks) -> models.Status:
+    """Accept raw binary data with metadata in headers for efficient transfer."""
     try:
-        polars_script.task_id = str(uuid.uuid4()) if polars_script.task_id is None else polars_script.task_id
-        default_cache_dir = create_and_get_default_cache_dir(polars_script.flowfile_flow_id)
+        # Read raw bytes directly from request body - no base64 decoding needed
+        polars_serializable_object = await request.body()
+
+        # Get metadata from headers
+        task_id = request.headers.get("X-Task-Id") or str(uuid.uuid4())
+        operation_type = request.headers.get("X-Operation-Type", "store")
+        flow_id = int(request.headers.get("X-Flow-Id", "1"))
+        node_id = request.headers.get("X-Node-Id", "-1")
+        # Try to parse node_id as int, fall back to string
+        try:
+            node_id = int(node_id)
+        except ValueError:
+            pass
+
+        logger.info(f"Processing query with operation: {operation_type}")
+
+        default_cache_dir = create_and_get_default_cache_dir(flow_id)
+        file_path = os.path.join(default_cache_dir, f"{task_id}.arrow")
+        result_type = "polars" if operation_type == "store" else "other"
 
-        polars_script.cache_dir = polars_script.cache_dir if polars_script.cache_dir is not None else default_cache_dir
-        polars_serializable_object = polars_script.polars_serializable_object()
-        file_path = os.path.join(polars_script.cache_dir, f"{polars_script.task_id}.arrow")
-        result_type = "polars" if polars_script.operation_type == "store" else "other"
         status = models.Status(
-            background_task_id=polars_script.task_id, status="Starting", file_ref=file_path, result_type=result_type
+            background_task_id=task_id, status="Starting", file_ref=file_path, result_type=result_type
         )
-        status_dict[polars_script.task_id] = status
+        status_dict[task_id] = status
+
         background_tasks.add_task(
             start_process,
             polars_serializable_object=polars_serializable_object,
-            task_id=polars_script.task_id,
-            operation=polars_script.operation_type,
+            task_id=task_id,
+            operation=operation_type,
             file_ref=file_path,
-            flowfile_flow_id=polars_script.flowfile_flow_id,
-            flowfile_node_id=polars_script.flowfile_node_id,
+            flowfile_flow_id=flow_id,
+            flowfile_node_id=node_id,
             kwargs={},
         )
-        logger.info(f"Started background task: {polars_script.task_id}")
+        logger.info(f"Started background task: {task_id}")
         return status
 
     except Exception as e:
@@ -57,32 +69,44 @@ def submit_query(polars_script: models.PolarsScript, background_tasks: Backgroun
 
 
 @router.post("/store_sample/")
-def store_sample(polars_script: models.PolarsScriptSample, background_tasks: BackgroundTasks) -> models.Status:
-    logger.info(f"Processing sample storage with size: {polars_script.sample_size}")
-
+async def store_sample(request: Request, background_tasks: BackgroundTasks) -> models.Status:
+    """Accept raw binary data with metadata in headers for efficient transfer."""
     try:
-        default_cache_dir = create_and_get_default_cache_dir(polars_script.flowfile_flow_id)
-        polars_script.task_id = str(uuid.uuid4()) if polars_script.task_id is None else polars_script.task_id
-        polars_script.cache_dir = polars_script.cache_dir if polars_script.cache_dir is not None else default_cache_dir
-        polars_serializable_object = polars_script.polars_serializable_object()
+        # Read raw bytes directly from request body - no base64 decoding needed
+        polars_serializable_object = await request.body()
+
+        # Get metadata from headers
+        task_id = request.headers.get("X-Task-Id") or str(uuid.uuid4())
+        sample_size = int(request.headers.get("X-Sample-Size", "100"))
+        flow_id = int(request.headers.get("X-Flow-Id", "1"))
+        node_id = request.headers.get("X-Node-Id", "-1")
+        # Try to parse node_id as int, fall back to string
+        try:
+            node_id = int(node_id)
+        except ValueError:
+            pass
+
+        logger.info(f"Processing sample storage with size: {sample_size}")
+
+        default_cache_dir = create_and_get_default_cache_dir(flow_id)
+        file_path = os.path.join(default_cache_dir, f"{task_id}.arrow")
 
-        file_path = os.path.join(polars_script.cache_dir, f"{polars_script.task_id}.arrow")
         status = models.Status(
-            background_task_id=polars_script.task_id, status="Starting", file_ref=file_path, result_type="other"
+            background_task_id=task_id, status="Starting", file_ref=file_path, result_type="other"
         )
-        status_dict[polars_script.task_id] = status
+        status_dict[task_id] = status
 
         background_tasks.add_task(
             start_process,
             polars_serializable_object=polars_serializable_object,
-            task_id=polars_script.task_id,
-            operation=polars_script.operation_type,
+            task_id=task_id,
+            operation="store_sample",
             file_ref=file_path,
-            flowfile_flow_id=polars_script.flowfile_flow_id,
-            flowfile_node_id=polars_script.flowfile_node_id,
-            kwargs={"sample_size": polars_script.sample_size},
+            flowfile_flow_id=flow_id,
+            flowfile_node_id=node_id,
+            kwargs={"sample_size": sample_size},
         )
-        logger.info(f"Started sample storage task: {polars_script.task_id}")
+        logger.info(f"Started sample storage task: {task_id}")
 
         return status
 
@@ -373,7 +397,8 @@ async def fetch_results(task_id: str):
         raise HTTPException(status_code=404, detail=f"An error occurred during processing: {status.error_message}")
     try:
         lf = pl.scan_parquet(status.file_ref)
-        return {"task_id": task_id, "result": encodebytes(lf.serialize()).decode()}
+        # Return raw bytes - Pydantic/FastAPI will handle serialization
+        return {"task_id": task_id, "result": lf.serialize()}
     except Exception as e:
         logger.error(f"Error reading results: {str(e)}")
         raise HTTPException(status_code=500, detail="Error reading results")

flowfile_worker/spawner.py

@@ -1,4 +1,5 @@
 import gc
+from base64 import b64encode
 from multiprocessing import Process
 from multiprocessing.queues import Queue
 from time import sleep
@@ -62,7 +63,12 @@ def handle_task(task_id: str, p: Process, progress: mp_context.Value, error_mess
                 if final_progress == 100:
                     status.status = "Completed"
                     if not q.empty():
-                        status.results = q.get()
+                        result = q.get()
+                        # b64-encode bytes for JSON-safe storage in status_dict (REST responses)
+                        if isinstance(result, bytes):
+                            status.results = b64encode(result).decode("ascii")
+                        else:
+                            status.results = result
                 elif final_progress != -1:
                     status_dict[task_id].status = "Unknown Error"
 

flowfile_worker/streaming.py

@@ -0,0 +1,335 @@
+"""
+WebSocket streaming endpoint for worker-core communication.
+
+Replaces the HTTP poll-based pattern with a single WebSocket connection per task:
+1. Core sends JSON metadata + binary payload
+2. Worker streams progress updates as JSON
+3. Worker sends result as binary frame (no base64 encoding)
+
+This eliminates:
+- HTTP polling latency (0.5s+ per poll cycle)
+- Base64 encode/decode overhead on result bytes
+- Multiple HTTP round-trips per task
+"""
+
+import asyncio
+import gc
+import os
+import threading
+import uuid
+from dataclasses import dataclass
+from multiprocessing import Process
+from multiprocessing.queues import Queue
+from typing import Any
+
+from base64 import b64encode
+from fastapi import APIRouter, WebSocket, WebSocketDisconnect
+
+from flowfile_worker import CACHE_DIR, funcs, models, mp_context, status_dict, status_dict_lock
+from flowfile_worker.configs import logger
+from flowfile_worker.spawner import process_manager
+
+streaming_router = APIRouter()
+
+# Maps operation type to result type for status tracking
+_POLARS_RESULT_OPERATIONS = frozenset({"store"})
+
+
+def _get_result_type(operation: str) -> str:
+    return "polars" if operation in _POLARS_RESULT_OPERATIONS else "other"
+
+
+# ---------------------------------------------------------------------------
+# Task context
+# ---------------------------------------------------------------------------
+
+@dataclass
+class _TaskContext:
+    """Parsed and resolved metadata for a WebSocket task."""
+    task_id: str
+    operation: str
+    flow_id: int
+    node_id: int | str
+    extra_kwargs: dict
+    file_path: str
+    result_type: str
+
+
+def _parse_metadata(metadata: dict) -> _TaskContext:
+    """Parse raw WebSocket metadata into a resolved TaskContext."""
+    task_id = metadata.get("task_id") or str(uuid.uuid4())
+    operation = metadata.get("operation", "store")
+    flow_id = int(metadata.get("flow_id", 1))
+    node_id = metadata.get("node_id", -1)
+    extra_kwargs = metadata.get("kwargs", {})
+
+    try:
+        node_id = int(node_id)
+    except (ValueError, TypeError):
+        pass
+
+    # Set up cache directory and file path
+    cache_dir = CACHE_DIR / str(flow_id)
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    file_path = os.path.join(str(cache_dir), f"{task_id}.arrow")
+    result_type = _get_result_type(operation)
+
+    return _TaskContext(
+        task_id=task_id,
+        operation=operation,
+        flow_id=flow_id,
+        node_id=node_id,
+        extra_kwargs=extra_kwargs,
+        file_path=file_path,
+        result_type=result_type,
+    )
+
+
+def _register_status(ctx: _TaskContext) -> None:
+    """Register the task in status_dict for REST compatibility."""
+    status_dict[ctx.task_id] = models.Status(
+        background_task_id=ctx.task_id,
+        status="Starting",
+        file_ref=ctx.file_path,
+        result_type=ctx.result_type,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Subprocess management
+# ---------------------------------------------------------------------------
+
+def _spawn_subprocess(ctx: _TaskContext, polars_bytes: bytes) -> tuple[Process, Any, Any, Queue]:
+    """Spawn a worker subprocess and return (process, progress, error_message, queue)."""
+    process_task = getattr(funcs, ctx.operation)
+
+    kwargs = dict(ctx.extra_kwargs)
+    kwargs["polars_serializable_object"] = polars_bytes
+
+    progress = mp_context.Value("i", 0)
+    error_message = mp_context.Array("c", 1024)
+    queue = mp_context.Queue(maxsize=1)
+
+    kwargs["progress"] = progress
+    kwargs["error_message"] = error_message
+    kwargs["queue"] = queue
+    kwargs["file_path"] = ctx.file_path
+    kwargs["flowfile_flow_id"] = ctx.flow_id
+    kwargs["flowfile_node_id"] = ctx.node_id
+
+    p = mp_context.Process(target=process_task, kwargs=kwargs)
+    p.start()
+    process_manager.add_process(ctx.task_id, p)
+
+    with status_dict_lock:
+        status_dict[ctx.task_id].status = "Processing"
+
+    logger.info(f"[WS] Started task {ctx.task_id} with operation: {ctx.operation}")
+    return p, progress, error_message, queue
+
+
+def _read_error_message(error_message) -> str:
+    """Extract error string from shared ctypes array."""
+    with error_message.get_lock():
+        return error_message.value.decode().rstrip("\x00")
+
+
+def _set_error_status(task_id: str, msg: str) -> None:
+    """Update status_dict to reflect an error."""
+    with status_dict_lock:
+        status_dict[task_id].status = "Error"
+        status_dict[task_id].error_message = msg
+
+
+# ---------------------------------------------------------------------------
+# Progress monitoring
+# ---------------------------------------------------------------------------
+
+async def _monitor_progress(websocket: WebSocket, p: Process, progress, error_message, task_id: str) -> bool:
+    """Stream progress updates while subprocess is alive.
+
+    Returns True if an error was detected and sent to the client.
+    """
+    last_progress = -1
+
+    while p.is_alive():
+        await asyncio.sleep(0.3)
+
+        with progress.get_lock():
+            current = progress.value
+
+        if current != last_progress:
+            try:
+                await websocket.send_json({"type": "progress", "progress": current})
+            except Exception:
+                return False
+            last_progress = current
+
+        if current == -1:
+            msg = _read_error_message(error_message)
+            _set_error_status(task_id, msg)
+            await websocket.send_json({"type": "error", "error_message": msg})
+            return True
+
+    return False
+
+
+# ---------------------------------------------------------------------------
+# Result handling
+# ---------------------------------------------------------------------------
+
+def _update_completed_status(task_id: str, result_data: Any) -> None:
+    """Update status_dict for a successfully completed task."""
+    with status_dict_lock:
+        status_dict[task_id].status = "Completed"
+        status_dict[task_id].progress = 100
+        if result_data is not None:
+            if isinstance(result_data, bytes):
+                status_dict[task_id].results = b64encode(result_data).decode("ascii")
+            else:
+                status_dict[task_id].results = result_data
+
+
+async def _send_completion(websocket: WebSocket, task_id: str, result_type: str,
+                           file_path: str, queue: Queue) -> None:
+    """Send completion message and result data over WebSocket."""
+    result_data = queue.get() if not queue.empty() else None
+    _update_completed_status(task_id, result_data)
+
+    has_result = result_data is not None
+    await websocket.send_json({
+        "type": "complete",
+        "result_type": result_type,
+        "file_ref": file_path,
+        "has_result": has_result,
+    })
+
+    if has_result:
+        if isinstance(result_data, bytes):
+            await websocket.send_bytes(result_data)
+        else:
+            await websocket.send_json({"type": "result_data", "data": result_data})
+
+    logger.info(f"[WS] Task {task_id} completed successfully")
+
+
+async def _send_final_error(websocket: WebSocket, task_id: str, progress, error_message) -> None:
+    """Handle error or unexpected termination after subprocess exits."""
+    with progress.get_lock():
+        final = progress.value
+
+    if final == -1:
+        msg = _read_error_message(error_message)
+        _set_error_status(task_id, msg)
+        await websocket.send_json({"type": "error", "error_message": msg})
+    else:
+        with status_dict_lock:
+            status_dict[task_id].status = "Unknown Error"
+        await websocket.send_json({
+            "type": "error",
+            "error_message": "Process ended unexpectedly",
+        })
+
+
+# ---------------------------------------------------------------------------
+# Disconnect handling
+# ---------------------------------------------------------------------------
+
+def _handoff_to_background(task_id: str, p: Process, progress, error_message, queue: Queue) -> None:
+    """Hand off a running subprocess to a background thread for REST status updates."""
+    from flowfile_worker.spawner import handle_task as _handle_task
+
+    threading.Thread(
+        target=_handle_task,
+        args=(task_id, p, progress, error_message, queue),
+        daemon=True,
+    ).start()
+
+
+# ---------------------------------------------------------------------------
+# Cleanup
+# ---------------------------------------------------------------------------
+
+def _cleanup_process(task_id: str | None, p: Process | None) -> None:
+    """Clean up subprocess resources."""
+    if p is not None:
+        if p.is_alive():
+            p.join(timeout=1)
+            if p.is_alive():
+                p.terminate()
+                p.join()
+        process_manager.remove_process(task_id)
+
+
+# ---------------------------------------------------------------------------
+# WebSocket endpoint
+# ---------------------------------------------------------------------------
+
+@streaming_router.websocket("/ws/submit")
+async def ws_submit(websocket: WebSocket):
+    """WebSocket endpoint for streaming task submission and result retrieval.
+
+    Protocol (Core -> Worker):
+        1. JSON message: task metadata (task_id, operation, flow_id, node_id, kwargs)
+        2. Binary message: serialized Polars LazyFrame bytes
+
+    Protocol (Worker -> Core):
+        - JSON: {"type": "progress", "progress": N}  (0-100, sent periodically)
+        - JSON: {"type": "complete", "result_type": "polars"|"other", "file_ref": "...", "has_result": bool}
+        - Binary: raw result bytes (only if has_result=True and result_type="polars")
+        - JSON: {"type": "result_data", "data": ...} (only if has_result=True and result_type="other")
+        - JSON: {"type": "error", "error_message": "..."}
+    """
+    await websocket.accept()
+    p = None
+    task_id = None
+    progress = None
+    error_message = None
+    queue = None
+
+    try:
+        # 1. Receive metadata + binary payload
+        metadata = await websocket.receive_json()
+        ctx = _parse_metadata(metadata)
+        task_id = ctx.task_id
+        _register_status(ctx)
+
+        polars_bytes = await websocket.receive_bytes()
+
+        # 2. Spawn subprocess
+        p, progress, error_message, queue = _spawn_subprocess(ctx, polars_bytes)
+
+        # 3. Monitor progress (returns True if error was sent)
+        had_error = await _monitor_progress(websocket, p, progress, error_message, task_id)
+        if had_error:
+            return
+
+        p.join()
+
+        # 4. Send result or error
+        with progress.get_lock():
+            final = progress.value
+
+        if final == 100:
+            await _send_completion(websocket, task_id, ctx.result_type, ctx.file_path, queue)
+        else:
+            await _send_final_error(websocket, task_id, progress, error_message)
+
+    except WebSocketDisconnect:
+        logger.warning(f"[WS] Client disconnected for task {task_id}")
+        if p is not None and p.is_alive() and queue is not None:
+            _handoff_to_background(task_id, p, progress, error_message, queue)
+            # Prevent finally block from cleaning up - handle_task owns these now
+            p = None
+            progress = None
+            error_message = None
+    except Exception as e:
+        logger.error(f"[WS] Error for task {task_id}: {e}", exc_info=True)
+        try:
+            await websocket.send_json({"type": "error", "error_message": str(e)})
+        except Exception:
+            pass
+    finally:
+        _cleanup_process(task_id, p)
+        del p, progress, error_message
+        gc.collect()

flowfile/web/static/assets/ContextMenu-366bf1b4.js

@@ -1,9 +0,0 @@
-import "./index-ca6799de.js";
-import "./DesignerView-a4466dab.js";
-import { _ as _sfc_main } from "./ContextMenu.vue_vue_type_script_setup_true_lang-774c517c.js";
-import "./PopOver-ddcfe4f6.js";
-import "./index-057d770d.js";
-import "./vue-codemirror.esm-8f46fb36.js";
-export {
-  _sfc_main as default
-};

flowfile/web/static/assets/ContextMenu-85cf5b44.js

@@ -1,9 +0,0 @@
-import "./index-ca6799de.js";
-import "./DesignerView-a4466dab.js";
-import { _ as _sfc_main } from "./ContextMenu.vue_vue_type_script_setup_true_lang-774c517c.js";
-import "./PopOver-ddcfe4f6.js";
-import "./index-057d770d.js";
-import "./vue-codemirror.esm-8f46fb36.js";
-export {
-  _sfc_main as default
-};

flowfile/web/static/assets/ContextMenu-9d28ae6d.js

@@ -1,9 +0,0 @@
-import "./index-ca6799de.js";
-import "./DesignerView-a4466dab.js";
-import { _ as _sfc_main } from "./ContextMenu.vue_vue_type_script_setup_true_lang-774c517c.js";
-import "./PopOver-ddcfe4f6.js";
-import "./index-057d770d.js";
-import "./vue-codemirror.esm-8f46fb36.js";
-export {
-  _sfc_main as default
-};

flowfile/web/static/assets/Join-28b5e18f.css

@@ -1,109 +0,0 @@
-
-/* Join Type Selector */
-.join-type-selector[data-v-8435fe68] {
-  display: flex;
-  align-items: center;
-  margin: 12px;
-  gap: 10px;
-}
-.join-type-label[data-v-8435fe68] {
-  font-size: 12px;
-  color: var(--color-text-primary);
-  font-weight: 500;
-  min-width: 70px;
-}
-
-/* Join Mapping Section */
-.table-wrapper[data-v-8435fe68] {
-  border: 1px solid #eee;
-  border-radius: 6px;
-  overflow: hidden;
-  margin: 5px;
-}
-.selectors-header[data-v-8435fe68] {
-  display: flex;
-  justify-content: space-between;
-  padding: 8px 16px;
-  background-color: #fafafa;
-  border-bottom: 1px solid #eee;
-}
-.selectors-title[data-v-8435fe68] {
-  flex: 1;
-  text-align: center;
-  font-size: 12px;
-  color: #666;
-  font-weight: 500;
-}
-.selectors-container[data-v-8435fe68] {
-  padding: 12px;
-  box-sizing: border-box;
-  width: 100%;
-  display: flex;
-  justify-content: space-between;
-  flex-direction: column;
-}
-.selectors-row[data-v-8435fe68] {
-  display: flex;
-  gap: 12px;
-  margin-bottom: 8px;
-  width: 100%;
-  display: flex;
-  justify-content: space-between;
-}
-.selectors-row[data-v-8435fe68]:last-child {
-  margin-bottom: 0;
-}
-
-/* Action Buttons */
-.action-buttons[data-v-8435fe68] {
-  display: flex;
-  gap: 4px;
-  min-width: 60px;
-  justify-content: center;
-}
-.action-button[data-v-8435fe68],
-.add-join-button[data-v-8435fe68],
-.remove-join-button[data-v-8435fe68] {
-  cursor: pointer;
-  width: 24px;
-  height: 24px;
-  border-radius: 4px;
-  border: 1px solid #ddd;
-  background-color: var(--color-background-primary);
-  display: flex;
-  align-items: center;
-  justify-content: center;
-  font-size: 14px;
-  transition: all 0.2s ease;
-}
-.add-join-button[data-v-8435fe68] {
-  color: #45a049;
-  border-color: #45a049;
-}
-.add-join-button[data-v-8435fe68]:hover {
-  background-color: #45a049;
-  color: #fff;
-}
-.remove-join-button[data-v-8435fe68] {
-  color: #d32f2f;
-  border-color: #d32f2f;
-}
-.remove-join-button[data-v-8435fe68]:hover {
-  background-color: #d32f2f;
-  color: #fff;
-}
-
-/* Custom scrollbar */
-.selectors-container[data-v-8435fe68]::-webkit-scrollbar {
-  width: 8px;
-}
-.selectors-container[data-v-8435fe68]::-webkit-scrollbar-track {
-  background: transparent;
-}
-.selectors-container[data-v-8435fe68]::-webkit-scrollbar-thumb {
-  background-color: rgba(0, 0, 0, 0.1);
-  border-radius: 4px;
-}
-.selectors-container[data-v-8435fe68]::-webkit-scrollbar-thumb:hover {
-  background-color: rgba(0, 0, 0, 0.2);
-}