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

flowfile_core/flowfile/flow_data_engine/subprocess_operations/streaming.py

@@ -0,0 +1,206 @@
+"""
+WebSocket streaming client for core-to-worker communication.
+
+Replaces the HTTP poll-based pattern with a single WebSocket connection:
+- Sends task metadata + serialized LazyFrame as binary
+- Receives progress updates as JSON pushes
+- Receives result as raw binary frame (no base64 encoding)
+
+Falls back to REST automatically if the worker doesn't support WebSocket.
+"""
+
+import io
+import json
+from base64 import b64encode
+from typing import Any
+
+import polars as pl
+from websockets.sync.client import connect
+
+from flowfile_core.configs.settings import WORKER_URL
+from flowfile_core.flowfile.flow_data_engine.subprocess_operations.models import Status
+
+
+def _get_ws_url() -> str:
+    """Convert HTTP worker URL to WebSocket URL."""
+    return WORKER_URL.replace("http://", "ws://").replace("https://", "wss://")
+
+
+# ---------------------------------------------------------------------------
+# Message building
+# ---------------------------------------------------------------------------
+
+def _build_metadata(
+    task_id: str,
+    operation_type: str,
+    flow_id: int,
+    node_id: int | str,
+    kwargs: dict | None,
+) -> dict:
+    """Build the JSON metadata message for the WebSocket protocol."""
+    metadata = {
+        "task_id": task_id,
+        "operation": operation_type,
+        "flow_id": flow_id,
+        "node_id": node_id,
+    }
+    if kwargs:
+        metadata["kwargs"] = kwargs
+    return metadata
+
+
+# ---------------------------------------------------------------------------
+# Message receiving
+# ---------------------------------------------------------------------------
+
+def _handle_complete_message(data: dict, task_id: str) -> Status:
+    """Build a partial Status from a 'complete' protocol message.
+
+    The ``results`` field is set to None here and populated by the caller
+    once the actual result payload has been received.
+    """
+    return Status(
+        background_task_id=task_id,
+        status="Completed",
+        file_ref=data.get("file_ref", ""),
+        result_type=data.get("result_type", "polars"),
+        progress=100,
+        results=data.get("results", None),
+    )
+
+
+def _receive_raw_result(ws, task_id: str) -> tuple[Any, Status | None]:
+    """Receive messages from the worker until a raw result or error arrives.
+
+    Returns the result **without** deserializing it so the caller can both
+    populate ``Status.results`` (b64-encoded bytes for polars) and
+    deserialize into the in-memory object.
+
+    Returns:
+        (raw_result, status) where raw_result is ``bytes`` for polars
+        results, arbitrary data for "other" results, or ``None``.
+    """
+    raw_result = None
+    status = None
+
+    while True:
+        msg = ws.recv()
+
+        if isinstance(msg, bytes):
+            raw_result = msg
+            break
+
+        data = json.loads(msg)
+        msg_type = data.get("type")
+
+        if msg_type == "progress":
+            continue
+
+        if msg_type == "complete":
+            status = _handle_complete_message(data, task_id)
+            if not data.get("has_result", False):
+                break
+            continue
+
+        if msg_type == "result_data":
+            raw_result = data.get("data")
+            break
+
+        if msg_type == "error":
+            raise Exception(data.get("error_message", "Unknown worker error"))
+
+    return raw_result, status
+
+
+def _deserialize_and_populate_status(
+    raw_result: Any, status: Status
+) -> tuple[Any, Status]:
+    """Deserialize the raw result and fill ``status.results``.
+
+    For polars results (bytes): deserializes into a LazyFrame and stores
+    the b64-encoded bytes in ``status.results`` (matching REST behaviour).
+    For other results: stores the value directly in ``status.results``.
+    """
+    if raw_result is None:
+        return None, status
+
+    if isinstance(raw_result, bytes):
+        status.results = b64encode(raw_result).decode("ascii")
+        return pl.LazyFrame.deserialize(io.BytesIO(raw_result)), status
+
+    status.results = raw_result
+    return raw_result, status
+
+
+def streaming_start(
+    task_id: str,
+    operation_type: str,
+    flow_id: int,
+    node_id: int | str,
+    lf_bytes: bytes,
+    kwargs: dict | None = None,
+):
+    """Open a WebSocket connection and send the task.
+
+    Returns the **open** connection.  The caller must eventually call
+    :func:`streaming_receive` (which closes the connection) or close it
+    manually.
+
+    Raises immediately on connection failure or send error.
+    """
+    ws_url = _get_ws_url() + "/ws/submit"
+    metadata = _build_metadata(task_id, operation_type, flow_id, node_id, kwargs)
+
+    ws = connect(ws_url)
+    try:
+        ws.send(json.dumps(metadata))
+        ws.send(lf_bytes)
+    except Exception:
+        ws.close()
+        raise
+    return ws
+
+
+def streaming_receive(ws, task_id: str) -> tuple[Any, Status]:
+    """Block until the worker sends back a result, then close the connection.
+
+    The returned ``Status`` object is fully populated (including
+    ``results``) so it is equivalent to what the REST polling path
+    would produce.
+
+    Returns:
+        Tuple of (result, Status)
+    """
+    try:
+        raw_result, status = _receive_raw_result(ws, task_id)
+    finally:
+        ws.close()
+
+    if status is None:
+        status = Status(
+            background_task_id=task_id,
+            status="Completed",
+            file_ref="",
+            result_type="polars",
+            progress=100,
+            results=None,
+        )
+
+    return _deserialize_and_populate_status(raw_result, status)
+
+
+def streaming_submit(
+    task_id: str,
+    operation_type: str,
+    flow_id: int,
+    node_id: int | str,
+    lf_bytes: bytes,
+    kwargs: dict | None = None,
+) -> tuple[Any, Status]:
+    """Submit a task via WebSocket and block until the result arrives.
+
+    Convenience wrapper around :func:`streaming_start` +
+    :func:`streaming_receive`.
+    """
+    ws = streaming_start(task_id, operation_type, flow_id, node_id, lf_bytes, kwargs)
+    return streaming_receive(ws, task_id)

flowfile_core/flowfile/flow_data_engine/subprocess_operations/subprocess_operations.py

@@ -1,7 +1,7 @@
 # Standard library imports
 import io
 import threading
-from base64 import decodebytes, encodebytes
+from base64 import b64decode
 from time import sleep
 from typing import Any, Literal
 from uuid import uuid4
@@ -18,6 +18,11 @@ from flowfile_core.flowfile.flow_data_engine.subprocess_operations.models import
     PolarsOperation,
     Status,
 )
+from flowfile_core.flowfile.flow_data_engine.subprocess_operations.streaming import (
+    streaming_receive,
+    streaming_start,
+    streaming_submit,
+)
 from flowfile_core.flowfile.sources.external_sources.sql_source.models import (
     DatabaseExternalReadSettings,
     DatabaseExternalWriteSettings,
@@ -30,15 +35,15 @@ from flowfile_core.utils.arrow_reader import read
 def trigger_df_operation(
     flow_id: int, node_id: int | str, lf: pl.LazyFrame, file_ref: str, operation_type: OperationType = "store"
 ) -> Status:
-    encoded_operation = encodebytes(lf.serialize()).decode()
-    _json = {
-        "task_id": file_ref,
-        "operation": encoded_operation,
-        "operation_type": operation_type,
-        "flowfile_flow_id": flow_id,
-        "flowfile_node_id": node_id,
+    # Send raw bytes directly - no base64 encoding overhead
+    headers = {
+        "Content-Type": "application/octet-stream",
+        "X-Task-Id": file_ref,
+        "X-Operation-Type": operation_type,
+        "X-Flow-Id": str(flow_id),
+        "X-Node-Id": str(node_id),
     }
-    v = requests.post(url=f"{WORKER_URL}/submit_query/", json=_json)
+    v = requests.post(url=f"{WORKER_URL}/submit_query/", data=lf.serialize(), headers=headers)
     if not v.ok:
         raise Exception(f"trigger_df_operation: Could not cache the data, {v.text}")
     return Status(**v.json())
@@ -47,16 +52,16 @@ def trigger_df_operation(
 def trigger_sample_operation(
     lf: pl.LazyFrame, file_ref: str, flow_id: int, node_id: str | int, sample_size: int = 100
 ) -> Status:
-    encoded_operation = encodebytes(lf.serialize()).decode()
-    _json = {
-        "task_id": file_ref,
-        "operation": encoded_operation,
-        "operation_type": "store_sample",
-        "sample_size": sample_size,
-        "flowfile_flow_id": flow_id,
-        "flowfile_node_id": node_id,
+    # Send raw bytes directly - no base64 encoding overhead
+    headers = {
+        "Content-Type": "application/octet-stream",
+        "X-Task-Id": file_ref,
+        "X-Operation-Type": "store_sample",
+        "X-Sample-Size": str(sample_size),
+        "X-Flow-Id": str(flow_id),
+        "X-Node-Id": str(node_id),
     }
-    v = requests.post(url=f"{WORKER_URL}/store_sample/", json=_json)
+    v = requests.post(url=f"{WORKER_URL}/store_sample/", data=lf.serialize(), headers=headers)
     if not v.ok:
         raise Exception(f"trigger_sample_operation: Could not cache the data, {v.text}")
     return Status(**v.json())
@@ -70,8 +75,9 @@ def trigger_fuzzy_match_operation(
     flow_id: int,
     node_id: int | str,
 ) -> Status:
-    left_serializable_object = PolarsOperation(operation=encodebytes(left_df.serialize()))
-    right_serializable_object = PolarsOperation(operation=encodebytes(right_df.serialize()))
+    # Use raw bytes - Pydantic will handle single base64 encoding for JSON transport
+    left_serializable_object = PolarsOperation(operation=left_df.serialize())
+    right_serializable_object = PolarsOperation(operation=right_df.serialize())
     fuzzy_join_input = FuzzyJoinInput(
         left_df_operation=left_serializable_object,
         right_df_operation=right_serializable_object,
@@ -80,7 +86,6 @@ def trigger_fuzzy_match_operation(
         flowfile_flow_id=flow_id,
         flowfile_node_id=node_id,
     )
-    print("fuzzy join input", fuzzy_join_input)
     v = requests.post(f"{WORKER_URL}/add_fuzzy_join", data=fuzzy_join_input.model_dump_json())
     if not v.ok:
         raise Exception(f"trigger_fuzzy_match_operation: Could not cache the data, {v.text}")
@@ -181,9 +186,9 @@ def clear_task_from_worker(file_ref: str) -> bool:
         return False
 
 
-def get_df_result(encoded_df: str) -> pl.LazyFrame:
-    r = decodebytes(encoded_df.encode())
-    return pl.LazyFrame.deserialize(io.BytesIO(r))
+def get_df_result(result_b64: str) -> pl.LazyFrame:
+    # Results are base64-encoded string from JSON response, decode once
+    return pl.LazyFrame.deserialize(io.BytesIO(b64decode(result_b64)))
 
 
 def get_external_df_result(file_ref: str) -> pl.LazyFrame | None:
@@ -240,6 +245,9 @@ class BaseFetcher:
         self._stop_event = threading.Event()
         self._thread = None
 
+        # WebSocket connection for non-blocking streaming mode
+        self._ws = None
+
         # State variables - use properties for thread-safe access
         self._result: Any | None = None
         self._started: bool = False
@@ -390,6 +398,16 @@ class BaseFetcher:
         """
         logger.warning("Cancelling the operation")
 
+        # Close WebSocket if streaming (causes recv thread to exit)
+        with self._lock:
+            ws = self._ws
+            self._ws = None
+        if ws is not None:
+            try:
+                ws.close()
+            except Exception:
+                pass
+
         # Cancel on the worker side
         try:
             cancel_task(self.file_ref)
@@ -451,6 +469,79 @@ class BaseFetcher:
         with self._lock:
             return self._error_code, self._error_description
 
+    def _execute_streaming(
+        self,
+        operation_type: str,
+        flow_id: int,
+        node_id: int | str,
+        lf_bytes: bytes,
+        kwargs: dict | None = None,
+        blocking: bool = True,
+    ) -> None:
+        """Execute via WebSocket streaming - no polling, binary result transfer.
+
+        Args:
+            blocking: If True (default), blocks until the result is available
+                and sets self._result directly.  If False, opens the
+                connection, sends the task, and hands off to a background
+                thread that will set self._result when done.
+
+        Raises on connection or send error so the caller can fall back to REST.
+        """
+        if blocking:
+            result, status = streaming_submit(
+                task_id=self.file_ref,
+                operation_type=operation_type,
+                flow_id=flow_id,
+                node_id=node_id,
+                lf_bytes=lf_bytes,
+                kwargs=kwargs,
+            )
+            with self._lock:
+                self._result = result
+                self._running = False
+                self._started = True
+            self.status = status
+        else:
+            # Non-blocking: open connection, send task, receive in background
+            ws = streaming_start(
+                task_id=self.file_ref,
+                operation_type=operation_type,
+                flow_id=flow_id,
+                node_id=node_id,
+                lf_bytes=lf_bytes,
+                kwargs=kwargs,
+            )
+            with self._lock:
+                self._ws = ws
+                self._running = True
+                self._started = True
+                self._thread = threading.Thread(
+                    target=self._ws_receive_thread,
+                    args=(ws,),
+                    daemon=True,
+                )
+                self._thread.start()
+
+    def _ws_receive_thread(self, ws) -> None:
+        """Background thread that receives results over an open WebSocket."""
+        try:
+            result, status = streaming_receive(ws, self.file_ref)
+            with self._condition:
+                self._result = result
+                self._running = False
+                self._ws = None
+                self._condition.notify_all()
+            self.status = status
+        except Exception as e:
+            logger.exception("Error in WebSocket receive thread")
+            with self._condition:
+                self._error_code = -1
+                self._error_description = str(e)
+                self._running = False
+                self._ws = None
+                self._condition.notify_all()
+
 
 class ExternalDfFetcher(BaseFetcher):
     status: Status | None = None
@@ -467,6 +558,21 @@ class ExternalDfFetcher(BaseFetcher):
     ):
         super().__init__(file_ref=file_ref)
         lf = lf.lazy() if isinstance(lf, pl.DataFrame) else lf
+
+        # Try WebSocket streaming first (blocking or non-blocking)
+        try:
+            self._execute_streaming(
+                operation_type=operation_type,
+                flow_id=flow_id,
+                node_id=node_id,
+                lf_bytes=lf.serialize(),
+                blocking=wait_on_completion,
+            )
+            return
+        except Exception as e:
+            logger.debug(f"WebSocket streaming unavailable ({e}), falling back to REST")
+
+        # REST fallback (original behavior)
         r = trigger_df_operation(
             lf=lf, file_ref=self.file_ref, operation_type=operation_type, node_id=node_id, flow_id=flow_id
         )
@@ -490,6 +596,22 @@ class ExternalSampler(BaseFetcher):
     ):
         super().__init__(file_ref=file_ref)
         lf = lf.lazy() if isinstance(lf, pl.DataFrame) else lf
+
+        # Try WebSocket streaming first (blocking or non-blocking)
+        try:
+            self._execute_streaming(
+                operation_type="store_sample",
+                flow_id=flow_id,
+                node_id=node_id,
+                lf_bytes=lf.serialize(),
+                kwargs={"sample_size": sample_size},
+                blocking=wait_on_completion,
+            )
+            return
+        except Exception as e:
+            logger.debug(f"WebSocket streaming unavailable ({e}), falling back to REST")
+
+        # REST fallback (original behavior)
         r = trigger_sample_operation(
             lf=lf, file_ref=file_ref, sample_size=sample_size, node_id=node_id, flow_id=flow_id
         )