Flowfile 0.2.2__py3-none-any.whl

flowfile_core/flowfile/sources/external_sources/base_class.py

@@ -0,0 +1,39 @@
+from abc import ABC, abstractmethod
+from typing import Generator, Callable, List, Any, Optional, Dict
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn
+import polars as pl
+
+
+class ExternalDataSource(ABC):
+    schema: Optional[List[FlowfileColumn]]
+    data_getter: Optional[Callable]
+    is_collected: bool
+    cache_store: Any
+    _type: str
+    initial_data_getter: Optional[Callable]
+
+    @abstractmethod
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def get_initial_data(self) -> List[Dict[str, Any]]:
+        pass
+
+    @abstractmethod
+    def get_iter(self) -> Generator[Dict[str, Any], None, None]:
+        pass
+
+    @abstractmethod
+    def get_sample(self, n: int = 10000) -> Generator[Dict[str, Any], None, None]:
+        pass
+
+    @abstractmethod
+    def get_pl_df(self) -> pl.DataFrame:
+        pass
+
+    @staticmethod
+    @abstractmethod
+    def parse_schema(*args, **kwargs) -> List[FlowfileColumn]:
+        pass
+

flowfile_core/flowfile/sources/external_sources/custom_external_sources/__init__.py

@@ -0,0 +1,2 @@
+from flowfile_core.flowfile.sources.external_sources.custom_external_sources import google_sheet
+from flowfile_core.flowfile.sources.external_sources.custom_external_sources import sample_users

flowfile_core/flowfile/sources/external_sources/custom_external_sources/external_source.py

@@ -0,0 +1,100 @@
+from typing import Any, Dict, Generator, List, Optional, Callable
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn
+from flowfile_core.schemas import input_schema
+from flowfile_core.flowfile.sources.external_sources.base_class import ExternalDataSource
+import polars as pl
+
+
+class CustomExternalSourceSettings:
+    data_getter: Generator
+    initial_data_getter: Optional[Callable] = None
+    orientation: str = 'row'
+
+    def __init__(self, data_getter: Generator, initial_data_getter: Optional[Callable] = None, orientation: str = 'row'):
+        self.data_getter = data_getter
+        self.initial_data_getter = initial_data_getter
+        self.orientation = orientation
+
+
+class CustomExternalSource(ExternalDataSource):
+    data_getter: Generator = None
+    schema: Optional[List[FlowfileColumn]] = None
+    cache_store: List = None
+    is_collected: bool = False
+
+    def __init__(self, data_getter: Generator[Any, None, None],
+                 initial_data_getter: Callable = None,
+                 orientation: str = 'row',
+                 schema: List = None,
+                 **kwargs):
+        self.cache_store = list()
+        self.data_getter = data_getter
+        self.collected = False
+        if schema is not None:
+            try:
+                self.schema = self.parse_schema(schema)
+            except ValueError:
+                self.schema = None
+        else:
+            self.schema = None
+
+        if not initial_data_getter and orientation == 'row':
+            def initial_data_getter():
+                if len(self.cache_store) == 0:
+                    self.cache_store.append(next(data_getter, None))
+                return self.cache_store
+            self.initial_data_getter = initial_data_getter
+        elif initial_data_getter:
+            self.initial_data_getter = initial_data_getter
+        elif self.schema:
+            def initial_data_getter():
+                return [{d.column_name: None for d in self.schema}]
+            self.initial_data_getter = initial_data_getter
+        else:
+            self.initial_data_getter = None
+
+    @staticmethod
+    def parse_schema(schema: List[Any]) -> List[FlowfileColumn]:
+        if len(schema) == 0:
+            return []
+        first_col = schema[0]
+        if isinstance(first_col, dict):
+            return [FlowfileColumn(**col) for col in schema]
+        elif isinstance(first_col, (list, tuple)):
+            return [FlowfileColumn.from_input(column_name=col[0], data_type=col[1]) for col in schema]
+        elif isinstance(first_col, str):
+            return [FlowfileColumn.from_input(column_name=col, data_type='varchar') for col in schema]
+        elif isinstance(first_col, input_schema.MinimalFieldInfo):
+            return [FlowfileColumn.from_input(column_name=col.name, data_type=col.data_type) for col in schema]
+        elif isinstance(first_col, FlowfileColumn):
+            return schema
+        else:
+            raise ValueError("Schema is not a valid type")
+
+    def get_initial_data(self):
+        if self.initial_data_getter:
+            return self.initial_data_getter()
+        return []
+
+    def get_iter(self) -> Generator[Dict[str, Any], None, None]:
+        if self.collected:
+            return
+        for data in self.cache_store:
+            yield data
+        for data in self.data_getter:
+            self.cache_store.append(data)
+            yield data
+        self.is_collected = True
+        return
+
+    def get_sample(self, n: int = 10000):
+        data = self.get_iter()
+        for i in range(n):
+            try:
+                yield next(data)
+            except StopIteration:
+                break
+
+    def get_pl_df(self) -> pl.DataFrame:
+        data = self.get_iter()
+        return pl.DataFrame(list(data))

flowfile_core/flowfile/sources/external_sources/custom_external_sources/google_sheet.py

@@ -0,0 +1,74 @@
+import gspread
+from google.oauth2.credentials import Credentials
+from googleapiclient.discovery import build
+from typing import Dict, Any, Generator, List
+from flowfile_core.schemas.input_schema import GoogleSheet
+
+
+def getter(data: GoogleSheet) -> Generator[Dict[str, Any], None, None]:
+    creds = Credentials(token=data.access_token.get_secret_value())
+    gc = gspread.authorize(credentials=creds)
+    worksheet = gc.open_by_key(data.sheet_id).worksheet(data.worksheet_name)
+    all_values = worksheet.get_values()
+    headers = all_values[0]
+    for i, value in enumerate(all_values[1:]):
+        yield {k: v for k, v in zip(headers, value)}
+
+
+def initial_getter(data: GoogleSheet):
+    def inner_func():
+        creds = Credentials(token=data.access_token.get_secret_value())
+        service = build('sheets', 'v4', credentials=creds)
+        _range = f'{data.worksheet_name}!1:2'
+        result = service.spreadsheets().values().get(spreadsheetId=data.sheet_id, range=_range).execute().get('values', [])
+        if len(result) > 1:
+            align_data(result[0], result[1:])
+        return [{k: v for k, v in zip(*result)}]
+    return inner_func
+
+
+def align_data(headers: List[str], values: List[List[str]]) -> None:
+    """
+    Ensures that the number of columns in 'headers' matches the maximum row length in 'values'.
+    If 'headers' has fewer columns, it appends 'unknown_column_{i}' for the missing columns.
+    Then aligns all rows in 'values' and 'headers' to have the same length.
+
+    Args:
+        headers (List[str]): A list of column names (headers).
+        values (List[List[str]]): A list of rows, where each row is a list of strings.
+
+    Returns:
+        None: The function modifies 'headers' and 'values' in place.
+    """
+    # Find the maximum number of values in any row
+    max_number_of_values = max(len(row) for row in values)
+
+    # Find the current number of columns
+    number_of_cols = len(headers)
+
+    # If headers have fewer columns than the maximum number of values, append missing columns
+    if number_of_cols < max_number_of_values:
+        headers.extend(f'unknown_column_{i}' for i in range(number_of_cols, max_number_of_values))
+
+    # Align all rows and headers to have the same length
+    align_list_len([headers] + values)
+
+
+def align_list_len(values: List[List[str]], default_value: str = '') -> None:
+    """
+    Aligns the size of each sublist in a list of lists by appending the default value to shorter sublists.
+    Modifies the list of lists in place.
+
+    Args:
+        values (List[List[str]]): A list of lists where each sublist may have different lengths.
+        default_value (str): The value to append to each sublist to align their lengths (default is an empty string).
+
+    Returns:
+        None
+    """
+    # Determine the maximum length of the sublists
+    max_len = max(len(sublist) for sublist in values)
+
+    # Extend each sublist to the maximum length by appending the default value
+    for sublist in values:
+        sublist.extend([default_value] * (max_len - len(sublist)))

flowfile_core/flowfile/sources/external_sources/custom_external_sources/sample_users.py

@@ -0,0 +1,29 @@
+from typing import Dict, Any, Generator
+from time import sleep
+from flowfile_core.schemas.input_schema import SampleUsers
+import requests
+
+
+def getter(data: SampleUsers) -> Generator[Dict[str, Any], None, None]:
+    """
+    Sample users generator function. This is a minimal example of a generator function that yields user data and can
+    be used in a flowfile. The function simulates a delay to mimic the behavior of an external data source.
+    Args:
+        data ():
+
+    Returns:
+
+    """
+    index_pos = 0
+    for i in range(data.size):
+        sleep(0.01)
+        headers = {
+            'x-api-key': 'reqres-free-v1'
+        }
+
+        response = requests.get("https://reqres.in/api/users", headers=headers).json()
+        for v in response['data']:
+            v['index'] = index_pos
+            index_pos += 1
+            yield v
+

flowfile_core/flowfile/sources/external_sources/factory.py

@@ -0,0 +1,22 @@
+from flowfile_core.flowfile.sources.external_sources.custom_external_sources.external_source import CustomExternalSource
+from flowfile_core.flowfile.sources.external_sources.airbyte_sources.airbyte import AirbyteSource
+
+
+def data_source_factory(source_type: str, **kwargs) -> CustomExternalSource | AirbyteSource:
+    """
+    Factory function to generate either CustomExternalSource or AirbyteSource.
+
+    Args:
+        source_type (str): The type of source to create ("custom" or "airbyte").
+        **kwargs: The keyword arguments required for the specific source type.
+
+    Returns:
+        Union[CustomExternalSource, AirbyteSource]: An instance of the selected data source type.
+    """
+    if source_type == "custom":
+        return CustomExternalSource(**kwargs)
+    elif source_type == "airbyte":
+        return AirbyteSource(**kwargs)
+    else:
+        raise ValueError(f"Unknown source type: {source_type}")
+

flowfile_core/flowfile/sources/external_sources/sql_source/models.py

@@ -0,0 +1,90 @@
+from typing import Literal, Optional, TYPE_CHECKING
+from pydantic import BaseModel, SecretStr
+from flowfile_core.schemas.input_schema import (DatabaseConnection,
+                                                NodeDatabaseReader,
+                                                FullDatabaseConnection,
+                                                NodeDatabaseWriter)
+import base64
+import polars as pl
+
+
+class ExtDatabaseConnection(DatabaseConnection):
+    """Database connection configuration with password handling."""
+    password: str = None
+
+
+class DatabaseExternalWriteSettings(BaseModel):
+    """Settings for SQL sink."""
+    connection: ExtDatabaseConnection
+    table_name: str
+    if_exists: Optional[Literal['append', 'replace', 'fail']] = 'append'
+    flowfile_flow_id: int = 1
+    flowfile_node_id: int | str = -1
+    operation: str
+
+    @classmethod
+    def create_from_from_node_database_writer(cls, node_database_writer: NodeDatabaseWriter,
+                                              password: str,
+                                              table_name: str,
+                                              lf: pl.LazyFrame,
+                                              database_reference_settings: FullDatabaseConnection = None,
+                                              ) -> 'DatabaseExternalWriteSettings':
+        """
+        Create DatabaseExternalWriteSettings from NodeDatabaseWriter.
+        Args:
+            node_database_writer (NodeDatabaseWriter): an instance of NodeDatabaseWriter
+            password (str): the password for the database connection
+            table_name (str): the table name to be used for writing
+            lf (pl.LazyFrame): the LazyFrame to be written to the database
+            database_reference_settings (FullDatabaseConnection): optional database reference settings
+        Returns:
+            DatabaseExternalReadSettings: an instance of DatabaseExternalReadSettings
+        """
+        if node_database_writer.database_write_settings.connection_mode == "inline":
+            database_connection = node_database_writer.database_settings.database_connection.model_dump()
+        else:
+            database_connection = {k: v for k, v in database_reference_settings.model_dump().items() if k != "password"}
+
+        ext_database_connection = ExtDatabaseConnection(**database_connection,
+                                                        password=password)
+        return cls(connection=ext_database_connection,
+                   table_name=table_name,
+                   if_exists=node_database_writer.database_write_settings.if_exists,
+                   flowfile_flow_id=node_database_writer.flow_id,
+                   flowfile_node_id=node_database_writer.node_id,
+                   operation=base64.b64encode(lf.serialize()).decode())
+
+
+class DatabaseExternalReadSettings(BaseModel):
+    """Settings for SQL source."""
+    connection: ExtDatabaseConnection
+    query: str
+    flowfile_flow_id: int = 1
+    flowfile_node_id: int | str = -1
+
+    @classmethod
+    def create_from_from_node_database_reader(cls, node_database_reader: NodeDatabaseReader,
+                                              password: str,
+                                              query: str,
+                                              database_reference_settings: FullDatabaseConnection = None) -> 'DatabaseExternalReadSettings':
+        """
+        Create DatabaseExternalReadSettings from NodeDatabaseReader.
+        Args:
+            node_database_reader (NodeDatabaseReader): an instance of NodeDatabaseReader
+            password (str): the password for the database connection
+            query (str): the SQL query to be executed
+            database_reference_settings (FullDatabaseConnection): optional database reference settings
+        Returns:
+            DatabaseExternalReadSettings: an instance of DatabaseExternalReadSettings
+        """
+        if node_database_reader.database_settings.connection_mode == "inline":
+            database_connection = node_database_reader.database_settings.database_connection.model_dump()
+        else:
+            database_connection = {k: v for k, v in database_reference_settings.model_dump().items() if k != "password"}
+
+        ext_database_connection = ExtDatabaseConnection(**database_connection,
+                                                        password=password)
+        return cls(connection=ext_database_connection,
+                   query=query,
+                   flowfile_flow_id=node_database_reader.flow_id,
+                   flowfile_node_id=node_database_reader.node_id)