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

cloe_nessy/object_manager/table_manager.py

@@ -1,58 +1,288 @@
+import functools
+import logging
+from dataclasses import dataclass, field
+
+from delta import DeltaTable  # type: ignore
+
 from ..logging import LoggerMixin
+from ..models import Table
 from ..session import SessionManager
 
 
+@dataclass
+class TableManagerLogs:
+    """Dataclass defining the table manager logs table."""
+
+    logger_name = "Tabular:TableManager"
+    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",
+        }
+    )
+
+
+def table_log_decorator(operation: str):
+    """Creates a decorator that logs the start, failure (if any), and completion of a table operation.
+
+    The created decorator wraps a function that performs an operation on a table. The decorator logs
+    the start of the operation, calls the original function, logs if there was an exception, and logs
+    the completion of the operation. Functions that are wrapped must support the self._table_logger
+    attribute.
+
+    Args:
+        operation: The name of the operation to be logged. This will be included in the log messages.
+
+    Returns:
+        inner_decorator: A decorator that can be used to wrap a function that performs an operation on a table.
+
+    Example:
+        ```python
+        @table_log_decorator(operation='delete_physical_data_for_table')
+        def _delete_physical_data(self, table_identifier: str):
+            self._dbutils.fs.rm(table_location, recurse=True)
+        ```
+    """
+
+    def inner_decorator(func):
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            table_identifier = kwargs.get("table_identifier") or kwargs.get("table") or args[0]
+            if isinstance(table_identifier, Table):
+                table_identifier = table_identifier.identifier
+            self._tabular_logger.info(
+                "operation:%s | identifier:%s | status:start | error:''",
+                operation,
+                table_identifier,
+            )
+            try:
+                func(self, *args, **kwargs)
+            except Exception as e:
+                self._tabular_logger.error(
+                    "operation:%s | identifier:%s | status:failed | error:%s",
+                    operation,
+                    table_identifier,
+                    e,
+                )
+                raise e
+            else:
+                self._tabular_logger.info(
+                    "operation:%s | identifier:%s | status:completed | error:''",
+                    operation,
+                    table_identifier,
+                )
+
+        return wrapper
+
+    return inner_decorator
+
+
 class TableManager(LoggerMixin):
-    """TableManager class for managing tables in the catalog."""
+    """TableManager class for managing tables."""
 
-    def __init__(self):
+    def __init__(self, tabular_logger: logging.Logger | None = None):
         self._spark = SessionManager.get_spark_session()
-        self._utils = SessionManager.get_utils()
         self._console_logger = self.get_console_logger()
         self._console_logger.debug("TableManager initialized...")
-        self._tabular_logger = self.get_tabular_logger(uc_table_name="TableManager")
+        self._tabular_logger = tabular_logger or self.get_tabular_logger(**TableManagerLogs().__dict__)
         self._tabular_logger.debug("message:TableManager initialized.")
 
-    @staticmethod
-    def create_table():
-        """Create a table in the catalog."""
-        raise NotImplementedError
+    @table_log_decorator(operation="create")
+    def create_table(
+        self,
+        table: Table,
+        ignore_if_exists: bool = False,
+        replace: bool = False,
+    ) -> None:
+        """Creates a Table in the catalog.
 
-    def drop_table(self, table_identifier: str, delete_physical_data: bool = False):
+        Args:
+            table: A Table object representing the Delta table.
+            ignore_if_exists: If set to True, the function will return early
+                without doing anything if the table already exists.
+            replace: If set to True, the function will replace the table if it
+                already exists.
+        """
+        if ignore_if_exists and self.table_exists(table):
+            return
+        self._console_logger.info(f"Creating table: {table.identifier}")
+        self._spark.sql(f"USE CATALOG {table.catalog};")
+        self._spark.sql(f"USE SCHEMA {table.schema};")
+        for statement in table.get_create_statement(replace=replace).split(";"):
+            if statement and statement != "\n":
+                self._spark.sql(statement)
+
+    def drop_table(
+        self,
+        table: Table | None = None,
+        storage_location: str | None = None,
+        table_identifier: str | None = None,
+        delete_physical_data: bool = False,
+    ):
         """Deletes a Table. For security reasons you are forced to pass the table_name.
 
         If delete_physical_data is True the actual physical data on the ADLS will be deleted.
         Use with caution!
 
         Args:
+            table: The Table object representing the Delta table.
+            storage_location: The location of the Delta table on the ADLS.
             table_identifier: The table identifier in the catalog. Must be in the format 'catalog.schema.table'.
             delete_physical_data: If set to True, deletes not only the metadata
                                   within the Catalog but also the physical data.
+
+        Raises:
+            ValueError: If neither table nor table_identifier is provided, or if both are provided.
+            ValueError: If the table storage path is not provided by the table object.
         """
         self._console_logger.info(f"Deleting table [ '{table_identifier}' ] ...")
-        if not isinstance(table_identifier, str):
-            raise NotImplementedError("table_identifier must be a string, can be a Table object in the future.")
-
+        if table is not None and (table_identifier is not None or storage_location is not None):
+            raise ValueError("Either table or table_identifier and storage_location must be provided, but not both.")
+        if table is not None:
+            table_identifier = table.identifier
+            storage_location = str(table.storage_path)
         if delete_physical_data:
-            self._delete_physical_data()
-        self.drop_table_from_catalog(table_identifier)
+            self._delete_physical_data(location=storage_location)
+        self.drop_table_from_catalog(table_identifier=table_identifier)
 
-    def drop_table_from_catalog(self, table_identifier: str) -> None:
+    def drop_table_from_catalog(self, table_identifier: str | None = None, table: Table | None = None) -> None:
         """Removes a table from the catalog. Physical data is retained.
 
         Args:
             table_identifier: The table identifier in the catalog. Must be in the format 'catalog.schema.table'.
+            table: The Table object representing the Delta table.
+
+        Raises:
+            ValueError: If neither table nor table_identifier is provided, or if both are provided.
         """
+        if (table is None and table_identifier is None) or (table is not None and table_identifier is not None):
+            raise ValueError("Either table or table_identifier must be provided, but not both.")
+        if table is not None:
+            table_identifier = table.identifier
         self._console_logger.info(f"... deleting table [ '{table_identifier}' ] from Catalog.")
-        if not isinstance(table_identifier, str):
-            raise NotImplementedError("table_identifier must be a string, can be a Table object in the future.")
         self._spark.sql(f"DROP TABLE IF EXISTS {table_identifier};")
 
-    def _delete_physical_data(self):
+    def _delete_physical_data(self, table: Table | None = None, location: str | None = None):
         """Removes the physical data on the ADLS for the location of this table.
 
+        Args:
+            table: The Table object representing the Delta table to be deleted.
+            location: The location of the Delta table to be deleted.
+
         Raises:
-            NotImplementedError: This can be implemented, once a Table object is available.
+            ValueError: If neither table nor location is provided, or if both are provided.
+            ValueError: If the table storage path is not provided by the table object.
         """
-        self._console_logger.info("... deleting physical data for table [ '' ] from Catalog.")
-        raise NotImplementedError("This can be implemented, once a Table object is available.")
+        if (table is None and location is None) or (table is not None and location is not None):
+            raise ValueError("Either table or location must be provided, but not both.")
+        if table is not None:
+            if table.storage_path is None:
+                raise ValueError("Table storage path must be provided.")
+            location = str(table.storage_path)
+        SessionManager.get_utils().fs.rm(location, recurse=True)
+        self._console_logger.info("... deleting physical data.")
+
+    def get_delta_table(self, table: Table | None = None, location: str | None = None, spark=None) -> DeltaTable:
+        """Get the DeltaTable object from the Table objects location or a location string.
+
+        Args:
+            table: A Table object representing the Delta table.
+            location: A string representing the table location.
+            spark: An optional Spark session. If not provided, the current Spark session will be used.
+
+        Returns:
+            The DeltaTable object corresponding to the given Table object or location string.
+
+        Raises:
+            ValueError: If neither table nor location is provided, or if both are provided.
+        """
+        if (table is None and location is None) or (table is not None and location is not None):
+            raise ValueError("Either table or location must be provided, but not both.")
+
+        if table is not None:
+            location = str(table.storage_path)
+        self._console_logger.info(f"Getting DeltaTable object for location: {location}")
+        return DeltaTable.forPath(spark or self._spark, str(location))
+
+    def table_exists(self, table: Table | None = None, table_identifier: str | None = None) -> bool:
+        """Checks if a table exists in the catalog.
+
+        Args:
+            table: A Table object representing the Delta table.
+            table_identifier: A string representing the table identifier in the format 'catalog.schema.table'.
+
+        Returns:
+            True if the table exists, else False.
+
+        Raises:
+            ValueError: If neither table nor table_identifier is provided, or if both are provided.
+            ValueError: If the table_identifier is not in the format 'catalog.schema.table'.
+        """
+        if (table is None and table_identifier is None) or (table is not None and table_identifier is not None):
+            raise ValueError("Either table or table_identifier must be provided, but not both.")
+
+        if table is not None:
+            catalog = table.catalog
+            schema = table.schema
+            table_name = table.name
+        else:
+            assert table_identifier is not None, "table_identifier must be provided."
+            catalog, schema, table_name = table_identifier.split(".")
+            if not all([catalog, schema, table_name]):
+                raise ValueError("Invalid table identifier format. Expected 'catalog.schema.table'.")
+
+        query_result = self._spark.sql(
+            f"""
+                SELECT 1 FROM {catalog}.information_schema.tables
+                WHERE table_name = '{table_name}'
+                AND table_schema = '{schema}'
+                LIMIT 1""",
+        )
+        result = query_result.count() > 0
+        self._console_logger.info(f"Table [ '{catalog}.{schema}.{table_name}' ] exists: {result}")
+        return result is True
+
+    @table_log_decorator(operation="refresh")
+    def refresh_table(self, table: Table | None = None, table_identifier: str | None = None):
+        """Refreshes the metadata of a Delta table.
+
+        Args:
+            table: A Table object representing the Delta table.
+            table_identifier: The identifier of the Delta table in the format 'catalog.schema.table'.
+
+        Raises:
+            ValueError: If neither table nor table_identifier is provided, or if both are provided.
+        """
+        if (table is None and table_identifier is None) or (table is not None and table_identifier is not None):
+            raise ValueError("Either table or table_identifier must be provided, but not both.")
+
+        if table is not None:
+            table_identifier = f"{table.catalog}.{table.schema}.{table.name}"
+
+        self._console_logger.info(f"Refreshing table: {table_identifier}")
+        self._spark.sql(f"REFRESH TABLE {table_identifier};")
+
+    @table_log_decorator(operation="truncate")
+    def truncate_table(
+        self,
+        table: Table | None = None,
+        table_identifier: str | None = None,
+    ):
+        """Truncates a table.
+
+        Args:
+            table: A Table object representing the Delta table.
+            table_identifier: The identifier of the Delta table in the format 'catalog.schema.table'.
+
+        Raises:
+            ValueError: If neither table nor table_identifier is provided, or if both are provided.
+        """
+        if (table is None and table_identifier is None) or (table is not None and table_identifier is not None):
+            raise ValueError("Either table or table_identifier must be provided, but not both.")
+
+        if table is not None:
+            table_identifier = table.escaped_identifier
+
+        self._console_logger.info(f"Truncating table: {table_identifier}")
+        self._spark.sql(f"TRUNCATE TABLE {table_identifier};")

cloe_nessy/object_manager/volume_manager.py

@@ -0,0 +1,70 @@
+import logging
+
+from ..logging import LoggerMixin
+from ..models import Volume
+from ..session import SessionManager
+
+
+class VolumeManager(LoggerMixin):
+    """VolumeManager class for managing volumes."""
+
+    def __init__(self, console_logger: logging.Logger | None = None):
+        self._spark = SessionManager.get_spark_session()
+        self._console_logger = console_logger or self.get_console_logger()
+
+    def create_volume(self, volume: Volume):
+        """Creates a Volume in the catalog.
+
+        Args:
+            volume: A Volume object representing the UC object.
+        """
+        self._console_logger.info(f"Creating volume: {volume.identifier}")
+        self._spark.sql(f"USE CATALOG {volume.catalog};")
+        self._spark.sql(f"USE SCHEMA {volume.schema_name};")
+        for statement in volume.get_create_statement().split(";"):
+            if statement and statement != "\n":
+                self._spark.sql(statement)
+
+    def drop_volume(self, volume: Volume, if_exists: bool = True):
+        """Delete the volume.
+
+        Args:
+            volume: The volume to be deleted.
+            if_exists: If False, an error will be raised if the volume does not exist.
+        """
+        self._console_logger.info(f"Deleting volume: [' {volume.identifier}' ]")
+        self._spark.sql(f"DROP VOLUME {'IF EXISTS' if if_exists else ''} {volume.escaped_identifier};")
+        self._console_logger.info(f"Volume [' {volume.identifier}' ] has been deleted.")
+
+    def volume_exists(self, volume: Volume | None = None, volume_identifier: str | None = None) -> bool:
+        """Check if the volume exists.
+
+        Args:
+            volume: The volume to check.
+            volume_identifier: The identifier of the volume to check.
+
+        Raises:
+            ValueError: If both volume and volume_identifier are provided.
+
+        Returns:
+            True if the volume exists, False otherwise.
+        """
+        if volume and volume_identifier:
+            raise ValueError("Only one of volume or volume_identifier should be provided.")
+        if volume:
+            volume_identifier = volume.identifier
+
+        assert volume_identifier is not None
+
+        if volume_identifier.count(".") != 2:
+            raise ValueError("The identifier must be in the format 'catalog.schema.volume_name'.")
+        catalog, volume_schema, table_name = volume_identifier.split(".")
+        query_result = self._spark.sql(
+            f"""
+                SELECT 1 FROM {catalog}.information_schema.volumes
+                WHERE volume_name = '{table_name}'
+                AND volume_schema = '{volume_schema}'
+                LIMIT 1""",
+        )
+        result = query_result.count() > 0
+        return result is True

cloe_nessy/pipeline/actions/__init__.py

@@ -14,6 +14,7 @@ from .transform_distinct import TransformDistinctAction
 from .transform_filter import TransformFilterAction
 from .transform_generic_sql import TransformSqlAction
 from .transform_group_aggregate import TransformGroupAggregate
+from .transform_hash_columns import TransformHashColumnsAction
 from .transform_join import TransformJoinAction
 from .transform_json_normalize import TransformJsonNormalize
 from .transform_rename_columns import TransformRenameColumnsAction
@@ -21,6 +22,9 @@ from .transform_replace_values import TransformReplaceValuesAction
 from .transform_select_columns import TransformSelectColumnsAction
 from .transform_union import TransformUnionAction
 from .write_catalog_table import WriteCatalogTableAction
+from .write_delta_append import WriteDeltaAppendAction
+from .write_delta_merge import WriteDeltaMergeAction
+from .write_file import WriteFileAction
 
 # Get all subclasses of PipelineAction defined in this submodule
 pipeline_actions = {cls.name: cls for cls in PipelineAction.__subclasses__()}
@@ -35,7 +39,6 @@ __all__ = [
     "ReadExcelAction",
     "ReadFilesAction",
     "ReadMetadataYAMLAction",
-    "WriteCatalogTableAction",
     "PipelineActionType",
     "TransformFilterAction",
     "TransformUnionAction",
@@ -51,4 +54,9 @@ __all__ = [
     "TransformRenameColumnsAction",
     "TransformReplaceValuesAction",
     "TransformSelectColumnsAction",
+    "WriteCatalogTableAction",
+    "WriteDeltaAppendAction",
+    "WriteDeltaMergeAction",
+    "WriteFileAction",
+    "TransformHashColumnsAction",
 ]

cloe_nessy/pipeline/actions/read_catalog_table.py

@@ -61,8 +61,5 @@ class ReadCatalogTableAction(PipelineAction):
             raise ValueError("Table name must be specified or a valid Table object with identifier must be set.")
 
         table_reader = CatalogReader()
-        df = table_reader.read(
-            table_identifier=table_identifier,
-            **options,
-        )
+        df = table_reader.read(table_identifier=table_identifier, options=options)
         return context.from_existing(data=df)

cloe_nessy/pipeline/actions/transform_hash_columns.py

@@ -0,0 +1,209 @@
+from typing import Any
+
+from pydantic import BaseModel, Field, model_validator
+from pyspark.sql import functions as F
+
+from ..pipeline_action import PipelineAction
+from ..pipeline_context import PipelineContext
+
+SUPPORTED_ALGORITHMS = {"hash", "md5", "sha1", "sha2", "xxhash64", "crc32"}
+VALID_SHA2_BITS = {224, 256, 384, 512}
+
+
+class HashSettings(BaseModel):
+    """Represents the settings for hashing columns.
+
+    Attributes:
+        columns: List of column names to hash.
+        algorithm: Hashing algorithm to use. Must be one of
+            "hash", "md5", "sha1", "sha2", "xxhash64", or "crc32".
+        bits: Bit length for the 'sha2' algorithm. Optional.
+    """
+
+    columns: list[str]
+    algorithm: str = Field(..., description="Hashing algorithm to use")
+    bits: int | None = Field(default=None, description="Only required for sha2")
+
+    @model_validator(mode="before")
+    def validate_all(cls, values):
+        """Validates the input values for a hashing operation before model instantiation.
+
+        This method performs the following checks:
+
+        1. Ensures the specified hashing algorithm is supported.
+        2. Validates that at least one column is provided and that the columns parameter is a non-empty list.
+        3. Checks that hashing multiple columns is only supported for the 'hash' and 'xxhash64' algorithms.
+        4. For the 'sha2' algorithm, ensures that the 'bits' parameter is one of the valid options.
+        5. Ensures that the 'bits' parameter is not provided for algorithms other than 'sha2'.
+
+        Raises:
+            ValueError: If the algorithm is unsupported, no columns are provided, the columns parameter is invalid,
+                        or the 'bits' parameter is invalid for the specified algorithm.
+            NotImplementedError: If multiple columns are provided and the algorithm does not support hashing multiple columns.
+
+        Args:
+            cls: The class being validated.
+            values: A dictionary of input values containing 'algorithm', 'columns', and 'bits'.
+
+        Returns:
+            The validated input values.
+        """
+        algorithm = values.get("algorithm")
+        columns = values.get("columns")
+        bits = values.get("bits")
+
+        if algorithm not in SUPPORTED_ALGORITHMS:
+            raise ValueError(
+                f"Unsupported hashing algorithm '{algorithm}'. Supported algorithms are: {SUPPORTED_ALGORITHMS}."
+            )
+
+        if not columns or not isinstance(columns, list) or len(columns) == 0:
+            raise ValueError("At least one column must be provided.")
+
+        if len(columns) > 1 and algorithm not in {"hash", "xxhash64"}:
+            raise NotImplementedError(
+                f"Hashing multiple columns is only supported for 'hash' and 'xxhash64'. Algorithm '{algorithm}' does not support this."
+            )
+
+        if algorithm == "sha2":
+            if bits not in VALID_SHA2_BITS:
+                raise ValueError(f"'bits' must be one of {VALID_SHA2_BITS} when using 'sha2'.")
+        elif bits is not None:
+            raise ValueError("'bits' is only allowed when algorithm is 'sha2'.")
+
+        return values
+
+
+class HashConfig(BaseModel):
+    """A configuration model for defining hash settings for specific columns.
+
+    Attributes:
+        hash_config: A dictionary where the keys are column names
+            (as strings) and the values are `HashSettings` objects that define
+            the hash settings for each column.
+
+    Methods:
+        validate_config: Validates the hash configuration to ensure it contains
+            at least one entry and that all column names are valid strings. Raises a
+            `ValueError` if the configuration is invalid.
+    """
+
+    hash_config: dict[str, HashSettings]
+
+    @model_validator(mode="before")
+    def validate_config(cls, values):
+        """Validates the hash configuration provided in the model.
+
+        This method is executed in "before" mode to ensure that the `hash_config`
+        field in the input values meets the required criteria:
+
+        - It must be a dictionary.
+        - It must contain at least one entry.
+        - Each key in the dictionary must be a non-empty string.
+
+        Raises:
+            ValueError: If `hash_config` is missing, not a dictionary, empty, or
+                        contains invalid column names.
+
+        Args:
+            cls: The class to which this validator is applied.
+            values: The input values to validate.
+
+        Returns:
+            The validated input values.
+        """
+        config = values.get("hash_config")
+        if not config or not isinstance(config, dict) or len(config) == 0:
+            raise ValueError("Hash configuration must contain at least one entry.")
+        for new_col in config:
+            if not new_col or not isinstance(new_col, str):
+                raise ValueError(f"Invalid column name '{new_col}' in hash configuration.")
+        return values
+
+
+class TransformHashColumnsAction(PipelineAction):
+    """Hashes specified columns in a DataFrame using a chosen algorithm.
+
+    Given the following `hash_config`:
+
+    Example:
+        ```yaml
+        Hash Columns:
+            action: TRANSFORM_HASH_COLUMNS
+            options:
+                hash_config:
+                - hashed_column1:
+                    columns: ["column1", "column2"]
+                    algorithm: "sha2"
+                    bits: 224
+                - hashed_column2:
+                    columns: ["column1"]
+                    algorithm: "crc32"
+        ```
+
+    Given a DataFrame `df` with the following structure:
+
+    | column1 | column2 | column3 |
+    |---------|---------|---------|
+    |   foo   |   bar   |   baz   |
+
+    After running the action, the resulting DataFrame will look like:
+
+    | column1 | column2 | column3 |                 hashed_column1                            | hashed_column2 |
+    |---------|---------|---------|-----------------------------------------------------------|----------------|
+    |   foo   |   bar   |   baz   |  17725b837e9c896e7123b142eb980131dcc0baa6160db45d4adfdb21 |  1670361220    |
+
+
+    !!! note "Hash values might vary"
+        The actual hash values will depend on the hashing algorithm used and the input data.
+    """
+
+    name: str = "TRANSFORM_HASH_COLUMNS"
+
+    def run(
+        self,
+        context: PipelineContext,
+        *,
+        hash_config: HashConfig | None = None,
+        **_: Any,
+    ) -> PipelineContext:
+        """Hashes the specified columns in the DataFrame.
+
+        Args:
+            context: Context in which this Action is executed.
+            hash_config: Dictionary that contains the configuration for executing the hashing.
+
+        Returns:
+            Updated PipelineContext with hashed columns.
+
+        Raises:
+            ValueError: If columns are missing, data is None, or algorithm/bits are invalid.
+            ValueError: If the hash configuration is invalid.
+        """
+        if context.data is None:
+            raise ValueError("Context data is required for hashing.")
+
+        if not hash_config:
+            raise ValueError("Hash configuration is required.")
+
+        df = context.data
+
+        hash_functions = {
+            "hash": lambda cols: F.hash(*[F.col(c) for c in cols]).cast("string"),
+            "xxhash64": lambda cols: F.xxhash64(F.concat_ws("||", *[F.col(c) for c in cols])).cast("string"),
+            "md5": lambda cols: F.md5(F.concat_ws("||", *[F.col(c) for c in cols])).cast("string"),
+            "sha1": lambda cols: F.sha1(F.concat_ws("||", *[F.col(c) for c in cols])).cast("string"),
+            "sha2": lambda cols, bits: F.sha2(F.concat_ws("||", *[F.col(c) for c in cols]), bits).cast("string"),
+            "crc32": lambda cols: F.crc32(F.concat_ws("||", *[F.col(c) for c in cols])).cast("string"),
+        }
+        default_sha2_bits = 256
+
+        config_obj = HashConfig.model_validate({"hash_config": hash_config})
+        for new_col, config in config_obj.hash_config.items():
+            hash_func = hash_functions[config.algorithm]
+            if config.algorithm == "sha2":
+                df = df.withColumn(new_col, hash_func(config.columns, config.bits or default_sha2_bits))  # type: ignore
+            else:
+                df = df.withColumn(new_col, hash_func(config.columns))  # type: ignore
+
+        return context.from_existing(data=df)