import time
import pandas
from typing import List, Dict, Callable, Optional, Generator, Mapping, Any
from collections import deque
from queue import Queue
from feldera.rest.errors import FelderaAPIError
from feldera.enums import PipelineStatus
from feldera.rest.pipeline import Pipeline as InnerPipeline
from feldera.rest.feldera_client import FelderaClient
from feldera._callback_runner import _CallbackRunnerInstruction, CallbackRunner
from feldera.output_handler import OutputHandler
from feldera._helpers import ensure_dataframe_has_columns, chunk_dataframe
[docs]
class Pipeline:
def __init__(self, name: str, client: FelderaClient):
self.name = name
self.client: FelderaClient = client
self._inner: InnerPipeline | None = None
self.views_tx: List[Dict[str, Queue]] = []
def __setup_output_listeners(self):
"""
Internal function used to set up the output listeners.
:meta private:
"""
for view_queue in self.views_tx:
for view_name, queue in view_queue.items():
# sends a message to the callback runner to start listening
queue.put(_CallbackRunnerInstruction.PipelineStarted)
# block until the callback runner is ready
queue.join()
[docs]
def status(self) -> PipelineStatus:
"""
Return the current status of the pipeline.
"""
try:
inner = self.client.get_pipeline(self.name)
self._inner = inner
return PipelineStatus.from_str(inner.deployment_status)
except FelderaAPIError as err:
if err.status_code == 404:
return PipelineStatus.NOT_FOUND
else:
raise err
[docs]
def listen(self, view_name: str) -> OutputHandler:
"""
Listen to the output of the provided view so that it is available in the notebook / python code.
When the pipeline is shutdown, these listeners are dropped.
:param view_name: The name of the view to listen to.
"""
queue: Optional[Queue] = None
if self.status() not in [PipelineStatus.PAUSED, PipelineStatus.RUNNING]:
queue = Queue(maxsize=1)
self.views_tx.append({view_name: queue})
handler = OutputHandler(self.client, self.name, view_name, queue)
handler.start()
return handler
[docs]
def foreach_chunk(
self, view_name: str, callback: Callable[[pandas.DataFrame, int], None]
):
"""
Run the given callback on each chunk of the output of the specified view.
:param view_name: The name of the view.
:param callback: The callback to run on each chunk. The callback should take two arguments:
- **chunk** -> The chunk as a pandas DataFrame
- **seq_no** -> The sequence number. The sequence number is a monotonically increasing integer that
starts from 0. Note that the sequence number is unique for each chunk, but not necessarily contiguous.
Please note that the callback is run in a separate thread, so it should be thread-safe.
Please note that the callback should not block for a long time, as by default, backpressure is enabled and
will block the pipeline.
.. note::
- The callback must be thread-safe as it will be run in a separate thread.
"""
queue: Optional[Queue] = None
if self.status() not in [PipelineStatus.RUNNING, PipelineStatus.PAUSED]:
queue = Queue(maxsize=1)
self.views_tx.append({view_name: queue})
handler = CallbackRunner(self.client, self.name, view_name, callback, queue)
handler.start()
[docs]
def wait_for_completion(self, shutdown: bool = False):
"""
Block until the pipeline has completed processing all input records.
This method blocks until (1) all input connectors attached to the pipeline
have finished reading their input data sources and issued end-of-input
notifications to the pipeline, and (2) all inputs received from these
connectors have been fully processed and corresponding outputs have been
sent out through the output connectors.
This method will block indefinitely if at least one of the input
connectors attached to the pipeline is a streaming connector, such as
Kafka, that does not issue the end-of-input notification.
:param shutdown: If True, the pipeline will be shutdown after completion. False by default.
:raises RuntimeError: If the pipeline returns unknown metrics.
"""
if self.status() not in [
PipelineStatus.RUNNING,
PipelineStatus.INITIALIZING,
PipelineStatus.PROVISIONING,
]:
raise RuntimeError("Pipeline must be running to wait for completion")
while True:
metrics: dict = self.client.get_pipeline_stats(self.name).get(
"global_metrics"
)
pipeline_complete: bool = metrics.get("pipeline_complete")
if pipeline_complete is None:
raise RuntimeError(
"received unknown metrics from the pipeline, pipeline_complete is None"
)
if pipeline_complete:
break
time.sleep(1)
if shutdown:
self.shutdown()
[docs]
def start(self):
"""
.. _start:
Starts this pipeline.
:raises RuntimeError: If the pipeline returns unknown metrics.
"""
status = self.status()
if status != PipelineStatus.SHUTDOWN:
raise RuntimeError(
f"pipeline {self.name} in state: {str(status.name)} cannot be started"
)
self.pause()
self.__setup_output_listeners()
self.resume()
[docs]
def restart(self):
"""
Restarts the pipeline.
"""
self.shutdown()
self.start()
[docs]
def wait_for_idle(
self,
idle_interval_s: float = 5.0,
timeout_s: float = 600.0,
poll_interval_s: float = 0.2,
):
"""
Wait for the pipeline to become idle and then returns.
Idle is defined as a sufficiently long interval in which the number of
input and processed records reported by the pipeline do not change, and
they equal each other (thus, all input records present at the pipeline
have been processed).
:param idle_interval_s: Idle interval duration (default is 5.0 seconds).
:param timeout_s: Timeout waiting for idle (default is 600.0 seconds).
:param poll_interval_s: Polling interval, should be set substantially
smaller than the idle interval (default is 0.2 seconds).
:raises ValueError: If idle interval is larger than timeout, poll interval
is larger than timeout, or poll interval is larger than idle interval.
:raises RuntimeError: If the metrics are missing or the timeout was
reached.
"""
if idle_interval_s > timeout_s:
raise ValueError(
f"idle interval ({idle_interval_s}s) cannot be larger than timeout ({timeout_s}s)"
)
if poll_interval_s > timeout_s:
raise ValueError(
f"poll interval ({poll_interval_s}s) cannot be larger than timeout ({timeout_s}s)"
)
if poll_interval_s > idle_interval_s:
raise ValueError(
f"poll interval ({poll_interval_s}s) cannot be larger "
f"than idle interval ({idle_interval_s}s)"
)
start_time_s = time.monotonic()
idle_started_s = None
prev = (0, 0)
while True:
now_s = time.monotonic()
# Metrics retrieval
metrics: dict = self.client.get_pipeline_stats(self.name).get(
"global_metrics"
)
total_input_records: int | None = metrics.get("total_input_records")
total_processed_records: int | None = metrics.get("total_processed_records")
if total_input_records is None:
raise RuntimeError(
"total_input_records is missing from the pipeline metrics"
)
if total_processed_records is None:
raise RuntimeError(
"total_processed_records is missing from the pipeline metrics"
)
# Idle check
unchanged = (
prev[0] == total_input_records and prev[1] == total_processed_records
)
equal = total_input_records == total_processed_records
prev = (total_input_records, total_processed_records)
if unchanged and equal:
if idle_started_s is None:
idle_started_s = now_s
else:
idle_started_s = None
if idle_started_s is not None and now_s - idle_started_s >= idle_interval_s:
return
# Timeout
if now_s - start_time_s >= timeout_s:
raise RuntimeError(f"waiting for idle reached timeout ({timeout_s}s)")
time.sleep(poll_interval_s)
[docs]
def pause(self):
"""
Pause the pipeline.
"""
self.client.pause_pipeline(self.name)
[docs]
def shutdown(self):
"""
Shut down the pipeline.
"""
if len(self.views_tx) > 0:
for _, queue in self.views_tx.pop().items():
# sends a message to the callback runner to stop listening
queue.put(_CallbackRunnerInstruction.RanToCompletion)
# block until the callback runner has been stopped
queue.join()
self.client.shutdown_pipeline(self.name)
[docs]
def resume(self):
"""
Resumes the pipeline.
"""
self.client.start_pipeline(self.name)
[docs]
def delete(self):
"""
Deletes the pipeline.
"""
self.client.delete_pipeline(self.name)
[docs]
@staticmethod
def get(name: str, client: FelderaClient) -> "Pipeline":
"""
Get the pipeline if it exists.
:param name: The name of the pipeline.
:param client: The FelderaClient instance.
"""
try:
inner = client.get_pipeline(name)
pipeline = Pipeline(inner.name, client)
pipeline.__inner = inner
return pipeline
except FelderaAPIError as err:
if err.status_code == 404:
raise RuntimeError(f"Pipeline with name {name} not found")
[docs]
def query(self, query: str) -> Generator[Mapping[str, Any], None, None]:
"""
Executes an ad-hoc SQL query on this pipeline and returns the result in the specified format.
For ``INSERT`` and ``DELETE`` queries, consider using :meth:`.execute` instead.
Important:
This method is lazy. It returns a generator and is not evaluated until you consume the result.
:param query: The SQL query to be executed.
:return: A generator that yields the rows of the result as Python dictionaries.
"""
return self.client.query_as_json(self.name, query)
[docs]
def query_parquet(self, query: str, path: str):
"""
Executes an ad-hoc SQL query on this pipeline and saves the result to the specified path as a parquet file.
If the extension isn't `parquet`, it will be automatically appended to `path`.
:param query: The SQL query to be executed.
:param path: The path of the parquet file.
"""
self.client.query_as_parquet(self.name, query, path)
[docs]
def query_tabular(self, query: str) -> Generator[str, None, None]:
"""
Executes a SQL query on this pipeline and returns the result as a formatted string.
Important:
This method is lazy. It returns a generator and is not evaluated until you consume the result.
:param query: The SQL query to be executed.
:return: A generator that yields a string representing the query result in a human-readable, tabular format.
"""
return self.client.query_as_text(self.name, query)
[docs]
def execute(self, query: str):
"""
Executes an ad-hoc SQL query on the current pipeline, discarding its result.
Unlike the :meth:`.query` method which returns a generator for retrieving query results lazily,
this method processes the query eagerly and fully before returning.
This method is suitable for SQL operations like ``INSERT`` and ``DELETE``, where the user needs
confirmation of successful query execution, but does not require the query result.
If the query fails, an exception will be raised.
:param query: The SQL query to be executed.
"""
gen = self.query_tabular(query)
deque(gen, maxlen=0)