cloe-nessy 0.2.9__py3-none-any.whl

cloe_nessy/integration/reader/file_reader.py

@@ -0,0 +1,96 @@
+from typing import Any
+
+import pyspark.sql.functions as F
+from pyspark.sql import DataFrame
+
+from ...file_utilities import get_file_paths
+from .reader import BaseReader
+
+
+class FileReader(BaseReader):
+    """Utility class for reading a file into a DataFrame.
+
+    This class reads data from files and loads it into a Spark DataFrame.
+    """
+
+    def __init__(self):
+        """Initializes the FileReader object."""
+        super().__init__()
+
+    def read(
+        self,
+        location: str,
+        spark_format: str | None = None,
+        extension: str | None = None,
+        schema: str | None = None,
+        search_subdirs: bool = True,
+        options: dict | None = None,
+        add_metadata_column: bool = False,
+        **kwargs: Any,
+    ) -> DataFrame:
+        """Reads files from a specified location and returns a DataFrame.
+
+        Arguments:
+            location: Location of files to read.
+            spark_format: Format of files to read. If not provided, it will be inferred from the extension.
+            extension: File extension (csv, json, parquet, txt). Used if spark_format is not provided.
+            schema: Schema of the file. If None, schema will be inferred.
+            search_subdirs: Whether to include files in subdirectories.
+            options: Spark DataFrame reader options.
+            add_metadata_column: Whether to include __metadata column in the DataFrame.
+            kwargs: This method does not accept any additional keyword arguments.
+        """
+        if options is None:
+            options = {}
+
+        if not spark_format and not extension:
+            raise ValueError("Either spark_format or extension must be provided.")
+        self._console_logger.debug(f"Reading files from [ '{location}' ] ...")
+        extension_to_datatype_dict = {"csv": "csv", "json": "json", "parquet": "parquet", "txt": "text", "xml": "xml"}
+
+        if extension and not spark_format:
+            if extension not in extension_to_datatype_dict:
+                raise ValueError(f"Unsupported file extension: {extension}")
+            spark_format = extension_to_datatype_dict[extension]
+        self._console_logger.debug(f"Reading files with format: {spark_format}")
+        if extension:
+            file_paths = get_file_paths(location, extension, search_subdirs)
+        else:
+            file_paths = [location]
+        self._console_logger.debug(f"Found {len(file_paths)} files to read")
+        self._console_logger.debug(f"File paths: {file_paths}")
+        assert spark_format is not None
+
+        reader = self._spark.read.format(spark_format)
+        if schema:
+            reader.schema(schema)
+        else:
+            options["inferSchema"] = True
+
+        self._console_logger.debug(f"Setting options: {options}")
+        reader.options(**options)
+
+        try:
+            self._console_logger.debug("Loading files into DataFrame")
+            df = reader.load(file_paths)
+            self._console_logger.debug("Successfully loaded files into DataFrame")
+            if add_metadata_column:
+                df = self._add_metadata_column(df)
+        except Exception as e:
+            self._console_logger.error(f"Failed to read files from [ '{location}' ]: {e}")
+            raise
+        else:
+            self._console_logger.info(f"Successfully read files from [ '{location}' ]")
+            return df
+
+    def _add_metadata_column(self, df: DataFrame) -> DataFrame:
+        """Add all metadata columns to the DataFrame."""
+        # Extract metadata fields into separate columns
+        metadata_columns = df.select("_metadata.*").columns
+
+        entries = [(F.lit(field), F.col(f"_metadata.{field}")) for field in metadata_columns]
+        flat_list = [item for tup in entries for item in tup]
+
+        df = df.withColumn("__metadata", F.create_map(flat_list))
+
+        return df

cloe_nessy/integration/reader/reader.py

@@ -0,0 +1,34 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from pyspark.sql import DataFrame, SparkSession
+
+from ...logging.logger_mixin import LoggerMixin
+from ...session import SessionManager
+
+
+class BaseReader(ABC, LoggerMixin):
+    """Abstract base class for reading data into a Spark DataFrame.
+
+    This class provides a common interface for different types of data readers.
+
+    Attributes:
+        _spark: The Spark session used for creating DataFrames.
+    """
+
+    def __init__(self) -> None:
+        self._spark: SparkSession = SessionManager.get_spark_session()
+        self._console_logger = self.get_console_logger()
+
+    @abstractmethod
+    def read(self, *args: Any, **kwargs: Any) -> DataFrame:
+        """Abstract method to return a batch data frame.
+
+        Args:
+            *args: Arbitrary non-keyword arguments for reading data.
+            **kwargs: Arbitrary keyword arguments for reading data.
+
+        Returns:
+            DataFrame: The Spark DataFrame containing the read data.
+        """
+        pass

cloe_nessy/integration/writer/__init__.py

@@ -0,0 +1,3 @@
+from .catalog_writer import CatalogWriter
+
+__all__ = ["CatalogWriter"]

cloe_nessy/integration/writer/catalog_writer.py

@@ -0,0 +1,48 @@
+from pyspark.sql import DataFrame
+
+
+class CatalogWriter:
+    """A writer for Catalog tables."""
+
+    @staticmethod
+    def write_table(
+        df: DataFrame | None,
+        table_identifier: str | None,
+        partition_by: str | list[str] | None = None,
+        options: dict[str, str] | None = None,
+        mode: str = "append",
+    ) -> None:
+        """Write a table to the unity catalog.
+
+        Args:
+            df: The DataFrame to write.
+            table_identifier: The table identifier in the unity catalog in the
+                              format 'catalog.schema.table'.
+            mode: The write mode. One of append, overwrite, error, errorifexists, ignore.
+            partition_by: Names of the partitioning columns.
+            options: All other string options.
+
+        Notes:
+            append: Append contents of this DataFrame to existing data.
+            overwrite: Overwrite existing data.
+            error or errorifexists: Throw an exception if data already exists.
+            ignore: Silently ignore this operation if data already exists.
+
+        Raises:
+            ValueError: If the mode is not one of append, overwrite, error, errorifexists, ignore.
+            ValueError: If the table_identifier is not a string or not in the format 'catalog.schema.table'.
+            ValueError: If the DataFrame is None.
+        """
+        if mode not in ("append", "overwrite", "error", "errorifexists", "ignore"):
+            raise ValueError("mode must be one of append, overwrite, error, errorifexists, ignore")
+        if not table_identifier:
+            raise ValueError("table_identifier is required")
+        elif not isinstance(table_identifier, str):
+            raise ValueError("table_identifier must be a string")
+        elif len(table_identifier.split(".")) != 3:
+            raise ValueError("table_identifier must be in the format 'catalog.schema.table'")
+        if not df:
+            raise ValueError("df is required, but was None.")
+        if options is None:
+            options = {}
+        df.write.saveAsTable(table_identifier, mode=mode, partitionBy=partition_by, **options)

cloe_nessy/logging/__init__.py

@@ -0,0 +1,3 @@
+from .logger_mixin import LoggerMixin
+
+__all__ = ["LoggerMixin"]

cloe_nessy/logging/logger_mixin.py

@@ -0,0 +1,162 @@
+import logging
+import logging.handlers
+from typing import cast
+
+from cloe_logging import LoggerFactory
+
+from ..settings import LoggingSettings, NessySettings
+
+factory = LoggerFactory()
+
+DEFAULT_COLUMN_SPLIT_CHAR = "|"
+DEFAULT_KEY_VALUE_SPLIT_CHAR = ":"
+
+
+class LoggerMixin:
+    """LoggingMixin class to add logging functionality to classes."""
+
+    def get_console_logger(
+        self,
+        level: int | None = None,
+        log_format: str | None = None,
+    ) -> logging.Logger:
+        """Adds a console logger to the class.
+
+        Args:
+            level: The logging level for the console logger.
+            log_format: The format for the console logger.
+
+        Returns:
+            The logger with the console handler.
+        """
+        logging_settings: LoggingSettings = NessySettings().logging
+        logger = LoggerFactory.get_logger(
+            handler_types=["console"],
+            logger_name=f"Console:{self.__class__.__name__}",
+            logging_level=level if level is not None else logging_settings.log_level_console,
+            log_format=log_format if log_format is not None else logging_settings.log_format_console,
+        )
+        return cast(logging.Logger, logger)
+
+    def get_tabular_logger(
+        self,
+        logger_name: str | None = None,
+        handlers: list[str] | None = None,
+        level: int | None = None,
+        add_log_analytics_logger: bool | None = None,
+        add_unity_catalog_logger: bool | None = None,
+        # LAW
+        log_type: str | None = None,
+        workspace_id: str | None = None,
+        shared_key: str | None = None,
+        # UC
+        uc_workspace_url: str | None = None,
+        uc_warehouse_id: str | None = None,
+        uc_catalog_name: str | None = None,
+        uc_schema_name: str | None = None,
+        uc_table_name: str | None = None,
+        uc_table_columns: dict[str, str] | None = None,
+        column_split_char: str = DEFAULT_COLUMN_SPLIT_CHAR,
+        key_value_split_char: str = DEFAULT_KEY_VALUE_SPLIT_CHAR,
+    ) -> logging.Logger:
+        """Adds a tabular logger to the class.
+
+        Args:
+            logger_name: The name of the logger.
+            handlers: The list of handlers to add.
+            level: The logging level for the tabular logger. If not provided, the value from the settings will be used.
+            add_log_analytics_logger: Whether to add a LogAnalyticsHandler to the logger. If not provided, the value from the settings will be used.
+            add_unity_catalog_logger: Whether to add a UnityCatalogHandler to the logger. If not provided, the value from the settings will be used.
+            log_type: The log type for the Log Analytics workspace.
+            workspace_id: The workspace id for the Log Analytics workspace. If not provided, the value from the settings will be used.
+            shared_key: The shared key for the Log Analytics workspace.
+            uc_workspace_url: The workspace url for the Unity Catalog. If not provided, the value from the settings will be used.
+            uc_warehouse_id: The warehouse id for the Unity Catalog. If not provided, the value from the settings will be used.
+            uc_catalog_name: The catalog name for the Unity Catalog. If not provided, the value from the settings will be used.
+            uc_schema_name: The schema name for the Unity Catalog. If not provided, the value from the settings will be used.
+            uc_table_name: The table name for the Unity Catalog.
+            uc_table_columns: The columns for the Unity Catalog Table.
+            column_split_char: The column split character for the Log Analytics workspace and Unity Catalog. Defaults to "|".
+            key_value_split_char: The key value split character for the Log Analytics workspace and Unity Catalog. Defaults to ":".
+
+        Returns:
+            The logger with the added tabular handlers.
+        """
+        if handlers is None:
+            handlers = []
+        logging_settings = NessySettings().logging
+
+        if self.should_add_log_analytics_handler(logging_settings, add_log_analytics_logger):
+            handlers.append("log_analytics")
+
+        if self.should_add_unity_catalog_handler(logging_settings, add_unity_catalog_logger):
+            handlers.append("unity_catalog")
+
+        logger = LoggerFactory.get_logger(
+            handler_types=handlers,
+            logger_name=logger_name or f"Tabular:{self.__class__.__name__}",
+            level=level,
+            column_split_char=column_split_char,
+            key_value_split_char=key_value_split_char,
+            # UC Settings
+            uc_table_name=uc_table_name,
+            uc_catalog_name=uc_catalog_name or logging_settings.uc_catalog_name,
+            uc_schema_name=uc_schema_name or logging_settings.uc_schema_name,
+            uc_table_columns=uc_table_columns,
+            workspace_url=uc_workspace_url or logging_settings.uc_workspace_url,
+            warehouse_id=uc_warehouse_id or logging_settings.uc_warehouse_id,
+            # LAW Settings
+            workspace_id=workspace_id or logging_settings.log_analytics_workspace_id,
+            shared_key=shared_key or logging_settings.log_analytics_shared_key,
+            log_type=log_type,
+            test_connectivity=False,
+        )
+        return cast(logging.Logger, logger)
+
+    @staticmethod
+    def should_add_log_analytics_handler(
+        logging_settings: LoggingSettings,
+        add_log_analytics_logger: bool | None,
+        **kwargs,  # noqa: ARG004
+    ) -> bool:
+        """Determines if a LogAnalyticsHandler should be added to the logger.
+
+        The Logger will be added if the `target_log_analytics` setting is set to True or if the `add_log_analytics_logger`
+        argument is set to True.
+
+        Setting `target_log_analytics` to False will prevent the handler from being added.
+
+        Args:
+            logging_settings: The logging settings to use for the logger.
+            add_log_analytics_logger: Whether to add a LogAnalyticsHandler to the logger.
+            **kwargs: Additional keyword arguments. Not used.
+
+        Returns:
+            bool: True if the LogAnalyticsHandler should be added, False otherwise.
+        """
+        disable_overwrite = logging_settings.target_log_analytics is False
+        enable_logger = logging_settings.target_log_analytics or add_log_analytics_logger
+        return cast(bool, enable_logger and not disable_overwrite)
+
+    @staticmethod
+    def should_add_unity_catalog_handler(
+        logging_settings: LoggingSettings,
+        add_unity_catalog_logger: bool | None,
+    ) -> bool:
+        """Determines if a UnityCatalogHandler should be added to the logger.
+
+        The Logger will be added if the `target_unity_catalog_table` setting is set to True or if the `add_unity_catalog_logger`
+        argument is set to True.
+
+        Setting `target_unity_catalog_table` to False will prevent the handler from being added.
+
+        Args:
+            logging_settings: The logging settings to use for the logger.
+            add_unity_catalog_logger: Whether to add a UnityCatalogHandler to the logger.
+
+        Returns:
+            bool: True if the UnityCatalogHandler should be added, False otherwise.
+        """
+        disable_overwrite = logging_settings.target_unity_catalog_table is False
+        enable_logger = logging_settings.target_unity_catalog_table or add_unity_catalog_logger
+        return cast(bool, enable_logger and not disable_overwrite)

cloe_nessy/models/__init__.py

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

cloe_nessy/models/column.py

@@ -0,0 +1,65 @@
+import re
+from typing import Any
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+COLUMN_DATA_TYPE_LIST = {
+    "string",
+    "integer",
+    "int",
+    "smallint",
+    "float",
+    "boolean",
+    "bool",
+    "bigint",
+    "long",
+    "double",
+    "date",
+    "timestamp",
+    "array",
+    "map",
+    "variant",
+    "struct",
+}
+
+
+class Column(BaseModel):
+    """Represents a Column of a Table."""
+
+    name: str
+    data_type: str
+    nullable: bool
+    default_value: Any = None
+    generated: str | None = None
+    properties: dict[str, Any] = Field(default_factory=dict)
+    comment: str | None = None
+
+    @field_validator("data_type", mode="before")
+    def data_type_transform(cls, raw: str) -> str:
+        """Map potential aliases to the correct SQL data type.
+
+        Args:
+            raw: The value for the data type.
+        """
+        val = raw.lower()
+        base_data_types = re.findall(r"\b[a-z]+\b", val)
+        forbidden_characters = re.findall(r"[^a-z\<\>)]+", val)
+
+        if forbidden_characters:
+            raise ValueError(f"Forbidden characters in data type definition [ '{val}' ]: [' {forbidden_characters} ']")
+        for base_data_type in base_data_types:
+            if base_data_type not in COLUMN_DATA_TYPE_LIST:
+                raise ValueError(f"Unknown data type used in data type definition [ '{val}' ]")
+        return val
+
+    @model_validator(mode="before")
+    def _validate_generated_and_default_value(cls, v: Any) -> Any:
+        """Check if a column has a default value and is generated.
+
+        That doesn't make sense, so an error should be raised.
+        """
+        if v.get("default_value") and v.get("generated"):
+            raise ValueError("A column can't have a default value and be generated.")
+        if (v.get("default_value") or v.get("generated")) and v.get("nullable") is True:
+            raise ValueError("A column can't have a default value or be generated and be nullable.")
+        return v

cloe_nessy/models/constraint.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+
+class Constraint(BaseModel):
+    """Represents a Constraint on a Table."""
+
+    name: str
+    expression: str
+    description: str | None = None

cloe_nessy/models/foreign_key.py

@@ -0,0 +1,34 @@
+import os
+
+from pydantic import BaseModel, field_validator
+
+
+def _process_column_input(v):
+    if isinstance(v, str):
+        v = [col.strip() for col in v.split(",")]
+    return v
+
+
+class ForeignKey(BaseModel):
+    """Represents a ForeignKey."""
+
+    foreign_key_columns: list[str]
+    parent_table: str
+    parent_columns: list[str]
+    foreign_key_option: list[str] | None = None
+
+    @field_validator("foreign_key_columns", mode="before")
+    def _validate_foreign_key_columns(cls, v):
+        return _process_column_input(v)
+
+    @field_validator("parent_columns", mode="before")
+    def _validate_parent_columns(cls, v):
+        return _process_column_input(v)
+
+    @field_validator("parent_table", mode="before")
+    def _validate_identifier(cls, v):
+        if len(v.split(".")) != 3:
+            raise ValueError("The 'parent_table' must be in the format 'catalog.schema.table'")
+        if "<env>" in v:
+            v = v.replace("<env>", os.environ["PROJECT_ENVIRONMENT"])
+        return v

cloe_nessy/models/mixins/read_instance_mixin.py

@@ -0,0 +1,124 @@
+import os
+import pathlib
+import re
+from typing import Any, Self
+
+import yaml
+import yaml.parser
+import yaml.scanner
+from pydantic import BaseModel, ValidationError
+
+from ...session import SessionManager
+from ..types import ValidationErrorType
+
+
+class ReadInstancesMixin(BaseModel):
+    """This class defines the methods to read, validate and parse metadata definitions."""
+
+    @classmethod
+    def metadata_to_instance(cls, data: dict) -> tuple[Self | None, list[ValidationError]]:
+        """Parses a Dictionary to an instance.
+
+        Returns:
+            An instance and potentially a list of errors.
+        """
+        errors = []
+        try:
+            instance = cls(**data)
+        except ValidationError as e:
+            instance = None
+            errors.append(e)
+        return instance, errors
+
+    @classmethod
+    def read_instance_from_file(
+        cls,
+        instance_path: pathlib.Path,
+        **_: Any,  # allow subclasses to pass additional arguments
+    ) -> tuple[Self | None, list[ValidationErrorType]]:
+        """Read and instantiate a single YAML file for the given path.
+
+        Arguments:
+            instance_path: The path to the file to instantiate.
+
+        Return:
+            Returns a tuple of the instantiated model and errors.
+        """
+        errors: list[ValidationErrorType] = []
+        try:
+            with instance_path.open("r") as file:
+                raw_string = file.read()
+                yaml_str = cls._replace_variables(raw_string)
+                data = yaml.safe_load(yaml_str)
+                instance, sub_errors = cls.metadata_to_instance(data)
+                errors += sub_errors
+        except (ValidationError, yaml.parser.ParserError, yaml.scanner.ScannerError) as e:
+            instance = None
+            errors.append(e)
+        return instance, errors
+
+    @classmethod
+    def read_instances_from_directory(
+        cls,
+        instance_path: pathlib.Path,
+        fail_on_missing_subfolder: bool = True,
+        **_: Any,  # allow subclasses to pass additional arguments
+    ) -> tuple[list[Self], list[ValidationErrorType]]:
+        """Read and instantiate all *.yaml files for the given path.
+
+        Arguments:
+            instance_path: Path to the directory containing the instance definitions as YAML files.
+            fail_on_missing_subfolder: If False return a tuple with 2 empty
+                    lists. Otherwise raise a FileNotFoundError.
+
+        Return:
+            Returns a tuple of the instantiated models and errors.
+        """
+        instances: list[Self] = []
+        errors: list[ValidationErrorType] = []
+
+        if not instance_path.exists() or not instance_path.is_dir():
+            if fail_on_missing_subfolder:
+                raise FileNotFoundError(f"Directory not found: {instance_path}")
+            else:
+                return instances, errors
+
+        for instance_file in instance_path.iterdir():
+            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)
+                instances += [] if instance is None else [instance]
+            errors += sub_errors
+
+        return instances, errors
+
+    @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.
+
+        Args:
+            yaml_str (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

cloe_nessy/models/mixins/template_loader_mixin.py

@@ -0,0 +1,18 @@
+from pathlib import Path
+
+from jinja2 import FileSystemLoader
+from jinja2 import Template as JinjaTemplate
+from jinja2.environment import Environment
+
+
+class TemplateLoaderMixin:
+    """A Mixin to load Jinja Templates."""
+
+    @staticmethod
+    def get_template(template_path: Path, template_name: str) -> JinjaTemplate:
+        """Load the specified template."""
+        loader: FileSystemLoader = FileSystemLoader(template_path)
+
+        env = Environment(loader=loader, keep_trailing_newline=True)
+
+        return env.get_template(template_name)