cloe-nessy 0.2.9__py3-none-any.whl

cloe_nessy/pipeline/actions/write_catalog_table.py

@@ -0,0 +1,73 @@
+from typing import Any
+
+from ...integration.writer import CatalogWriter
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+
+class WriteCatalogTableAction(PipelineAction):
+    """Writes a DataFrame to a specified catalog table using [CatalogWriter][cloe_nessy.integration.writer.CatalogWriter].
+
+    Example:
+    ```yaml
+    Write Table to Catalog:
+        action: WRITE_CATALOG_TABLE
+        options:
+            table_identifier: my_catalog.business_schema.sales_table
+            mode: append
+            partition_by: day
+            options: <options for the writer>
+    ```
+    """
+
+    name: str = "WRITE_CATALOG_TABLE"
+
+    @staticmethod
+    def run(
+        context: PipelineContext,
+        *,
+        table_identifier: str | None = None,
+        mode: str = "append",
+        partition_by: str | list[str] | None = None,
+        options: dict[str, str] | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Writes a DataFrame to a specified catalog table.
+
+        Args:
+            context: Context in which this Action is executed.
+            table_identifier: The table identifier in the unity catalog in the
+                format 'catalog.schema.table'. If not provided, attempts to use the
+                context's table metadata.
+            mode: The write mode. One of 'append', 'overwrite', 'error',
+                'errorifexists', or 'ignore'.
+            partition_by: Names of the partitioning columns.
+            options: Additional options for the write operation.
+
+        Raises:
+            ValueError: If the table name is not specified or cannot be inferred from
+                the context.
+
+        Returns:
+            Context after the execution of this Action.
+        """
+        if not options:
+            options = dict()
+        if partition_by is None:
+            if hasattr(context.table_metadata, "partition_by"):
+                partition_by = context.table_metadata.partition_by  # type: ignore
+
+        if (table_metadata := context.table_metadata) and table_identifier is None:
+            table_identifier = table_metadata.identifier
+        if table_identifier is None:
+            raise ValueError("Table name must be specified or a valid Table object with identifier must be set.")
+
+        writer = CatalogWriter()
+        writer.write_table(
+            df=context.data,  # type: ignore
+            table_identifier=table_identifier,
+            mode=mode,
+            partition_by=partition_by,
+            options=options,
+        )
+        return context.from_existing()

cloe_nessy/pipeline/pipeline.py

@@ -0,0 +1,201 @@
+import os
+from collections import OrderedDict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from threading import Lock
+
+import matplotlib.pyplot as plt
+import networkx as nx
+
+from ..logging.logger_mixin import LoggerMixin
+from .pipeline_step import PipelineStep
+
+
+class Pipeline(LoggerMixin):
+    """A Pipeline represents the logical unit of one ETL process.
+
+    This class manages a directed acyclic graph (DAG) of steps, ensuring that
+    each step is executed in the correct order based on dependencies.
+
+    Attributes:
+        name: The name of the pipeline.
+        steps: An ordered dictionary of PipelineSteps that are part of the pipeline.
+    """
+
+    def __init__(self, name: str, steps: OrderedDict[str, "PipelineStep"] | None = None) -> None:
+        self.name: str = name
+        self.steps: OrderedDict[str, PipelineStep] = steps if steps is not None else OrderedDict()
+        self._console_logger = self.get_console_logger()
+        self._graph: nx.DiGraph = self._create_graph()
+        self._lock: Lock = Lock()
+
+    @property
+    def graph(self) -> nx.DiGraph:
+        """Get the pipeline graph."""
+        return self._graph
+
+    def _create_graph(self) -> nx.DiGraph:
+        """Creates a directed acyclic graph (DAG) representing the pipeline steps and their dependencies.
+
+        Each node in the graph represents a single step in the pipeline, and each edge represents a dependency.
+        """
+        g: nx.DiGraph = nx.DiGraph()
+        g.add_nodes_from(set([s.name for s in self.steps.values()]))
+        g.add_edges_from(set([(p, s.name) for s in self.steps.values() for p in s._predecessors if p]))
+
+        self._console_logger.debug(f"Graph created with {g.number_of_nodes()} nodes and {g.number_of_edges()} edges.")
+        return g
+
+    def _run_step(self, step_name: str) -> None:
+        """Executes the run method of the corresponding step in the pipeline."""
+        step = self.steps[step_name]
+
+        # Handle context and metadata references
+        if step._context_ref:
+            step.context = self.steps[step._context_ref].result
+        if step._table_metadata_ref:
+            step.context.table_metadata = self.steps[step._table_metadata_ref].result.table_metadata
+
+        try:
+            self._console_logger.info(f"Starting execution of step: {step.name}")
+            step.run()
+        except Exception as err:
+            self._console_logger.error(f"Execution of step {step.name} failed with error: {str(err)}")
+            raise err
+        else:
+            self._console_logger.info(f"Execution of step {step.name} succeeded.")
+
+    def _get_ready_to_run_steps(self, remaining_steps: list[str], g: nx.DiGraph) -> set[str]:
+        """Identifies and returns the steps that are ready to run.
+
+        This method checks the directed acyclic graph (DAG) to find steps that have no predecessors,
+        indicating that they are ready to be executed. It logs the remaining steps and the steps that
+        are ready to run.
+
+        Args:
+            remaining_steps: A list of step IDs that are yet to be executed.
+            g: The directed acyclic graph representing the pipeline.
+
+        Returns:
+            A set of step IDs that are ready to be executed.
+        """
+        with self._lock:
+            ready_to_run = set([step for step in remaining_steps if g.in_degree(step) == 0])
+            self._console_logger.debug(f"Remaining steps: {remaining_steps}")
+            self._console_logger.debug(f"Ready to run: {ready_to_run}")
+            return ready_to_run
+
+    def _submit_ready_steps(
+        self, ready_to_run: set[str], remaining_steps: list[str], executor: ThreadPoolExecutor, futures: dict
+    ):
+        """Submits the ready-to-run steps to the executor for execution.
+
+        This method takes the steps that are ready to run, removes them from the list of remaining steps,
+        and submits them to the executor for concurrent execution. It also updates the futures dictionary
+        to keep track of the submitted tasks.
+
+        Args:
+            ready_to_run: A set of steps that are ready to be executed.
+            remaining_steps: A list of steps that are yet to be executed.
+            executor: The executor that manages the concurrent execution of steps.
+            futures: A dictionary mapping futures to their corresponding step ID.
+        """
+        with self._lock:
+            for step in ready_to_run:
+                self._console_logger.debug(f"Submitting: {step}")
+                remaining_steps.remove(step)
+                future = executor.submit(self._run_step, step)
+                futures[future] = step
+
+    def _handle_completed_tasks(self, futures, g, remaining_steps):
+        """Handles the completion of tasks in the pipeline.
+
+        This method processes the futures that have completed execution. It removes the corresponding
+        steps from the directed acyclic graph (DAG) and checks if new steps are ready to run. If new
+        steps are ready, it returns True to indicate that the pipeline can continue execution.
+
+        Args:
+            futures: A dictionary mapping futures to their corresponding steps.
+            g: The directed acyclic graph representing the pipeline.
+            remaining_steps: A list of steps that are yet to be executed.
+
+        Returns:
+            True if new steps are ready to run, False otherwise.
+        """
+        # Wait for tasks to complete and free up dependencies
+        for future in as_completed(futures):
+            future.result()  # checks if the run was successful, otherwise throws an error and cancels remaining futures
+            step = futures[future]
+            del futures[future]
+            with self._lock:
+                g.remove_node(step)  # Mark the step as completed by removing it from the graph.
+            if len(set([step for step in remaining_steps if g.in_degree(step) == 0])) > 0:
+                self._console_logger.debug("New steps ready to run")
+                return True
+        self._console_logger.debug("No more steps to run")
+        return False
+
+    def run(self) -> None:
+        """Runs the pipeline by executing each step in the correct order."""
+        g = self._create_graph()
+        remaining_steps = list(g.nodes())
+        self._console_logger.info(f"Pipeline [' {self.name} '] started with {len(remaining_steps)} steps.")
+
+        with ThreadPoolExecutor(max_workers=int(os.environ.get("NESSY_MAX_WORKERS", 1))) as executor:
+            futures: dict = {}
+            try:
+                self._console_logger.debug(f"Remaining steps: {remaining_steps}")
+                while remaining_steps:
+                    ready_to_run = self._get_ready_to_run_steps(remaining_steps, g)
+                    if not ready_to_run:
+                        # If there are still steps to be executed, but all of them have predecessors there
+                        # must be a cyclic dependency in the graph.
+                        self._console_logger.error(
+                            f"Cyclic dependency detected in the pipeline. Remaining steps: {remaining_steps}"
+                        )
+                        raise RuntimeError("Cyclic dependency detected in the pipeline!")
+
+                    self._submit_ready_steps(ready_to_run, remaining_steps, executor, futures)
+
+                    if self._handle_completed_tasks(futures, g, remaining_steps):
+                        continue
+            except RuntimeError as e:
+                self._console_logger.error(f"Pipeline [' {self.name} '] failed due to cyclic dependency: {str(e)}")
+                raise e
+            except Exception as e:
+                self._console_logger.error(f"Pipeline [' {self.name} '] failed: {str(e)}")
+                raise e
+            finally:
+                # ensure that any futures are canceled (if successful, it finished anyway, if error, cancel still running futures)
+                for future in futures:
+                    future.cancel()  # Cancel remaining futures
+                self._graph = self._create_graph()  # recreate the graph after the run
+        self._console_logger.info(f"Pipeline [' {self.name} '] completed successfully.")
+
+    def plot_graph(self, save_path: str | None = None) -> None:
+        """Visualizes the graph of the pipeline using matplotlib.
+
+        Args:
+            save_path: If provided, the graph will be saved to this path. Otherwise, it will be shown.
+        """
+        pos = nx.spring_layout(self._graph)  # Position steps (nodes) using the spring layout
+        plt.figure(figsize=(12, 8))
+        nx.draw(
+            self._graph,
+            pos,
+            with_labels=True,
+            node_color="lightblue",
+            font_weight="bold",
+            node_size=3000,
+            font_size=10,
+            edge_color="gray",
+        )
+
+        # Draw edge labels if needed
+        edge_labels = nx.get_edge_attributes(self._graph, "label")
+        nx.draw_networkx_edge_labels(self._graph, pos, edge_labels=edge_labels)
+
+        if save_path:
+            plt.savefig(save_path)
+            self._console_logger.info(f"Graph visual saved to {save_path}")
+        else:
+            plt.show()

cloe_nessy/pipeline/pipeline_action.py

@@ -0,0 +1,62 @@
+import logging
+from abc import ABC, ABCMeta, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..logging import LoggerMixin
+from .pipeline_context import PipelineContext
+
+
+@dataclass
+class PipelineActionLogs:
+    """Dataclass defining the pipeline action logs table."""
+
+    table_name: str = "nessy_action_logs"
+    log_type: str = "nessy_action_logs"
+    columns: dict[str, str] = field(
+        default_factory=lambda: {
+            "action_name": "STRING",
+            "message": "STRING",
+        }
+    )
+
+
+class PipelineActionMeta(ABCMeta):
+    """Metaclass for PipelineAction to ensure that all subclasses have a 'name' attribute."""
+
+    def __init__(cls, name, bases, dct):
+        if cls.__name__ != "PipelineAction" and "name" not in dct:
+            raise TypeError(f"Class {name} is missing required 'name' attribute")
+        super().__init__(name, bases, dct)
+
+
+class PipelineAction(ABC, LoggerMixin, metaclass=PipelineActionMeta):
+    """Models the operation being executed against an Input.
+
+    Attributes:
+        name: The name of the action.
+    """
+
+    name: str
+
+    def __init__(self, tabular_logger: logging.Logger | None = None) -> None:
+        """Initializes the PipelineAction object.
+
+        Args:
+            tabular_logger: The tabular logger to use for dependency injection.
+        """
+        self._console_logger = self.get_console_logger()
+        self._tabular_logger = tabular_logger or self.get_tabular_logger(
+            logger_name="Tabular:PipelineAction",
+            uc_table_name=PipelineActionLogs().table_name,
+            uc_table_columns=PipelineActionLogs().columns,
+            log_type=PipelineActionLogs().log_type,
+        )
+
+    def __str__(self) -> str:
+        return self.__class__.__name__
+
+    @abstractmethod
+    def run(self, context: PipelineContext, **kwargs: Any) -> PipelineContext:
+        """Execute the pipeline action."""
+        pass

cloe_nessy/pipeline/pipeline_config.py

@@ -0,0 +1,92 @@
+import logging
+from collections import OrderedDict
+from typing import Any
+
+from pydantic import BaseModel, Field, ValidationError, model_validator
+
+
+class PipelineConfigBaseModel(BaseModel):
+    """The base model for Pipeline Config objects."""
+
+    @classmethod
+    def metadata_to_instance(cls, data: dict) -> Any:
+        """Parses a Dictionary to an instance.
+
+        Args:
+            data: The data to parse.
+
+        Returns:
+            An instance and potentially a list of errors.
+        """
+        errors = []
+        try:
+            instance = cls(**data)
+        except ValidationError as e:
+            instance = None
+            errors.append(e)
+        if errors:
+            PipelineConfig.handle_validation_errors(errors)
+        return instance
+
+    @staticmethod
+    def handle_validation_errors(errors: list[ValidationError]) -> None:
+        """Cleanly prints Pydantic validation errors and raises a ValueError.
+
+        Args:
+            errors: A list of Pydantic validation errors.
+
+        Raises:
+            ValueError: If any validation errors occurred.
+        """
+        logger = logging.getLogger(__name__)
+        for error in errors:
+            if isinstance(error, ValidationError):
+                logger.error(f"Validation errors for {error.title}:")
+                for err in error.errors():
+                    loc = ".".join(map(str, err["loc"]))
+                    msg = err["msg"]
+                    err_type = err["type"]
+                    input_value = err.get("input", "N/A")
+                    logger.error(f"  Location: {loc}")
+                    logger.error(f"  Error Message: {msg}")
+                    logger.error(f"  Error Type: {err_type}")
+                    logger.error(f"  Input Value: {input_value}")
+                    logger.error(f"  Further information: {err.get('ctx', {}).get('url', 'N/A')}")
+                    logger.error("")
+            else:
+                logger.error(error)
+        if errors:
+            raise ValueError(f"Validation errors occurred: {errors}")
+
+
+class PipelineActionConfig(PipelineConfigBaseModel):
+    """This class stores the configuration for a pipeline action."""
+
+    name: str
+
+    @model_validator(mode="before")
+    def validate_action(cls, v):
+        """The Pipeline Action must be a valid action type."""
+        # This validation was removed in favor of custom validations in YAML
+        # pipeline definitions.
+        # if v not in PipelineActionType.__members__: # noqa: ERA001
+        #     raise ValueError(f"Action '{v}' is not a valid action.") # noqa: ERA001
+        action_config = {"name": v}
+        return action_config
+
+
+class PipelineStepConfig(PipelineConfigBaseModel):
+    """This class stores the configuration for a pipeline step."""
+
+    action: PipelineActionConfig
+    is_successor: bool = True
+    context: str | None = None
+    table_metadata: str | None = None
+    options: dict = Field(default_factory=dict)
+
+
+class PipelineConfig(PipelineConfigBaseModel):
+    """This class stores the configuration for a pipeline."""
+
+    name: str
+    steps: OrderedDict[str, PipelineStepConfig]

cloe_nessy/pipeline/pipeline_context.py

@@ -0,0 +1,56 @@
+from typing import Any
+
+from pyspark.sql import DataFrame
+
+from ..models import Table
+
+
+class PipelineContext:
+    """A class that models the context of a pipeline.
+
+    The context consists of Table Metadata (the Table definition) and the actual data
+    as a DataFrame.
+
+    Attributes:
+        table_metadata: The Nessy-Table definition.
+        data: The data of the context.
+        runtime_info: Additional runtime information, e.g. streaming status.
+        status: The status of the context. Can be "initialized", "successful" or
+            "failed".
+
+    Note:
+        This is not a pydantic class, because Fabric does not support the type ConnectDataFrame.
+    """
+
+    def __init__(
+        self,
+        table_metadata: Table | None = None,
+        data: DataFrame | None = None,
+        runtime_info: dict[str, Any] | None = None,
+        status: str = "initialized",
+    ) -> None:
+        self.table_metadata = table_metadata
+        self.data = data
+        self.runtime_info = runtime_info if runtime_info is not None else {}
+        self.status = status
+
+    def from_existing(
+        self,
+        table_metadata: Table | None = None,
+        data: DataFrame | None = None,
+        runtime_info: dict[str, Any] | None = None,
+    ) -> "PipelineContext":
+        """Creates a new PipelineContext from an existing one.
+
+        Args:
+            table_metadata: The metadata of the new context.
+            data: The data of the new context.
+            runtime_info: The runtime_info of the new context.
+
+        Returns:
+            The new PipelineContext.
+        """
+        final_metadata = table_metadata or self.table_metadata
+        final_data = data or self.data
+        final_runtime_info = runtime_info or self.runtime_info or {}
+        return PipelineContext(table_metadata=final_metadata, data=final_data, runtime_info=final_runtime_info)

cloe_nessy/pipeline/pipeline_parsing_service.py

@@ -0,0 +1,156 @@
+import os
+import re
+from collections import OrderedDict
+from enum import Enum
+from pathlib import Path
+
+import yaml
+
+from ..logging import LoggerMixin
+from ..session import SessionManager
+from .actions import PipelineActionType, pipeline_actions
+from .pipeline import Pipeline
+from .pipeline_config import PipelineConfig
+from .pipeline_step import PipelineStep
+
+
+class PipelineParsingService:
+    """A service class that parses a YAML document or string into a Pipeline object."""
+
+    def __init__(self, custom_actions=None):
+        if custom_actions is not None:
+            for action in custom_actions:
+                self.register_pipeline_action(action)
+
+    @staticmethod
+    def register_pipeline_action(pipeline_action_class):
+        """Registers a custom pipeline action class.
+
+        !!! note
+            Registering an action enables the custom action to be used in the
+            pipeline YAML definition. This is automatically called, when the
+            PipelineParsingService is instantiated with (a list of) custom
+            actions.
+        """
+        console_logger = LoggerMixin().get_console_logger()
+        console_logger.info("Registering custom pipeline action [' %s ']", pipeline_action_class.name)
+        pipeline_actions[pipeline_action_class.name] = pipeline_action_class
+
+        global PipelineActionType
+        PipelineActionType = Enum("PipelineActionType", pipeline_actions)
+
+    @staticmethod
+    def parse(path: Path | None = None, yaml_str: str | None = None) -> Pipeline:
+        """Reads the YAML from a given Path and returns a Pipeline object.
+
+        Args:
+            path: Path to the YAML document.
+            yaml_str: A string that can be parsed in YAML format.
+
+        Raises:
+            ValueError: If neither 'path' nor 'yaml_str' has been provided.
+
+        Returns:
+            Pipeline: The resulting Pipeline instance.
+        """
+        console_logger = LoggerMixin().get_console_logger()
+        if not path and not yaml_str:
+            raise ValueError("Neither 'file_path' nor 'yaml_str' was provided. Please supply one of them.")
+        if path:
+            path_obj = Path(path)
+            with open(path_obj) as f:
+                yaml_str = f.read()
+        if not yaml_str:
+            raise ValueError("YAML content is empty.")
+
+        final_yaml_str = PipelineParsingService._replace_variables(yaml_str)
+        config = yaml.safe_load(final_yaml_str)
+        pipeline_config = PipelineConfig.metadata_to_instance(config)
+        steps = PipelineParsingService._get_steps(pipeline_config.steps)
+        pipeline = Pipeline(name=pipeline_config.name, steps=steps)  # type: ignore
+        console_logger.info("Pipeline [ '%s' ] parsed successfully with %d steps.", pipeline.name, len(pipeline.steps))
+        return pipeline
+
+    @staticmethod
+    def _replace_variables(yaml_str: str) -> str:
+        """Replace variable placeholders in a YAML string.
+
+        Replaces environment variables with the pattern `{{env:var-name}}`. Where
+        the var-name is the name of the environment variable. Replaces secret
+        references with the pattern `{{secret-scope-name:secret-key}}`. Where
+        scope-name is the name of the secret scope and secret-key is the key of
+        the secret.
+
+        Args:
+            yaml_str: A string that can be parsed in YAML format.
+
+        Returns:
+            The same YAML string with environment variable placeholders replaced.
+        """
+        env_var_pattern = r"\{\{env:([^}]+)\}\}"
+        secret_ref_pattern = r"\{\{(?!step|env)([^}]+):([^}]+)\}\}"
+
+        def replace_with_env_var(match):
+            env_var_name = match.group(1)
+            env_var_value = os.getenv(env_var_name)
+            return env_var_value
+
+        def replace_with_secret(match):
+            secret_scope_name = match.group(1)
+            secret_key = match.group(2)
+            return SessionManager.get_utils().secrets.get(scope=secret_scope_name, key=secret_key)
+
+        env_replaced_yaml_string = re.sub(env_var_pattern, replace_with_env_var, yaml_str)
+        final_yaml_string = re.sub(secret_ref_pattern, replace_with_secret, env_replaced_yaml_string)
+        return final_yaml_string
+
+    @staticmethod
+    def _get_steps(step_configs, last_step_name: str | None = None):
+        steps = OrderedDict()
+        for step_name, step_config in step_configs.items():
+            is_successor = step_config.is_successor
+            context_ref = step_config.context
+            if is_successor and not context_ref:
+                context_ref = last_step_name
+            action = PipelineActionType[step_config.action.name].value()
+            step = PipelineStep(
+                name=step_name,
+                action=action,
+                options=step_config.options,
+                _context_ref=context_ref,
+                _table_metadata_ref=step_config.table_metadata,
+            )
+            steps[step.name] = step
+            last_step_name = step_name
+        for step in steps.values():
+            steps[step.name] = PipelineParsingService._replace_step_refs(steps, step)
+        return steps
+
+    @staticmethod
+    def _replace_step_refs(steps: OrderedDict[str, PipelineStep], step: PipelineStep) -> PipelineStep:
+        step_ref_pattern = r"\(\(step:([^)]+)\)\)"
+
+        def _handle_string_value(value: str, option: str):
+            if match := re.match(step_ref_pattern, value):
+                dependency_step_name = match.group(1)
+                dependency_step = steps.get(dependency_step_name)
+                step.options[option] = dependency_step
+                step._predecessors.add(dependency_step_name)
+
+        def _handle_list_value(value: list, option: str):
+            for i, v in enumerate(value):
+                if isinstance(v, str):
+                    if match := re.match(step_ref_pattern, v):
+                        dependency_step_name = match.group(1)
+                        dependency_step = steps.get(dependency_step_name)
+                        step.options[option][i] = dependency_step
+                        step._predecessors.add(dependency_step_name)
+
+        if step.options:
+            for option, value in step.options.items():
+                if isinstance(value, str):
+                    _handle_string_value(value, option)
+                elif isinstance(value, list):
+                    _handle_list_value(value, option)
+
+        return step