Flowfile 0.3.9__py3-none-any.whl → 0.5.1__py3-none-any.whl

flowfile_core/flowfile/flow_graph.py

@@ -1,20 +1,27 @@
 import datetime
-import pickle
+
+import os
+import yaml
+import json
+
 import polars as pl
+from pathlib import Path
+
 import fastexcel
 from fastapi.exceptions import HTTPException
 from time import time
 from functools import partial
-from typing import List, Dict, Union, Callable, Any, Optional, Tuple
+from typing import List, Dict, Union, Callable, Any, Optional, Tuple, Literal
 from uuid import uuid1
 from copy import deepcopy
 from pyarrow.parquet import ParquetFile
 from flowfile_core.configs import logger
 from flowfile_core.configs.flow_logger import FlowLogger
 from flowfile_core.flowfile.sources.external_sources.factory import data_source_factory
-from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import cast_str_to_polars_type, FlowfileColumn
+from flowfile_core.flowfile.flow_data_engine.flow_file_column.main import FlowfileColumn, cast_str_to_polars_type
 
 from flowfile_core.flowfile.flow_data_engine.cloud_storage_reader import CloudStorageReader
+from flowfile_core.schemas.transform_schema import FuzzyMatchInputManager
 from flowfile_core.utils.arrow_reader import get_read_top_n
 from flowfile_core.flowfile.flow_data_engine.flow_data_engine import FlowDataEngine, execute_polars_code
 from flowfile_core.flowfile.flow_data_engine.read_excel_tables import (get_open_xlsx_datatypes,
@@ -47,6 +54,23 @@ from flowfile_core.flowfile.sources.external_sources.sql_source.sql_source impor
 from flowfile_core.flowfile.database_connection_manager.db_connections import (get_local_database_connection,
                                                                                get_local_cloud_connection)
 from flowfile_core.flowfile.util.calculate_layout import calculate_layered_layout
+from flowfile_core.flowfile.node_designer.custom_node import CustomNodeBase
+from importlib.metadata import version, PackageNotFoundError
+
+try:
+    __version__ = version("Flowfile")
+except PackageNotFoundError:
+    __version__ = "0.0.0-dev"
+
+
+def represent_list_json(dumper, data):
+    """Use inline style for short simple lists, block style for complex ones."""
+    if len(data) <= 10 and all(isinstance(item, (int, str, float, bool, type(None))) for item in data):
+        return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=True)
+    return dumper.represent_sequence('tag:yaml.org,2002:seq', data, flow_style=False)
+
+
+yaml.add_representer(list, represent_list_json)
 
 
 def get_xlsx_schema(engine: str, file_path: str, sheet_name: str, start_row: int, start_column: int,
@@ -147,7 +171,7 @@ def get_cloud_connection_settings(connection_name: str,
         HTTPException: If the connection settings cannot be found.
     """
     cloud_connection_settings = get_local_cloud_connection(connection_name, user_id)
-    if cloud_connection_settings is None and auth_mode in ("env_vars", "auto"):
+    if cloud_connection_settings is None and auth_mode in ("env_vars", transform_schema.AUTO_DATA_TYPE):
         # If the auth mode is aws-cli, we do not need connection settings
         cloud_connection_settings = FullCloudStorageConnection(storage_type="s3", auth_method="env_vars")
     elif cloud_connection_settings is None and auth_mode == "aws-cli":
@@ -175,11 +199,9 @@ class FlowGraph:
     schema: Optional[List[FlowfileColumn]] = None
     has_over_row_function: bool = False
     _flow_starts: List[Union[int, str]] = None
-    node_results: List[NodeResult] = None
     latest_run_info: Optional[RunInformation] = None
     start_datetime: datetime = None
     end_datetime: datetime = None
-    nodes_completed: int = 0
     _flow_settings: schemas.FlowSettings = None
     flow_logger: FlowLogger
 
@@ -206,11 +228,9 @@ class FlowGraph:
 
         self._flow_settings = flow_settings
         self.uuid = str(uuid1())
-        self.nodes_completed = 0
         self.start_datetime = None
         self.end_datetime = None
         self.latest_run_info = None
-        self.node_results = []
         self._flow_id = flow_settings.flow_id
         self.flow_logger = FlowLogger(flow_settings.flow_id)
         self._flow_starts: List[FlowNode] = []
@@ -222,7 +242,7 @@ class FlowGraph:
         self._node_ids = []
         self._node_db = {}
         self.cache_results = cache_results
-        self.__name__ = name if name else id(self)
+        self.__name__ = name if name else "flow_" + str(id(self))
         self.depends_on = {}
         if path_ref is not None:
             self.add_datasource(input_schema.NodeDatasource(file_path=path_ref))
@@ -436,6 +456,24 @@ class FlowGraph:
         node = self._node_db.get(node_id)
         if node is not None:
             return node
+        
+    def add_user_defined_node(self, *,
+                              custom_node: CustomNodeBase,
+                              user_defined_node_settings: input_schema.UserDefinedNode
+                              ):
+       
+        def _func(*fdes: FlowDataEngine) -> FlowDataEngine | None:
+            output = custom_node.process(*(fde.data_frame for fde in fdes))
+            if isinstance(output, pl.LazyFrame | pl.DataFrame):
+                return FlowDataEngine(output)
+            return None
+        
+        self.add_node_step(node_id=user_defined_node_settings.node_id,
+                           function=_func,
+                           setting_input=user_defined_node_settings,
+                           input_node_ids=user_defined_node_settings.depending_on_ids,
+                           node_type=custom_node.item,
+                           )
 
     def add_pivot(self, pivot_settings: input_schema.NodePivot):
         """Adds a pivot node to the graph.
@@ -714,11 +752,11 @@ class FlowGraph:
         """
 
         error = ""
-        if function_settings.function.field.data_type not in (None, "Auto"):
+        if function_settings.function.field.data_type not in (None, transform_schema.AUTO_DATA_TYPE):
             output_type = cast_str_to_polars_type(function_settings.function.field.data_type)
         else:
             output_type = None
-        if output_type not in (None, "Auto"):
+        if output_type not in (None, transform_schema.AUTO_DATA_TYPE):
             new_col = [FlowfileColumn.from_input(column_name=function_settings.function.field.name,
                                                  data_type=str(output_type))]
         else:
@@ -736,6 +774,7 @@ class FlowGraph:
                            setting_input=function_settings,
                            input_node_ids=[function_settings.depending_on_id]
                            )
+        # TODO: Add validation here
         if error != "":
             node = self.get_node(function_settings.node_id)
             node.results.errors = error
@@ -752,13 +791,11 @@ class FlowGraph:
         Returns:
             The `FlowGraph` instance for method chaining.
         """
-
         def _func(main: FlowDataEngine, right: FlowDataEngine) -> FlowDataEngine:
             for left_select in cross_join_settings.cross_join_input.left_select.renames:
                 left_select.is_available = True if left_select.old_name in main.schema else False
             for right_select in cross_join_settings.cross_join_input.right_select.renames:
                 right_select.is_available = True if right_select.old_name in right.schema else False
-
             return main.do_cross_join(cross_join_input=cross_join_settings.cross_join_input,
                                       auto_generate_selection=cross_join_settings.auto_generate_selection,
                                       verify_integrity=False,
@@ -781,13 +818,11 @@ class FlowGraph:
         Returns:
             The `FlowGraph` instance for method chaining.
         """
-
         def _func(main: FlowDataEngine, right: FlowDataEngine) -> FlowDataEngine:
             for left_select in join_settings.join_input.left_select.renames:
                 left_select.is_available = True if left_select.old_name in main.schema else False
             for right_select in join_settings.join_input.right_select.renames:
                 right_select.is_available = True if right_select.old_name in right.schema else False
-
             return main.join(join_input=join_settings.join_input,
                              auto_generate_selection=join_settings.auto_generate_selection,
                              verify_integrity=False,
@@ -814,18 +849,18 @@ class FlowGraph:
         def _func(main: FlowDataEngine, right: FlowDataEngine) -> FlowDataEngine:
             node = self.get_node(node_id=fuzzy_settings.node_id)
             if self.execution_location == "local":
-                return main.fuzzy_join(fuzzy_match_input=fuzzy_settings.join_input,
+                return main.fuzzy_join(fuzzy_match_input=deepcopy(fuzzy_settings.join_input),
                                        other=right,
                                        node_logger=self.flow_logger.get_node_logger(fuzzy_settings.node_id))
 
-            f = main.start_fuzzy_join(fuzzy_match_input=fuzzy_settings.join_input, other=right, file_ref=node.hash,
+            f = main.start_fuzzy_join(fuzzy_match_input=deepcopy(fuzzy_settings.join_input), other=right, file_ref=node.hash,
                                       flow_id=self.flow_id, node_id=fuzzy_settings.node_id)
             logger.info("Started the fuzzy match action")
             node._fetch_cached_df = f  # Add to the node so it can be cancelled and fetch later if needed
             return FlowDataEngine(f.get_result())
 
         def schema_callback():
-            fm_input_copy = deepcopy(fuzzy_settings.join_input)  # Deepcopy create an unique object per func
+            fm_input_copy = FuzzyMatchInputManager(fuzzy_settings.join_input)  # Deepcopy create an unique object per func
             node = self.get_node(node_id=fuzzy_settings.node_id)
             return calculate_fuzzy_match_schema(fm_input_copy,
                                                 left_schema=node.node_inputs.main_inputs[0].schema,
@@ -1112,7 +1147,6 @@ class FlowGraph:
         """
 
         def _func(df: FlowDataEngine):
-            output_file.output_settings.populate_abs_file_path()
             execute_remote = self.execution_location != 'local'
             df.output(output_fs=output_file.output_settings, flow_id=self.flow_id, node_id=output_file.node_id,
                       execute_remote=execute_remote)
@@ -1432,10 +1466,10 @@ class FlowGraph:
         Args:
             input_file: The settings for the read operation.
         """
-
-        if input_file.received_file.file_type in ('xlsx', 'excel') and input_file.received_file.sheet_name == '':
+        if (input_file.received_file.file_type in ('xlsx', 'excel') and
+                input_file.received_file.table_settings.sheet_name == ''):
             sheet_name = fastexcel.read_excel(input_file.received_file.path).sheet_names[0]
-            input_file.received_file.sheet_name = sheet_name
+            input_file.received_file.table_settings.sheet_name = sheet_name
 
         received_file = input_file.received_file
         input_file.received_file.set_absolute_filepath()
@@ -1444,7 +1478,7 @@ class FlowGraph:
             input_file.received_file.set_absolute_filepath()
             if input_file.received_file.file_type == 'parquet':
                 input_data = FlowDataEngine.create_from_path(input_file.received_file)
-            elif input_file.received_file.file_type == 'csv' and 'utf' in input_file.received_file.encoding:
+            elif input_file.received_file.file_type == 'csv' and 'utf' in input_file.received_file.table_settings.encoding:
                 input_data = FlowDataEngine.create_from_path(input_file.received_file)
             else:
                 input_data = FlowDataEngine.create_from_path_worker(input_file.received_file,
@@ -1481,12 +1515,12 @@ class FlowGraph:
                     # If the file is an Excel file, we need to use the openpyxl engine to read the schema
                     schema_callback = get_xlsx_schema_callback(engine='openpyxl',
                                                                file_path=received_file.file_path,
-                                                               sheet_name=received_file.sheet_name,
-                                                               start_row=received_file.start_row,
-                                                               end_row=received_file.end_row,
-                                                               start_column=received_file.start_column,
-                                                               end_column=received_file.end_column,
-                                                               has_headers=received_file.has_headers)
+                                                               sheet_name=received_file.table_settings.sheet_name,
+                                                               start_row=received_file.table_settings.start_row,
+                                                               end_row=received_file.table_settings.end_row,
+                                                               start_column=received_file.table_settings.start_column,
+                                                               end_column=received_file.table_settings.end_column,
+                                                               has_headers=received_file.table_settings.has_headers)
                 else:
                     schema_callback = None
         else:
@@ -1599,6 +1633,73 @@ class FlowGraph:
             self.reset()
         self.flow_settings.execution_location = execution_location
 
+    def validate_if_node_can_be_fetched(self, node_id: int) -> None:
+        flow_node = self._node_db.get(node_id)
+        if not flow_node:
+            raise Exception("Node not found found")
+        skip_nodes, execution_order = compute_execution_plan(
+            nodes=self.nodes, flow_starts=self._flow_starts+self.get_implicit_starter_nodes()
+        )
+        if flow_node.node_id in [skip_node.node_id for skip_node in skip_nodes]:
+            raise Exception("Node can not be executed because it does not have it's inputs")
+
+    def create_initial_run_information(self, number_of_nodes: int,
+                                       run_type: Literal["fetch_one", "full_run"]):
+        return RunInformation(
+            flow_id=self.flow_id, start_time=datetime.datetime.now(), end_time=None,
+            success=None, number_of_nodes=number_of_nodes, node_step_result=[],
+            run_type=run_type
+        )
+
+    def create_empty_run_information(self) -> RunInformation:
+        return RunInformation(
+            flow_id=self.flow_id, start_time=None, end_time=None,
+            success=None, number_of_nodes=0, node_step_result=[],
+            run_type="init"
+        )
+
+    def trigger_fetch_node(self, node_id: int) -> RunInformation | None:
+        """Executes a specific node in the graph by its ID."""
+        if self.flow_settings.is_running:
+            raise Exception("Flow is already running")
+        flow_node = self.get_node(node_id)
+        self.flow_settings.is_running = True
+        self.flow_settings.is_canceled = False
+        self.flow_logger.clear_log_file()
+        self.latest_run_info = self.create_initial_run_information(1, "fetch_one")
+        node_logger = self.flow_logger.get_node_logger(flow_node.node_id)
+        node_result = NodeResult(node_id=flow_node.node_id, node_name=flow_node.name)
+        logger.info(f'Starting to run: node {flow_node.node_id}, start time: {node_result.start_timestamp}')
+        try:
+            self.latest_run_info.node_step_result.append(node_result)
+            flow_node.execute_node(run_location=self.flow_settings.execution_location,
+                                   performance_mode=False,
+                                   node_logger=node_logger,
+                                   optimize_for_downstream=False,
+                                   reset_cache=True)
+            node_result.error = str(flow_node.results.errors)
+            if self.flow_settings.is_canceled:
+                node_result.success = None
+                node_result.success = None
+                node_result.is_running = False
+            node_result.success = flow_node.results.errors is None
+            node_result.end_timestamp = time()
+            node_result.run_time = int(node_result.end_timestamp - node_result.start_timestamp)
+            node_result.is_running = False
+            self.latest_run_info.nodes_completed += 1
+            self.latest_run_info.end_time = datetime.datetime.now()
+            self.flow_settings.is_running = False
+            return self.get_run_info()
+        except Exception as e:
+            node_result.error = 'Node did not run'
+            node_result.success = False
+            node_result.end_timestamp = time()
+            node_result.run_time = int(node_result.end_timestamp - node_result.start_timestamp)
+            node_result.is_running = False
+            node_logger.error(f'Error in node {flow_node.node_id}: {e}')
+        finally:
+            self.flow_settings.is_running = False
+
     def run_graph(self) -> RunInformation | None:
         """Executes the entire data flow graph from start to finish.
 
@@ -1614,20 +1715,23 @@ class FlowGraph:
         if self.flow_settings.is_running:
             raise Exception('Flow is already running')
         try:
+
             self.flow_settings.is_running = True
             self.flow_settings.is_canceled = False
             self.flow_logger.clear_log_file()
-            self.nodes_completed = 0
-            self.node_results = []
-            self.start_datetime = datetime.datetime.now()
-            self.end_datetime = None
-            self.latest_run_info = None
             self.flow_logger.info('Starting to run flowfile flow...')
-            skip_nodes, execution_order = compute_execution_plan(nodes=self.nodes, flow_starts=self._flow_starts+self.get_implicit_starter_nodes())
+
+            skip_nodes, execution_order = compute_execution_plan(
+                nodes=self.nodes,
+                flow_starts=self._flow_starts+self.get_implicit_starter_nodes()
+            )
+
+            self.latest_run_info = self.create_initial_run_information(len(execution_order), "full_run")
 
             skip_node_message(self.flow_logger, skip_nodes)
             execution_order_message(self.flow_logger, execution_order)
             performance_mode = self.flow_settings.execution_mode == 'Performance'
+
             for node in execution_order:
                 node_logger = self.flow_logger.get_node_logger(node.node_id)
                 if self.flow_settings.is_canceled:
@@ -1637,7 +1741,7 @@ class FlowGraph:
                     node_logger.info(f'Skipping node {node.node_id}')
                     continue
                 node_result = NodeResult(node_id=node.node_id, node_name=node.name)
-                self.node_results.append(node_result)
+                self.latest_run_info.node_step_result.append(node_result)
                 logger.info(f'Starting to run: node {node.node_id}, start time: {node_result.start_timestamp}')
                 node.execute_node(run_location=self.flow_settings.execution_location,
                                   performance_mode=performance_mode,
@@ -1663,7 +1767,8 @@ class FlowGraph:
                 if not node_result.success:
                     skip_nodes.extend(list(node.get_all_dependent_nodes()))
                 node_logger.info(f'Completed node with success: {node_result.success}')
-                self.nodes_completed += 1
+                self.latest_run_info.nodes_completed += 1
+            self.latest_run_info.end_time = datetime.datetime.now()
             self.flow_logger.info('Flow completed!')
             self.end_datetime = datetime.datetime.now()
             self.flow_settings.is_running = False
@@ -1681,22 +1786,17 @@ class FlowGraph:
         Returns:
             A RunInformation object with details about the last run.
         """
+        is_running = self.flow_settings.is_running
         if self.latest_run_info is None:
-            node_results = self.node_results
-            success = all(nr.success for nr in node_results)
-            self.latest_run_info = RunInformation(start_time=self.start_datetime, end_time=self.end_datetime,
-                                                  success=success,
-                                                  node_step_result=node_results, flow_id=self.flow_id,
-                                                  nodes_completed=self.nodes_completed,
-                                                  number_of_nodes=len(self.nodes))
-        elif self.latest_run_info.nodes_completed != self.nodes_completed:
-            node_results = self.node_results
-            self.latest_run_info = RunInformation(start_time=self.start_datetime, end_time=self.end_datetime,
-                                                  success=all(nr.success for nr in node_results),
-                                                  node_step_result=node_results, flow_id=self.flow_id,
-                                                  nodes_completed=self.nodes_completed,
-                                                  number_of_nodes=len(self.nodes))
-        return self.latest_run_info
+            return self.create_empty_run_information()
+
+        elif not is_running and self.latest_run_info.success is not None:
+            return self.latest_run_info
+
+        run_info = self.latest_run_info
+        if not is_running:
+            run_info.success = all(nr.success for nr in run_info.node_step_result)
+        return run_info
 
     @property
     def node_connections(self) -> List[Tuple[int, int]]:
@@ -1729,6 +1829,42 @@ class FlowGraph:
         node = self._node_db[node_id]
         return node.get_node_data(flow_id=self.flow_id, include_example=include_example)
 
+    def get_flowfile_data(self) -> schemas.FlowfileData:
+        start_node_ids = {v.node_id for v in self._flow_starts}
+
+        nodes = []
+        for node in self.nodes:
+            node_info = node.get_node_information()
+            flowfile_node = schemas.FlowfileNode(
+                id=node_info.id,
+                type=node_info.type,
+                is_start_node=node.node_id in start_node_ids,
+                description=node_info.description,
+                x_position=int(node_info.x_position),
+                y_position=int(node_info.y_position),
+                left_input_id=node_info.left_input_id,
+                right_input_id=node_info.right_input_id,
+                input_ids=node_info.input_ids,
+                outputs=node_info.outputs,
+                setting_input=node_info.setting_input,
+            )
+            nodes.append(flowfile_node)
+
+        settings = schemas.FlowfileSettings(
+            description=self.flow_settings.description,
+            execution_mode=self.flow_settings.execution_mode,
+            execution_location=self.flow_settings.execution_location,
+            auto_save=self.flow_settings.auto_save,
+            show_detailed_progress=self.flow_settings.show_detailed_progress,
+        )
+        return schemas.FlowfileData(
+            flowfile_version=__version__,
+            flowfile_id=self.flow_id,
+            flowfile_name=self.__name__,
+            flowfile_settings=settings,
+            nodes=nodes,
+        )
+
     def get_node_storage(self) -> schemas.FlowInformation:
         """Serializes the entire graph's state into a storable format.
 
@@ -1761,14 +1897,64 @@ class FlowGraph:
         for node in self.nodes:
             node.remove_cache()
 
+    def _handle_flow_renaming(self, new_name: str, new_path: Path):
+        """
+        Handle the rename of a flow when it is being saved.
+        """
+        if self.flow_settings and self.flow_settings.path and Path(self.flow_settings.path).absolute() != new_path.absolute():
+            self.__name__ = new_name
+            self.flow_settings.save_location = str(new_path.absolute())
+            self.flow_settings.name = new_name
+        if self.flow_settings and not self.flow_settings.save_location:
+            self.flow_settings.save_location = str(new_path.absolute())
+            self.__name__ = new_name
+            self.flow_settings.name = new_name
+
     def save_flow(self, flow_path: str):
         """Saves the current state of the flow graph to a file.
 
+        Supports multiple formats based on file extension:
+        - .yaml / .yml: New YAML format
+        - .json: JSON format
+
         Args:
             flow_path: The path where the flow file will be saved.
         """
-        with open(flow_path, 'wb') as f:
-            pickle.dump(self.get_node_storage(), f)
+        logger.info("Saving flow to %s", flow_path)
+        path = Path(flow_path)
+        os.makedirs(path.parent, exist_ok=True)
+        suffix = path.suffix.lower()
+        new_flow_name = path.name.replace(suffix, "")
+        self._handle_flow_renaming(new_flow_name, path)
+        self.flow_settings.modified_on = datetime.datetime.now().timestamp()
+        try:
+            if suffix == '.flowfile':
+                raise DeprecationWarning(
+                    f"The .flowfile format is deprecated. Please use .yaml or .json formats.\n\n"
+                    "Or stay on v0.4.1 if you still need .flowfile support.\n\n"
+                )
+            elif suffix in ('.yaml', '.yml'):
+                flowfile_data = self.get_flowfile_data()
+                data = flowfile_data.model_dump(mode='json')
+                with open(flow_path, 'w', encoding='utf-8') as f:
+                    yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
+            elif suffix == '.json':
+                flowfile_data = self.get_flowfile_data()
+                data = flowfile_data.model_dump(mode='json')
+                with open(flow_path, 'w', encoding='utf-8') as f:
+                    json.dump(data, f, indent=2, ensure_ascii=False)
+
+            else:
+                flowfile_data = self.get_flowfile_data()
+                logger.warning(f"Unknown file extension {suffix}. Defaulting to YAML format.")
+                data = flowfile_data.model_dump(mode='json')
+                with open(flow_path, 'w', encoding='utf-8') as f:
+                    yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
+
+        except Exception as e:
+            logger.error(f"Error saving flow: {e}")
+            raise
+
         self.flow_settings.path = flow_path
 
     def get_frontend_data(self) -> dict: