cloe-nessy 0.3.8__py3-none-any.whl → 0.3.10__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/adapter/unity_catalog_adapter.py

@@ -79,7 +79,11 @@ class UnityCatalogAdapter(LoggerMixin):
 
         for schema in schemas_df:
             schemas.append(
-                Schema(name=schema["schema_name"], catalog=catalog, comment=schema["comment"]),
+                Schema(
+                    name=schema["schema_name"],
+                    catalog=catalog,
+                    comment=schema["comment"],
+                ),
             )
         return schemas
 

cloe_nessy/models/schema.py

@@ -49,7 +49,7 @@ class Schema(ReadInstancesMixin):
                 instance_path=processed_instance_path.parents[0] / table_dir_name,
                 catalog_name=schema.catalog,
                 schema_name=schema.name,
-                schema_storage_path=Path(schema.storage_path),
+                schema_storage_path=schema.storage_path,
                 fail_on_missing_subfolder=fail_on_missing_subfolder,
             )
             schema.tables = tables

cloe_nessy/models/table.py

@@ -4,7 +4,13 @@ from typing import Any, Self
 import yaml
 import yaml.scanner
 from jinja2 import TemplateNotFound
-from pydantic import Field, ValidationError, ValidationInfo, field_validator, model_validator
+from pydantic import (
+    Field,
+    ValidationError,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
 
 from ..logging import LoggerMixin
 from ..utils.file_and_directory_handler import process_path
@@ -28,7 +34,7 @@ class Table(TemplateLoaderMixin, ReadInstancesMixin, LoggerMixin):
     properties: dict[str, str] = Field(default_factory=dict)
     constraints: list[Constraint] = Field(default_factory=list)
     foreign_keys: list[ForeignKey] = Field(default_factory=list)
-    storage_path: Path | None = None
+    storage_path: str | None = None
     business_properties: dict[str, str] = Field(default_factory=dict)
     comment: str | None = None
     data_source_format: str | None = None
@@ -93,6 +99,7 @@ class Table(TemplateLoaderMixin, ReadInstancesMixin, LoggerMixin):
         """If is_external is set to True, storage_path has to be set."""
         if table.is_external and table.storage_path is None:
             raise ValueError("is_external cannot be true while storage_path is None.")
+        return table
 
     @classmethod
     def read_instances_from_directory(
@@ -154,7 +161,7 @@ class Table(TemplateLoaderMixin, ReadInstancesMixin, LoggerMixin):
             sub_errors: list[ValidationErrorType] = []
             if instance_file.is_file() and instance_file.suffix in (".yaml", ".yml"):
                 instance, sub_errors = cls.read_instance_from_file(
-                    instance_file, catalog_name, schema_name, schema_storage_path
+                    instance_file, catalog_name, schema_name, str(schema_storage_path)
                 )
                 instances += [] if instance is None else [instance]
             errors += sub_errors
@@ -206,9 +213,9 @@ class Table(TemplateLoaderMixin, ReadInstancesMixin, LoggerMixin):
                 data["identifier"] = f"{catalog_name}.{schema_name}.{data['name']}"
                 if data.get("is_external"):
                     if storage_path := data.get("storage_path"):
-                        data["storage_path"] = Path(storage_path)
+                        data["storage_path"] = storage_path
                     elif schema_storage_path:
-                        data["storage_path"] = schema_storage_path / data["name"]
+                        data["storage_path"] = (schema_storage_path / data["name"]).as_posix()
                     else:
                         raise ValueError(
                             f"Neither storage path nor schema storage path of table {data['name']} has been provided."
@@ -216,22 +223,37 @@ class Table(TemplateLoaderMixin, ReadInstancesMixin, LoggerMixin):
 
                 instance, sub_errors = cls.metadata_to_instance(data)
                 errors += sub_errors
-        except (ValidationError, yaml.parser.ParserError, yaml.scanner.ScannerError) as e:
+        except (
+            ValidationError,
+            yaml.parser.ParserError,
+            yaml.scanner.ScannerError,
+        ) as e:
             instance = None
             errors.append(e)
         return instance, errors
 
     def get_create_statement(
         self,
-        templates: Path = Path("./src/cloe_nessy/models/templates/"),
-        template_name: str = "create_table.sql.j2",
         replace: bool = True,
     ):
-        """Get the create statement for the Table."""
+        """Get the create statement for the Table.
+
+        Args:
+            replace: Whether to use the REPLACE statement or not.
+
+        Returns:
+            The rendered create statement for the Table.
+
+        Raises:
+            TemplateNotFound: If the template file is not found.
+        """
+        templates = Path(__file__).parent / "templates"
+        template_name = "create_table.sql.j2"
+
         try:
             template = self.get_template(templates, template_name)
         except TemplateNotFound as err:
-            self._console_logger.error(f"Template [ {template_name} ] not found.")
+            self._console_logger.error("Template [ {template_name} ] not found.")
             raise err
         render = template.render(table=self, replace=replace)
         return render

cloe_nessy/models/volume.py

@@ -54,14 +54,23 @@ class Volume(TemplateLoaderMixin, LoggerMixin, BaseModel):
 
     def get_create_statement(
         self,
-        templates: Path = Path("./src/cloe_nessy/models/templates/"),
-        template_name: str = "create_volume.sql.j2",
+        if_not_exists: bool = True,
     ):
-        """Get the create statement for the Volume."""
+        """Get the create statement for the Volume.
+
+        Args:
+            if_not_exists: Whether to include the IF NOT EXISTS clause in the create statement
+
+        Returns:
+                The rendered create statement as a string.
+        """
+        template_name: str = "create_volume.sql.j2"
+        templates = Path(__file__).parent / "templates"
+
         try:
             template = self.get_template(templates, template_name)
         except TemplateNotFound as err:
             self._console_logger.error(f"Template [ {template_name} ] not found.")
             raise err
-        render = template.render(volume=self)
+        render = template.render(volume=self, if_not_exists=if_not_exists)
         return render