cloe-nessy 0.2.9__py3-none-any.whl

cloe_nessy/file_utilities/strategies/base_strategy.py

@@ -0,0 +1,59 @@
+from abc import ABC, abstractmethod
+
+
+class FileRetrievalStrategy(ABC):
+    """Abstract base class for file retrieval strategies.
+
+    This class defines the interface for strategies that retrieve file paths
+    based on certain criteria. Concrete implementations of this class should
+    provide the logic for retrieving file paths.
+    """
+
+    @staticmethod
+    @abstractmethod
+    def get_file_paths(location: str, extension: str | None = None, search_subdirs: bool = True) -> list[str]:
+        """Retrieves a list of file paths based on the specified criteria.
+
+        Args:
+            location: The location to search for files.
+            extension: The file extension to filter by. If None, no extension filtering is applied.
+                If an empty string, it matches files with no extension.
+            search_subdirs: Whether to search in subdirectories.
+
+        Returns:
+            list[str]: A list of file paths that match the specified criteria.
+        """
+        pass
+
+    @staticmethod
+    def _matches_extension(file_name: str, extension: str | None) -> bool:
+        """Determines if a file name ends with the specified extension.
+
+        This method checks whether the provided file name matches the given file extension. The comparison is case-insensitive.
+
+        If the `extension` is an empty string, it checks if the file name either does not contain a dot or ends with a dot,
+        which indicates a file with no extension. If the `extension` is `None`, it matches any file name regardless of extension.
+
+        If the `extension` contains a dot (e.g., ".txt"), it is compared directly against the end of the file name. Otherwise,
+        a dot is prefixed to the `extension` to create the expected file extension format (e.g., "txt" becomes ".txt").
+
+        Args:
+            file_name: The name of the file to check. This is converted to lowercase for case-insensitive comparison.
+            extension: The extension to match against. Can be a string with or without a leading dot.
+
+        Returns:
+            bool: True if the file name ends with the specified extension, False otherwise.
+        """
+        file_name_lower = file_name.lower()
+        matches = False
+
+        if extension == "":
+            matches = "." not in file_name_lower or file_name.endswith(".")
+        elif extension is None:
+            matches = True
+        elif "." in extension:
+            matches = file_name_lower.endswith(extension.lower())
+        else:
+            matches = file_name_lower.endswith(f".{extension.lower()}")
+
+        return matches

cloe_nessy/file_utilities/strategies/local_strategy.py

@@ -0,0 +1,51 @@
+import os
+
+from ..exceptions import FileUtilitiesError
+from .base_strategy import FileRetrievalStrategy
+
+
+class LocalDirectoryStrategy(FileRetrievalStrategy):
+    """Strategy for retrieving files from a local directory.
+
+    This strategy implements the file retrieval logic for local directories, including
+    optional recursive search through subdirectories and filtering by file extension.
+    """
+
+    @staticmethod
+    def get_file_paths(location: str, extension: str | None = None, search_subdirs: bool = True) -> list[str]:
+        """Recursively retrieves all files with a specified extension from a given directory and its subdirectories.
+
+        Args:
+            location: Top-level directory to read from, e.g., '/Volumes/my_volume/landing/example_landing/'.
+            extension: File extension, e.g., 'csv', 'json'. Input an empty string to get files without any
+                                    extension, input None to get all files.
+            search_subdirs: If True, function will also search within all subdirectories.
+
+        Returns:
+            List: List of files in the directory and its subdirectories with the given extension.
+
+        Raises:
+            ValueError: If the location is not provided.
+            FileUtilitiesError: For any other unexpected errors.
+        """
+        if not location:
+            raise ValueError("location is required")
+
+        if not os.path.isdir(location):
+            raise FileUtilitiesError(f"The provided path '{location}' is not a valid directory.")
+
+        file_list = []
+
+        try:
+            for root, _, files in os.walk(location):
+                if not search_subdirs and root != location:
+                    continue
+
+                for file_name in files:
+                    if FileRetrievalStrategy._matches_extension(file_name, extension):
+                        file_list.append(os.path.join(root, file_name))
+
+        except Exception as err:
+            raise FileUtilitiesError(f"An error occurred while retrieving file paths: {err}") from err
+
+        return file_list

cloe_nessy/file_utilities/strategies/onelake_strategy.py

@@ -0,0 +1,31 @@
+from .base_strategy import FileRetrievalStrategy
+from .local_strategy import LocalDirectoryStrategy
+
+
+class OneLakeStrategy(FileRetrievalStrategy):
+    """Strategy for retrieving files from the OneLake."""
+
+    @staticmethod
+    def get_file_paths(location: str, extension: str | None = None, search_subdirs: bool = True) -> list:
+        """Recursively retrieves all files with a specified extension from a given directory and its subdirectories.
+
+        Args:
+            location: Top-level directory to read from, e.g., '/Volumes/my_volume/landing/example_landing/'.
+            extension: File extension, e.g., 'csv', 'json'. Input an empty string to get files without any
+                                    extension, input None to get all files.
+            search_subdirs: If True, function will also search within all subdirectories.
+
+        Returns:
+            List: List of files in the directory and its subdirectories with the given extension.
+
+        Raises:
+            ValueError: If the location is not provided.
+            Exception: For any other unexpected errors.
+        """
+        if not location:
+            raise ValueError("location is required")
+
+        file_paths = LocalDirectoryStrategy.get_file_paths(location, extension, search_subdirs)
+
+        shortened_file_paths = [p.replace("/lakehouse/default/", "") for p in file_paths]
+        return shortened_file_paths

cloe_nessy/file_utilities/strategies/utils_strategy.py

@@ -0,0 +1,72 @@
+from ...session import SessionManager
+from ..exceptions import FileUtilitiesError
+from .base_strategy import FileRetrievalStrategy
+
+
+class UtilsStrategy(FileRetrievalStrategy):
+    """Strategy for retrieving files using DButils (in Databricks) and mssparkutils (in Fabric).
+
+    This strategy implements the file retrieval logic using utils, including
+    recursive search through directories and filtering by file extension.
+    """
+
+    @staticmethod
+    def get_file_paths(location: str, extension: str | None = None, search_subdirs: bool = True) -> list:
+        """Recursively retrieves all files with a specified extension from a given directory and its subdirectories.
+
+        Args:
+            location: Top-level directory to read from, e.g., '/Volumes/my_volume/landing/example_landing/'.
+            extension: File extension, e.g., 'csv', 'json'. Input an empty string to get files without any
+                                    extension, input None to get all files.
+            search_subdirs: If True, function will also search within all subdirectories.
+
+        Returns:
+            List: List of files in the directory and its subdirectories with the given extension.
+
+        Raises:
+            ValueError: If the location is not provided.
+            Exception: For any other unexpected errors.
+        """
+        if not location:
+            raise ValueError("location is required")
+
+        utils = SessionManager.get_utils()
+
+        def _inner_loop(directory: str) -> list:
+            """Inner loop that recursively traverses directories to find all files with a given extension.
+
+            Args:
+                directory: The directory to start searching in.
+
+            Returns:
+                List: List of all files in the directory and its subdirectories with the given extension.
+            """
+            try:
+                dirs = utils.fs.ls(directory)
+            except Exception as err:
+                raise FileUtilitiesError(
+                    f"An error occurred while listing files in directory '{directory}': {err}"
+                ) from err
+
+            file_list = [file for file in dirs if FileRetrievalStrategy._matches_extension(file.name, extension)]
+
+            if search_subdirs:
+                for p in dirs:
+                    if p.isDir() and p.path != directory:
+                        try:
+                            sub_dir_files = _inner_loop(p.path)
+                            file_list.extend(sub_dir_files)
+                        except Exception as err:
+                            raise FileUtilitiesError(
+                                f"An error occurred while processing subdirectory '{p.path}': {err}"
+                            ) from err
+
+            return file_list
+
+        try:
+            file_list = _inner_loop(location)
+        except Exception as err:
+            raise FileUtilitiesError(f"An error occurred while retrieving file paths: {err}") from err
+
+        file_list = [p.path for p in file_list if not p.isDir()]
+        return file_list

cloe_nessy/integration/reader/__init__.py

@@ -0,0 +1,6 @@
+from .api_reader import APIReader
+from .catalog_reader import CatalogReader
+from .excel_reader import ExcelDataFrameReader
+from .file_reader import FileReader
+
+__all__ = ["APIReader", "CatalogReader", "FileReader", "ExcelDataFrameReader"]

cloe_nessy/integration/reader/api_reader.py

@@ -0,0 +1,141 @@
+import json
+from typing import Any
+
+import pyspark.sql.functions as F
+from pyspark.sql import DataFrame
+from requests.auth import AuthBase
+
+from cloe_nessy.clients.api_client.api_response import APIResponse
+
+from ...clients.api_client import APIClient
+from ...clients.api_client.exceptions import (
+    APIClientConnectionError,
+    APIClientError,
+    APIClientHTTPError,
+    APIClientTimeoutError,
+)
+from .reader import BaseReader
+
+
+class APIReader(BaseReader):
+    """Utility class for reading an API into a DataFrame.
+
+    This class uses an APIClient to fetch data from an API and load it into a Spark DataFrame.
+
+    Attributes:
+        api_client: The client for making API requests.
+    """
+
+    def __init__(self, base_url: str, auth: AuthBase | None, default_headers: dict[str, str] | None = None):
+        """Initializes the APIReader object.
+
+        Args:
+            base_url : The base URL for the API.
+            auth: The authentication method for the API.
+            default_headers: Default headers to include in requests.
+        """
+        super().__init__()
+        self.api_client = APIClient(base_url, auth, default_headers)
+
+    def read(
+        self,
+        endpoint: str = "",
+        method: str = "GET",
+        key: str | None = None,
+        timeout: int = 30,
+        params: dict[str, str] | None = None,
+        headers: dict[str, str] | None = None,
+        data: dict[str, str] | None = None,
+        json_body: dict[str, str] | None = None,
+        max_retries: int = 0,
+        options: dict[str, str] | None = None,
+        add_metadata_column: bool = False,
+        **kwargs: Any,
+    ) -> DataFrame:
+        """Reads data from an API endpoint and returns it as a DataFrame.
+
+        Args:
+            endpoint: The endpoint to send the request to.
+            method: The HTTP method to use for the request.
+            key: The key to extract from the JSON response.
+            timeout: The timeout for the request in seconds.
+            params: The query parameters for the request.
+            headers: The headers to include in the request.
+            data: The form data to include in the request.
+            json_body: The JSON data to include in the request.
+            max_retries: The maximum number of retries for the request.
+            options: Additional options for the createDataFrame function.
+            add_metadata_column: If set, adds a __metadata column containing metadata about the API response.
+            kwargs: This method does not accept any additional keyword arguments.
+
+        Returns:
+            DataFrame: The Spark DataFrame containing the read data in the json_object column.
+
+        Raises:
+            RuntimeError: If there is an error with the API request or reading the data.
+        """
+        options = options or {}
+        try:
+            response = self.api_client.request(
+                method=method,
+                endpoint=endpoint,
+                timeout=timeout,
+                params=params,
+                headers=headers,
+                data=data,
+                json=json_body,
+                max_retries=max_retries,
+            )
+            data_list = response.to_dict(key)
+            json_string = json.dumps(data_list)
+            df: DataFrame = self._spark.createDataFrame(data={json_string}, schema=["json_string"], **options)  # type: ignore
+            row = df.select("json_string").head()
+            if row is not None:
+                schema = F.schema_of_json(row[0])
+            else:
+                raise RuntimeError("It was not possible to infer the schema of the JSON data.")
+            df_result = df.withColumn("json_object", F.from_json("json_string", schema)).select("json_object")
+            if add_metadata_column:
+                df_result = self._add_metadata_column(df_result, response)
+            return df_result
+
+        except (APIClientHTTPError, APIClientConnectionError, APIClientTimeoutError) as e:
+            raise RuntimeError(f"API request failed: {e}") from e
+        except APIClientError as e:
+            raise RuntimeError(f"An error occurred while reading the API data: {e}") from e
+        except Exception as e:
+            raise RuntimeError(f"An unexpected error occurred: {e}") from e
+
+    def _add_metadata_column(self, df: DataFrame, response: APIResponse):
+        """Adds a metadata column to a DataFrame.
+
+        This method appends a column named `__metadata` to the given DataFrame, containing a map
+        of metadata related to an API response. The metadata includes the current timestamp,
+        the base URL of the API, the URL of the request, the HTTP status code, the reason phrase,
+        and the elapsed time of the request in seconds.
+
+        Args:
+            df: The DataFrame to which the metadata column will be added.
+            response: The API response object containing the metadata to be added.
+
+        Returns:
+            DataFrame: The original DataFrame with an added `__metadata` column containing the API response metadata.
+        """
+        df = df.withColumn(
+            "__metadata",
+            F.create_map(
+                F.lit("timestamp"),
+                F.current_timestamp(),
+                F.lit("base_url"),
+                F.lit(self.api_client.base_url),
+                F.lit("url"),
+                F.lit(response.url),
+                F.lit("status_code"),
+                F.lit(response.status_code),
+                F.lit("reason"),
+                F.lit(response.reason),
+                F.lit("elapsed"),
+                F.lit(response.elapsed),
+            ),
+        )
+        return df

cloe_nessy/integration/reader/catalog_reader.py

@@ -0,0 +1,49 @@
+from typing import Any
+
+from pyspark.sql import DataFrame
+from pyspark.sql.utils import AnalysisException
+
+from .exceptions import ReadOperationFailedError
+from .reader import BaseReader
+
+
+class CatalogReader(BaseReader):
+    """A reader for Unity Catalog objects.
+
+    This class reads data from a Unity Catalog table and loads it into a Spark DataFrame.
+    """
+
+    def __init__(self):
+        """Initializes the CatalogReader object."""
+        super().__init__()
+
+    def read(self, table_identifier: str = "", **kwargs: Any) -> DataFrame:
+        """Reads a table from the Unity Catalog.
+
+        Args:
+            table_identifier: The table identifier in the Unity Catalog in the format 'catalog.schema.table'.
+            **kwargs: This method does not accept any additional keyword arguments.
+
+        Returns:
+            The Spark DataFrame containing the read data.
+
+        Raises:
+            ValueError: If the table_identifier is not provided, is not a string, or is not in the correct format.
+            Exception: For any other unexpected errors.
+        """
+        if not table_identifier:
+            raise ValueError("table_identifier is required")
+        if not isinstance(table_identifier, str):
+            raise ValueError("table_identifier must be a string")
+        if len(table_identifier.split(".")) != 3:
+            raise ValueError("table_identifier must be in the format 'catalog.schema.table'")
+
+        try:
+            df = self._spark.read.table(table_identifier)
+            return df
+        except AnalysisException as err:
+            raise ValueError(f"Table not found: {table_identifier}") from err
+        except Exception as err:
+            raise ReadOperationFailedError(
+                f"An error occurred while reading the table '{table_identifier}': {err}"
+            ) from err

cloe_nessy/integration/reader/excel_reader.py

@@ -0,0 +1,170 @@
+from collections.abc import Callable
+from typing import Any
+
+import pandas as pd
+import pyspark.sql.functions as F
+from pyspark.sql import DataFrame
+
+from .reader import BaseReader
+
+
+class ExcelDataFrameReader(BaseReader):
+    """Utility class for reading an Excel file into a DataFrame.
+
+    This class uses the Pandas API on Spark to read Excel files to a DataFrame.
+    More information can be found in the [official
+    documentation](https://spark.apache.org/docs/latest/api/python/reference/pyspark.pandas/index.html).
+    """
+
+    def __init__(self):
+        """Initializes the ExcelDataFrameReader object."""
+        super().__init__()
+
+    def read_stream(self) -> DataFrame:
+        """Currently not implemented."""
+        raise NotImplementedError("Currently not implemented.")
+
+    def read(
+        self,
+        location: str,
+        *,
+        sheet_name: str | int | list = 0,
+        header: int | list[int] = 0,
+        index_col: int | list[int] | None = None,
+        usecols: int | str | list | Callable | None = None,
+        true_values: list | None = None,
+        false_values: list | None = None,
+        nrows: int | None = None,
+        na_values: list[str] | dict[str, list[str]] | None = None,
+        keep_default_na: bool = True,
+        parse_dates: bool | list | dict = False,
+        date_parser: Callable | None = None,
+        thousands: str | None = None,
+        options: dict | None = None,
+        load_as_strings: bool = False,
+        add_metadata_column: bool = False,
+        **kwargs: Any,
+    ) -> DataFrame:
+        """Reads Excel file on specified location and returns DataFrame.
+
+        Args:
+            location: Location of files to read.
+            sheet_name: Strings are used for sheet names.
+                Integers are used in zero-indexed sheet positions. Lists of
+                strings/integers are used to request multiple sheets. Specify None
+                to get all sheets.
+            header: Row to use for column labels. If a
+                list of integers is passed those row positions will be combined. Use
+                None if there is no header.
+            index_col: Column to use as the row labels of the
+                DataFrame. Pass None if there is no such column. If a list is
+                passed, those columns will be combined.
+            usecols: Return a subset of the columns. If
+                None, then parse all columns. If str, then indicates comma separated
+                list of Excel column letters and column ranges (e.g. “A:E” or
+                “A,C,E:F”). Ranges are inclusive of both sides. nIf list of int,
+                then indicates list of column numbers to be parsed. If list of
+                string, then indicates list of column names to be parsed. If
+                Callable, then evaluate each column name against it and parse the
+                column if the Callable returns True.
+            true_values: Values to consider as True.
+            false_values: Values to consider as False.
+            nrows: Number of rows to parse.
+            na_values: Additional strings to recognize as
+                NA/NaN. If dict passed, specific per-column NA values.
+            keep_default_na: If na_values are specified and
+                keep_default_na is False the default NaN values are overridden,
+                otherwise they're appended to.
+            parse_dates: The behavior is as follows:
+                - bool. If True -> try parsing the index.
+                - list of int or names. e.g. If [1, 2, 3] -> try parsing columns 1, 2, 3 each as a separate date column.
+                - list of lists. e.g. If [[1, 3]] -> combine columns 1 and 3 and parse as a single date column.
+                - dict, e.g. {{"foo" : [1, 3]}} -> parse columns 1, 3 as date and call result "foo"
+                If a column or index contains an unparseable date, the entire column or index will be returned unaltered as an object data type.
+            date_parser: Function to use for converting a sequence of
+                string columns to an array of datetime instances. The default uses
+                dateutil.parser.parser to do the conversion.
+            thousands: Thousands separator for parsing string columns to
+                numeric. Note that this parameter is only necessary for columns
+                stored as TEXT in Excel, any numeric columns will automatically be
+                parsed, regardless of display format.
+            options: Optional keyword arguments passed to
+                pyspark.pandas.read_excel and handed to TextFileReader.
+            load_as_strings: If True, converts all columns to string type to avoid datatype conversion errors in Spark.
+            add_metadata_column: If True, adds a metadata column containing the file location and sheet name.
+            kwargs: This method does not accept any additional keyword arguments.
+        """
+        if options is None:
+            options = {}
+        if ".xls" not in location:
+            raise ValueError(
+                "The excel reader can only be used for files with extension .xls. Use FileReader or some other reader instead."
+            )
+        try:
+            df = pd.read_excel(  # type: ignore
+                location,
+                sheet_name=sheet_name,
+                header=header,
+                index_col=index_col,
+                usecols=usecols,
+                true_values=true_values,
+                false_values=false_values,
+                nrows=nrows,
+                na_values=na_values,
+                keep_default_na=keep_default_na,
+                parse_dates=parse_dates,
+                date_parser=date_parser,
+                thousands=thousands,
+                dtype="string" if load_as_strings else None,
+                **options,
+            )
+            if isinstance(df, dict):
+                # in case pandas.read_excel returns a dict, union to single df
+                df = pd.concat(list(df.values()), ignore_index=True)
+
+        except FileNotFoundError:
+            self._console_logger.error(f"No xls(x) file was found at the specified location [ '{location}' ].")
+            raise
+        except Exception as e:
+            self._console_logger.error(f"read file [ '{location}' ] failed. Error: {e}")
+        else:
+            self._console_logger.info(f"Read file [ '{location}' ] succeeded.")
+
+        spark_df = self._spark.createDataFrame(df)
+        if add_metadata_column:
+            spark_df = self._add_metadata_column(df=spark_df, location=location, sheet_name=sheet_name)
+        return spark_df
+
+    def _add_metadata_column(self, df: DataFrame, location: str, sheet_name: str | int | list):
+        """Adds a metadata column to a DataFrame.
+
+        This method appends a column named `__metadata` to the given DataFrame, containing a map
+        of metadata related to the Excel file read operation. The metadata includes the current
+        timestamp, the location of the Excel file, and the sheet name(s) from which the data was read.
+
+        Args:
+            df: The DataFrame to which the metadata column will be added.
+            location: The file path of the Excel file.
+            sheet_name: The sheet name or sheet index used when reading the Excel file.
+
+        Returns:
+            DataFrame: The original DataFrame with an added `__metadata` column containing the Excel file metadata.
+        """
+        # Convert sheet_name to string if it is not already a string
+        if isinstance(sheet_name, list):
+            sheet_name = ", ".join(map(str, sheet_name))
+        else:
+            sheet_name = str(sheet_name)
+
+        df = df.withColumn(
+            "__metadata",
+            F.create_map(
+                F.lit("timestamp"),
+                F.current_timestamp(),
+                F.lit("file_location"),
+                F.lit(location),
+                F.lit("sheet_name"),
+                F.lit(sheet_name),
+            ),
+        )
+        return df

cloe_nessy/integration/reader/exceptions.py

@@ -0,0 +1,10 @@
+class ReaderError(Exception):
+    """Base class for reader exceptions."""
+
+    pass
+
+
+class ReadOperationFailedError(ReaderError):
+    """Exception raised for general Read Operation errors."""
+
+    pass