cloe-nessy 0.3.5__py3-none-any.whl → 0.3.9__py3-none-any.whl

cloe_nessy/integration/writer/delta_writer/delta_writer_base.py

@@ -0,0 +1,210 @@
+import logging
+from abc import ABC
+from dataclasses import dataclass, field
+
+from pyspark.sql import DataFrame, Row
+from pyspark.sql.functions import col, concat, concat_ws, format_string, lit
+
+from ....object_manager import TableManager
+from ....session import SessionManager
+from ..writer import BaseWriter
+from .delta_table_operation_type import DeltaTableOperationType
+from .exceptions import EmptyDataframeError
+
+
+@dataclass
+class DeltaWriterLogs:
+    """Dataclass defining the delta writer logs table."""
+
+    logger_name = "Tabular:DeltaWriter"
+    log_type: str = "nessy_simple_logs"
+    uc_table_name: str = "nessy_simple_logs"
+    uc_table_columns: dict[str, str] = field(
+        default_factory=lambda: {
+            "message": "STRING",
+        }
+    )
+
+
+@dataclass
+class TableOperationMetricsLogs:
+    """Dataclass defining the table operation metrics logs table."""
+
+    logger_name = "Tabular:TableOperationMetrics"
+    log_type: str = "nessy_table_operation_metrics"
+    uc_table_name: str = "nessy_table_operation_metrics"
+    uc_table_columns: dict[str, str] = field(
+        default_factory=lambda: {
+            "timestamp": "TIMESTAMP",
+            "table_identifier": "STRING",
+            "operation_type": "STRING",
+            "metric": "STRING",
+            "value": "STRING",
+            "user_name": "STRING",
+            "job_id": "STRING",
+            "job_run_id": "STRING",
+            "run_id": "STRING",
+            "notebook_id": "STRING",
+            "cluster_id": "STRING",
+        }
+    )
+
+
+class BaseDeltaWriter(BaseWriter, ABC):
+    """A class for writing DataFrames to Delta tables."""
+
+    def __init__(
+        self,
+        tabular_logger: logging.Logger | None = None,
+        table_operation_metrics_logger: logging.Logger | None = None,
+    ):
+        super().__init__()
+        self._spark = SessionManager.get_spark_session()
+        self._dbutils = SessionManager.get_utils()
+        self._table_operation_metrics_logger = table_operation_metrics_logger or self.get_tabular_logger(
+            **DeltaWriterLogs().__dict__
+        )
+        self.table_manager = TableManager()
+        self._tabular_logger = tabular_logger or self.get_tabular_logger(**DeltaWriterLogs().__dict__)
+
+    def _delta_operation_log(self, table_identifier: str, operation_type: DeltaTableOperationType) -> dict:
+        """Returns a dictionary containing the most recent delta log of a Delta table for given operation type.
+
+        Args:
+            table_identifier: The identifier of the Delta table in the format 'catalog.schema.table'.
+            operation_type: A DeltaTableOperationType
+                object specifying the type of operation for which metrics should
+                be retrieved (UPDATE, DELETE, MERGE or WRITE).
+
+        Returns:
+            dict: A dictionary containing the operation log.
+        """
+        delta_history = self._spark.sql(f"DESCRIBE HISTORY {table_identifier}")
+
+        try:
+            operation_log: dict = (
+                delta_history.filter(col("operation") == operation_type.name.replace("_", " "))
+                .orderBy("version", ascending=False)
+                .collect()[0]
+                .asDict()
+            )
+        except IndexError:
+            operation_log = {}
+
+        return operation_log
+
+    def _report_delta_table_operation_metrics(
+        self, table_identifier: str, operation_type: DeltaTableOperationType
+    ) -> None:
+        """Logs the most recent metrics of a Delta table for given operation type.
+
+        Args:
+            table_identifier: The identifier of the Delta table in the format 'catalog.schema.table'.
+            operation_type: A DeltaTableOperationType object specifying the type
+                of operation for which metrics should be retrieved (UPDATE, DELETE,
+                MERGE or WRITE).
+        """
+        operation_log = self._delta_operation_log(table_identifier, operation_type)
+        timestamp = operation_log.get("timestamp")
+        user_name = operation_log.get("userName")
+        job_id = (operation_log.get("job") or Row(jobId=None)).asDict().get("jobId")
+        job_run_id = (operation_log.get("job") or Row(jobRunId=None)).asDict().get("jobRunId")
+        run_id = (operation_log.get("job") or Row(runId=None)).asDict().get("runId")
+        notebook_id = (operation_log.get("notebook") or Row(notebook_id=None)).asDict().get("notebookId")
+        cluster_id = operation_log.get("clusterId")
+        affected_rows = {
+            k: v for k, v in operation_log.get("operationMetrics", {}).items() if k in operation_type.value
+        }
+        for metric, value in affected_rows.items():
+            log_message = f"""timestamp: {timestamp} |
+                                table_identifier: {table_identifier} |
+                                operation_type: {operation_type.name} |
+                                metric_name: {metric} |
+                                metric_value: {value} |
+                                user_name: {user_name} |
+                                job_id: {job_id} |
+                                job_run_id: {job_run_id} |
+                                run_id: {run_id} |
+                                notebook_id: {notebook_id} |
+                                cluster_id: {cluster_id}
+                            """
+            self._table_operation_metrics_logger.info(log_message)
+
+    @staticmethod
+    def _merge_match_conditions(columns: list[str]) -> str:
+        """Merges match conditions of the given columns into a single string.
+
+        This function is used to generate an SQL query to match rows between two tables based on
+        the specified columns.
+
+        Args:
+            columns: A list of strings representing the names of the columns to match.
+
+        Returns:
+            A string containing the match conditions, separated by " AND "
+
+        Example:
+            ```python
+            _merge_match_conditions(["column1", "column2"]) # "target.column1 <=> source.column1 AND target.column2 <=> source.column2"
+            ```
+        """
+        return " AND ".join([f"target.`{c}` <=> source.`{c}`" for c in columns])
+
+    @staticmethod
+    def _partition_pruning_conditions(df, partition_cols: list[str] | None) -> str:
+        """Generates partition pruning conditions for an SQL query.
+
+        This function is used to optimize the performance of an SQL query by only scanning the
+        necessary partitions in a table, based on the specified partition columns and the data
+        in a Spark dataframe.
+
+        Args:
+            df: A Spark dataframe containing the data to generate the partition pruning
+                conditions from.
+            partition_cols: A list of strings representing the names of the
+                partition columns.
+
+        Returns:
+            A string, representing the partition pruning conditions.
+
+        Example:
+            ```python
+            _partition_pruning_conditions(df, ["column1", "column2"])
+            "(target.column1 = 'value1' AND target.column2 = 'value2') OR (target.column1 = 'value3'
+                AND target.column2 = 'value4')"
+            ```
+        """
+        if not partition_cols:
+            return ""
+        pruning_conditions = (
+            df.select(*partition_cols)
+            .distinct()
+            .select([format_string("target.`%s` = '%s'", lit(c), col(c)).alias(c) for c in partition_cols])
+            .withColumn("result", concat(lit("("), concat_ws(" AND ", *partition_cols), lit(")")))
+            .select("result")
+            .toPandas()
+            .result.str.cat(sep=" OR ")
+        )
+        pruning_conditions = "(" + pruning_conditions + ")"
+
+        return str(pruning_conditions)
+
+    def _empty_dataframe_check(self, df: DataFrame, ignore_empty_df: bool) -> bool | None:
+        """Checks if a DataFrame is empty and raises an exception if it is not expected to be empty.
+
+        Args:
+            df: The DataFrame to check for emptiness.
+            ignore_empty_df: If True, the function will return without raising
+                an exception if the DataFrame is empty. If False, an EmptyDataframeException
+                will be raised.
+
+        Raises:
+            EmptyDataframeException: If the DataFrame is empty and ignore_empty_df is False.
+        """
+        if df.isEmpty():
+            if ignore_empty_df:
+                return True
+            raise EmptyDataframeError(
+                "EMPTY DATAFRAME, nothing to write. If this is expected, consider setting `ignore_empty_df` to True.",
+            )
+        return None

cloe_nessy/integration/writer/delta_writer/exceptions.py

@@ -0,0 +1,4 @@
+class EmptyDataframeError(Exception):
+    """When a dataframe is empty when it should not be."""
+
+    pass

cloe_nessy/integration/writer/file_writer.py

@@ -0,0 +1,132 @@
+from typing import Any
+
+from pyspark.sql import DataFrame, DataFrameWriter
+from pyspark.sql.streaming import DataStreamWriter
+
+from .writer import BaseWriter
+
+
+class FileWriter(BaseWriter):
+    """Utility class for writing a DataFrame to a file."""
+
+    def __init__(self):
+        super().__init__()
+
+    def _get_writer(self, df: DataFrame) -> DataFrameWriter:
+        """Returns a DataFrameWriter."""
+        return df.write
+
+    def _get_stream_writer(self, df: DataFrame) -> DataStreamWriter:
+        """Returns a DataStreamWriter."""
+        return df.writeStream
+
+    def _log_operation(self, location: str, status: str, error: str | None = None):
+        """Logs the status of an operation."""
+        if status == "start":
+            self._console_logger.info(f"Starting to write to {location}")
+        elif status == "succeeded":
+            self._console_logger.info(f"Successfully wrote to {location}")
+        elif status == "failed":
+            self._console_logger.error(f"Failed to write to {location}: {error}")
+
+    def _validate_trigger(self, trigger_dict: dict[str, Any]):
+        """Validates the trigger type."""
+        triggers = ["processingTime", "once", "continuous", "availableNow"]
+        if not any(trigger in trigger_dict for trigger in triggers):
+            raise ValueError(f"Invalid trigger type. Supported types are: {', '.join(triggers)}")
+
+    def write_stream(
+        self,
+        data_frame: DataFrame | None = None,
+        location: str | None = None,
+        format: str = "delta",
+        checkpoint_location: str | None = None,
+        partition_cols: list[str] | None = None,
+        mode: str = "append",
+        trigger_dict: dict | None = None,
+        options: dict[str, Any] | None = None,
+        await_termination: bool = False,
+        **_: Any,
+    ):
+        """Writes a dataframe to specified location in specified format as a stream.
+
+        Args:
+            data_frame: The DataFrame to write.
+            location: The location to write the DataFrame to.
+            format: The format to write the DataFrame in.
+            checkpoint_location: Location of checkpoint. If None, defaults
+                to the location of the table being written, with '_checkpoint_'
+                added before the name.
+            partition_cols: Columns to partition by.
+            mode: The write mode.
+            trigger_dict: A dictionary specifying the trigger configuration for the streaming query.
+                Supported keys include:
+
+                - "processingTime": Specifies a time interval (e.g., "10 seconds") for micro-batch processing.
+                - "once": Processes all available data once and then stops.
+                - "continuous": Specifies a time interval (e.g., "1 second") for continuous processing.
+                - "availableNow": Processes all available data immediately and then stops.
+
+                If nothing is provided, the default is {"availableNow": True}.
+            options: Additional options for writing.
+            await_termination: If True, the function will wait for the streaming
+                query to finish before returning. This is useful for ensuring that
+                the data has been fully written before proceeding with other
+                operations.
+        """
+        if not location or not data_frame:
+            raise ValueError("Location and data_frame are required for streaming.")
+
+        self._log_operation(location, "start")
+        try:
+            options = options or {}
+            trigger_dict = trigger_dict or {"availableNow": True}
+            checkpoint_location = self._get_checkpoint_location(location, checkpoint_location)
+            self._validate_trigger(trigger_dict)
+            stream_writer = self._get_stream_writer(data_frame)
+
+            stream_writer.trigger(**trigger_dict)
+            stream_writer.format(format)
+            stream_writer.outputMode(mode)
+            stream_writer.options(**options).option("checkpointLocation", checkpoint_location)
+            if partition_cols:
+                stream_writer.partitionBy(partition_cols)
+
+            query = stream_writer.start(location)
+            if await_termination is True:
+                query.awaitTermination()
+        except Exception as e:
+            self._log_operation(location, "failed", str(e))
+            raise e
+        else:
+            self._log_operation(location, "succeeded")
+
+    def write(
+        self,
+        data_frame: DataFrame,
+        location: str | None = None,
+        format: str = "delta",
+        partition_cols: list[str] | None = None,
+        mode: str = "append",
+        options: dict[str, Any] | None = None,
+        **_: Any,
+    ):
+        """Writes a dataframe to specified location in specified format."""
+        if not location:
+            raise ValueError("Location is required for writing to file.")
+
+        self._log_operation(location, "start")
+        try:
+            options = options or {}
+            df_writer = self._get_writer(data_frame)
+            df_writer.format(format)
+            df_writer.mode(mode)
+            if partition_cols:
+                df_writer.partitionBy(partition_cols)
+            df_writer.options(**options)
+            df_writer.save(str(location))
+        except Exception as e:
+            self._log_operation(location, "failed", str(e))
+            raise e
+        else:
+            self._log_operation(location, "succeeded")

cloe_nessy/integration/writer/writer.py

@@ -0,0 +1,54 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+from pyspark.sql import DataFrame
+
+from ...logging import LoggerMixin
+
+
+class BaseWriter(ABC, LoggerMixin):
+    """BaseWriter class to write data."""
+
+    def __init__(self):
+        self._console_logger = self.get_console_logger()
+
+    @abstractmethod
+    def write_stream(self, **kwargs: Any):
+        """Writes a DataFrame stream."""
+        pass
+
+    @abstractmethod
+    def write(
+        self,
+        data_frame: DataFrame,
+        **kwargs: Any,
+    ):
+        """Writes a DataFrame."""
+        pass
+
+    def log_operation(self, operation: str, identifier: str | Path, status: str, error: str = ""):
+        """Logs the metrics for one operation on the given identifier.
+
+        Args:
+            operation: Describes the type of operation, e.g. 'read_api'.
+            identifier: An identifier for the object that's being interacted with.
+            status: The status of the operation. Must be one of "start", "failed", "succeeded".
+            error: The error message, if any. Defaults to ''.
+        """
+        self._console_logger.info(
+            "operation:%s | identifier:%s | status:%s | error:%s",
+            operation,
+            identifier,
+            status,
+            error,
+        )
+
+    def _get_checkpoint_location(self, location: str, checkpoint_location: str | None) -> str:
+        """Generates the checkpoint location if not provided."""
+        if checkpoint_location is None:
+            location_path = Path(location)
+            checkpoint_location = str(location_path.parent / f"_checkpoint_{location_path.name}").replace(
+                "abfss:/", "abfss://"
+            )
+        return checkpoint_location

cloe_nessy/models/__init__.py

@@ -1,13 +1,17 @@
+from .catalog import Catalog
 from .column import Column
 from .constraint import Constraint
 from .foreign_key import ForeignKey
 from .schema import Schema
 from .table import Table
+from .volume import Volume
 
 __all__ = [
+    "Catalog",
     "Column",
     "Constraint",
     "Table",
     "Schema",
     "ForeignKey",
+    "Volume",
 ]

cloe_nessy/models/adapter/__init__.py

@@ -0,0 +1,3 @@
+from .unity_catalog_adapter import UnityCatalogAdapter
+
+__all__ = ["UnityCatalogAdapter"]