from threading import Thread

from typing import Callable, Optional

import pandas as pd

from feldera import FelderaClient

from feldera._helpers import dataframe_from_response

from feldera.enums import PipelineFieldSelector

class CallbackRunner(Thread):

def __init__(

self,

client: FelderaClient,

pipeline_name: str,

view_name: str,

callback: Callable[[pd.DataFrame, int], None],

):

super().__init__()

self.daemon = True

self.client: FelderaClient = client

self.pipeline_name: str = pipeline_name

self.view_name: str = view_name

self.callback: Callable[[pd.DataFrame, int], None] = callback

self.schema: Optional[dict] = None

def run(self):

"""

The main loop of the thread. Listens for data and calls the callback function on each chunk of data received.

:meta private:

"""

pipeline = self.client.get_pipeline(

self.pipeline_name, PipelineFieldSelector.ALL

)

schemas = pipeline.tables + pipeline.views

for schema in schemas:

if schema.name == self.view_name:

self.schema = schema

break

if self.schema is None:

raise ValueError(

f"Table or View {self.view_name} not found in the pipeline schema."

)

gen_obj = self.client.listen_to_pipeline(

self.pipeline_name,

self.view_name,

format="json",

case_sensitive=self.schema.case_sensitive,

)

iterator = gen_obj()

for chunk in iterator:

chunk: dict = chunk

data: Optional[list[dict]] = chunk.get("json_data")

seq_no: Optional[int] = chunk.get("sequence_number")

if data is not None and seq_no is not None:

self.callback(

dataframe_from_response([data], self.schema.fields), seq_no

)