sibi-dst 0.3.44__py3-none-any.whl → 0.3.46__py3-none-any.whl

sibi_dst/v2/df_helper/_df_helper.py

@@ -0,0 +1,214 @@
+import warnings
+from typing import Any, Dict, Type, TypeVar, Union
+
+import dask.dataframe as dd
+import fsspec
+import pandas as pd
+from pydantic import BaseModel
+
+from sibi_dst.v2.utils import Logger
+from sibi_dst.v2.df_helper.core import QueryConfig, ParamsConfig, FilterHandler
+from sibi_dst.v2.df_helper.backends.sqlalchemy import SqlAlchemyConnectionConfig, SqlAlchemyLoadFromDb
+from sibi_dst.v2.df_helper.backends.sqlmodel import SQLModelConnectionConfig, SQLModelLoadFromDb
+
+# Define a generic type variable for BaseModel subclasses
+T = TypeVar("T", bound=BaseModel)
+
+# Suppress warnings about protected member access
+warnings.filterwarnings(
+    "ignore",
+    message="Access to a protected member _meta",
+    category=UserWarning,
+)
+
+
+class DfHelper:
+    df: Union[dd.DataFrame, pd.DataFrame] = None
+    default_config = {
+        'parquet_storage_path': None,
+        'dt_field': None,
+        'as_pandas': False,
+        'filesystem': 'file',
+        'filesystem_options': {},
+        'fs': fsspec.filesystem('file')
+    }
+
+    def __init__(self, **kwargs: Any) -> None:
+        # Merge default configuration with any provided kwargs
+        config = {**self.default_config.copy(), **kwargs}
+        self.backend = config.setdefault('backend', 'sqlalchemy')
+        self.debug = config.setdefault('debug', False)
+        self.as_pandas = config.setdefault('as_pandas', False)
+        self.logger = config.setdefault(
+            'logger',
+            Logger.default_logger(logger_name=self.__class__.__name__, debug=self.debug)
+        )
+        self.logger.debug("Logger initialized in DEBUG mode.")
+
+        # Propagate logger and debug settings to all components
+        config.setdefault('logger', self.logger)
+        config.setdefault('debug', self.debug)
+
+        self._initialize_backend_config(**config)
+
+    def __str__(self) -> str:
+        return self.__class__.__name__
+
+    def _extract_config_vars(self, model: Type[T], kwargs: Dict[str, Any]) -> T:
+        """
+        Extracts and initializes a Pydantic model using only the keys that the model accepts.
+        The recognized keys are removed from kwargs.
+        """
+        recognized_keys = set(model.__annotations__.keys())
+        self.logger.debug(f"Recognized keys for {model.__name__}: {recognized_keys}")
+        model_kwargs = {k: kwargs.pop(k) for k in list(kwargs.keys()) if k in recognized_keys}
+        self.logger.debug(f"Initializing {model.__name__} with: {model_kwargs}")
+        return model(**model_kwargs)
+
+    def _initialize_backend_config(self, **kwargs: Any) -> None:
+        """
+        Initializes the backend configurations by extracting the settings required for queries,
+        parameters, and SQLAlchemy connections.
+        """
+        self.logger.debug("Initializing backend configuration.")
+        self._backend_query = self._extract_config_vars(QueryConfig, kwargs)
+        self._backend_params = self._extract_config_vars(ParamsConfig, kwargs)
+        if self.backend == "sqlalchemy":
+            self.backend_connection_config = self._extract_config_vars(SqlAlchemyConnectionConfig, kwargs)
+        elif self.backend == "sqlmodel":
+            self.backend_connection_config = self._extract_config_vars(SQLModelConnectionConfig, kwargs)
+        else:
+            raise ValueError(f"Unsupported backend: {self.backend}")
+
+    def load(self, **options: Any) -> Union[dd.DataFrame, pd.DataFrame]:
+        """
+        Loads the data using the underlying SQLAlchemy loader. Returns a pandas DataFrame
+        if 'as_pandas' is True; otherwise returns a dask DataFrame.
+        """
+        df = self._load(**options)
+        return df.compute() if self.as_pandas else df
+
+    def _load(self, **options: Any) -> Union[dd.DataFrame, pd.DataFrame]:
+        self._backend_params.parse_params(options)
+        if self.backend == "sqlalchemy":
+            return self._load_from_sqlalchemy(**options)
+        elif self.backend == "sqlmodel":
+            return self._load_from_sqlmodel(**options)
+        else:
+            raise ValueError(f"Unsupported backend: {self.backend}")
+
+    def _load_from_sqlalchemy(self, **options: Any) -> Union[dd.DataFrame, pd.DataFrame]:
+        """
+        Loads data from a SQLAlchemy source. On failure, logs the error and returns an empty
+        DataFrame wrapped as a dask DataFrame.
+        """
+        try:
+            db_loader = SqlAlchemyLoadFromDb(
+                self.backend_connection_config,
+                self._backend_query,
+                self._backend_params,
+                self.debug,
+                self.logger,
+                **options
+            )
+            self.df = db_loader.build_and_load()
+            self._process_loaded_data()
+            self._post_process_df()
+            self.logger.debug("Data successfully loaded from SQLAlchemy database.")
+        except Exception as e:
+            self.logger.error(f"Failed to load data from SQLAlchemy database: {e}. Options: {options}")
+            # Optionally re-raise the exception if in debug mode
+            if self.debug:
+                raise
+            self.df = dd.from_pandas(pd.DataFrame(), npartitions=1)
+        return self.df
+
+    def _load_from_sqlmodel(self, **options: Any) -> Union[dd.DataFrame, pd.DataFrame]:
+        try:
+            db_loader = SQLModelLoadFromDb(
+                self.backend_connection_config,
+                self._backend_query,
+                self._backend_params,
+                self.debug,
+                self.logger,
+                **options
+            )
+            self.df = db_loader.build_and_load()
+            self._process_loaded_data()
+            self._post_process_df()
+            self.logger.debug("Data successfully loaded from SQLModel database.")
+        except Exception as e:
+            self.logger.error(f"Failed to load data from SQLModel database: {e}. Options: {options}")
+            if self.debug:
+                raise
+            self.df = dd.from_pandas(pd.DataFrame(), npartitions=1)
+        return self.df
+
+    def _post_process_df(self) -> None:
+        """
+        Post-processes the DataFrame by filtering columns, renaming them, setting the index,
+        and converting the index to datetime if requested.
+        """
+        df_params = self._backend_params.df_params
+        fieldnames = df_params.get("fieldnames")
+        index_col = df_params.get("index_col")
+        datetime_index = df_params.get("datetime_index", False)
+        column_names = df_params.get("column_names")
+
+        # Filter columns based on fieldnames
+        if fieldnames:
+            valid_fieldnames = [col for col in fieldnames if col in self.df.columns]
+            self.df = self.df[valid_fieldnames]
+
+        # Rename columns if column_names are provided
+        if column_names is not None:
+            if not fieldnames or len(fieldnames) != len(column_names):
+                raise ValueError(
+                    f"Length mismatch: fieldnames ({len(fieldnames) if fieldnames else 0}) and "
+                    f"column_names ({len(column_names)}) must match."
+                )
+            rename_mapping = dict(zip(fieldnames, column_names))
+            self.df = self.df.map_partitions(self._rename_columns, mapping=rename_mapping)
+
+        # Set the index column if specified
+        if index_col is not None:
+            if index_col in self.df.columns:
+                self.df = self.df.set_index(index_col)
+            else:
+                raise ValueError(f"Index column '{index_col}' not found in DataFrame.")
+
+        # Convert the index to datetime if required
+        if datetime_index and self.df.index.dtype != 'datetime64[ns]':
+            self.df = self.df.map_partitions(self._convert_index_to_datetime)
+
+        self.logger.debug("Post-processing of DataFrame completed.")
+
+    def _process_loaded_data(self) -> None:
+        """
+        Applies renaming logic based on the field map configuration.
+        Logs a warning for any missing columns, and only renames existing columns.
+        """
+        self.logger.debug(f"Processing loaded data; DataFrame type: {type(self.df)}")
+        if self.df.map_partitions(len).compute().sum() > 0:
+            field_map = self._backend_params.field_map or {}
+            if isinstance(field_map, dict):
+                rename_mapping = {k: v for k, v in field_map.items() if k in self.df.columns}
+                missing_columns = [k for k in field_map if k not in self.df.columns]
+                if missing_columns:
+                    self.logger.warning(
+                        f"The following columns in field_map are not in the DataFrame: {missing_columns}"
+                    )
+                if rename_mapping:
+                    self.df = self.df.map_partitions(self._rename_columns, mapping=rename_mapping)
+        self.logger.debug("Processing of loaded data completed.")
+
+    @staticmethod
+    def _rename_columns(df: pd.DataFrame, mapping: Dict[str, str]) -> pd.DataFrame:
+        """Helper function to rename columns in a DataFrame."""
+        return df.rename(columns=mapping)
+
+    @staticmethod
+    def _convert_index_to_datetime(df: pd.DataFrame) -> pd.DataFrame:
+        """Helper function to convert the DataFrame index to datetime."""
+        df.index = pd.to_datetime(df.index, errors='coerce')
+        return df

sibi_dst/v2/df_helper/backends/sqlalchemy/__init__.py

@@ -0,0 +1,10 @@
+from ._db_connection import SqlAlchemyConnectionConfig
+from ._model_builder import SqlAlchemyModelBuilder
+from ._load_from_db import SqlAlchemyLoadFromDb
+
+__all__ = [
+    'SqlAlchemyConnectionConfig',
+    'SqlAlchemyModelBuilder',
+    'SqlAlchemyLoadFromDb',
+]
+

sibi_dst/v2/df_helper/backends/sqlalchemy/_db_connection.py

@@ -0,0 +1,82 @@
+from typing import Any, Optional
+
+from pydantic import BaseModel, model_validator, ConfigDict
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine import Engine
+from sqlalchemy.exc import OperationalError
+
+from sibi_dst.v2.utils import Logger
+from ._model_builder import SqlAlchemyModelBuilder
+
+
+class SqlAlchemyConnectionConfig(BaseModel):
+    """
+    Configuration for establishing an SQLAlchemy database connection and dynamically building
+    an ORM model for a specific table.
+
+    Attributes:
+        connection_url (str): The URL used to connect to the database.
+        table_name (Optional[str]): The name of the table for which the model will be built.
+        model (Any): The dynamically built SQLAlchemy model.
+        engine (Optional[Engine]): The SQLAlchemy engine instance.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+    connection_url: str
+    table: Optional[str] = None
+    model: Any = None
+    engine: Optional[Engine] = None
+    debug: bool = False
+    logger: Optional[Logger] = None
+    add_relationships: bool = False
+    export_models: bool = False
+    export_file_name: str = 'models.py'
+
+    @model_validator(mode="after")
+    def validate_and_initialize(self) -> "SqlAlchemyConnectionConfig":
+        """
+        Validate the configuration, initialize the engine, test the connection, and build the model.
+
+        Raises:
+            ValueError: If `connection_url` or `table_name` is missing, or if the connection or model
+                        building fails.
+        """
+        self.logger = self.logger or Logger.default_logger(logger_name="sqlalchemy_connection", debug=self.debug)
+        self.logger.debug("Validating and initializing SQLAlchemy connection configuration.")
+        if not self.connection_url:
+            raise ValueError("`connection_url` must be provided.")
+
+        # Initialize the engine.
+        self.engine = create_engine(self.connection_url)
+        self.logger.debug(f"Engine created for URL")
+
+        # Validate the connection.
+        self.validate_connection()
+
+        if not self.table:
+            raise ValueError("`table` must be provided to build the model.")
+
+        try:
+            builder = SqlAlchemyModelBuilder(self.engine, self.table, self.add_relationships, self.debug, self.logger)
+            self.model = builder.build_model()
+            if self.export_models:
+                builder.export_models_to_file(self.export_file_name)
+            self.logger.debug(f"Successfully built model for table: {self.table}")
+        except Exception as e:
+            raise ValueError(f"Failed to build model for table {self.table}: {e}")
+
+        return self
+
+    def validate_connection(self) -> None:
+        """
+        Test the database connection by executing a simple query.
+
+        Raises:
+            ValueError: If the connection cannot be established.
+        """
+        try:
+            with self.engine.connect() as connection:
+                connection.execute(text("SELECT 1"))
+            self.logger.debug("Database connection validated.")
+        except OperationalError as e:
+            raise ValueError(f"Failed to connect to the database: {e}")

sibi_dst/v2/df_helper/backends/sqlalchemy/_io_dask.py

@@ -0,0 +1,135 @@
+import itertools
+
+import dask.dataframe as dd
+import pandas as pd
+from sqlalchemy import create_engine, inspect, select
+from sqlalchemy.orm import sessionmaker
+
+from sibi_dst.v2.df_helper.core import FilterHandler
+from sibi_dst.v2.utils import Logger
+
+
+class SQLAlchemyDask:
+    def __init__(self, model, filters, engine_url, chunk_size=1000, logger=None, debug=False):
+        """
+        Initialize with an SQLAlchemy query and database engine URL.
+
+        :param model: SQLAlchemy ORM model.
+        :param filters: Filters to apply on the query.
+        :param engine_url: Database connection string for SQLAlchemy engine.
+        :param chunk_size: Number of records per chunk for Dask partitions.
+        :param logger: Logger instance for logging.
+        :param debug: Whether to print detailed logs.
+        """
+        self.query = None
+        self.model = model
+        self.filters = filters
+        self.chunk_size = chunk_size
+        self.debug = debug
+        self.engine = create_engine(engine_url)
+        self.Session = sessionmaker(bind=self.engine)
+        self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
+        self.logger.set_level(logger.DEBUG if debug else logger.INFO)
+
+    @staticmethod
+    def infer_dtypes_from_model(model):
+        """
+        Infer data types for Dask DataFrame based on SQLAlchemy ORM model columns.
+        """
+        mapper = inspect(model)
+        sqlalchemy_to_dask_dtype = {
+            'INTEGER': 'Int64',
+            'SMALLINT': 'Int64',
+            'BIGINT': 'Int64',
+            'FLOAT': 'float64',
+            'NUMERIC': 'float64',
+            'BOOLEAN': 'bool',
+            'VARCHAR': 'object',
+            'TEXT': 'object',
+            'DATE': 'datetime64[ns]',
+            'DATETIME': 'datetime64[ns]',
+            'TIME': 'object',
+            'UUID': 'object',
+        }
+
+        dtypes = {}
+        for column in mapper.columns:
+            dtype = sqlalchemy_to_dask_dtype.get(str(column.type).upper(), 'object')
+            dtypes[column.name] = dtype
+
+        return dtypes
+
+    def read_frame(self, fillna_value=None):
+        """
+        Load data from an SQLAlchemy query into a Dask DataFrame.
+
+        :param fillna_value: Value to replace NaN or NULL values with, if any.
+        :return: Dask DataFrame.
+        """
+        with self.Session() as session:
+            try:
+                # Build query
+                self.query = select(self.model)
+                if self.filters:
+                    self.query = FilterHandler(backend="sqlalchemy", logger=self.logger, debug=self.debug).apply_filters(self.query,
+                                                                                                       model=self.model,
+                                                                                                       filters=self.filters)
+                else:
+                    n_records = 100
+                    self.query = self.query.limit(n_records)
+                self.logger.debug(f"query:{self.query}")
+                # Infer dtypes
+                dtypes = self.infer_dtypes_from_model(self.model)
+                # Get the column order from the SQLAlchemy model
+                ordered_columns = [column.name for column in self.model.__table__.columns]
+
+                # Execute query and fetch results in chunks
+                result_proxy = session.execute(self.query)
+                results = result_proxy.scalars().all()  # Fetch all rows
+                iterator = iter(results)
+
+                partitions = []
+
+                while True:
+                    chunk = list(itertools.islice(iterator, self.chunk_size))
+                    if not chunk:
+                        break
+
+                    # Convert chunk to Pandas DataFrame
+                    df = pd.DataFrame.from_records(
+                        [row._asdict() if hasattr(row, '_asdict') else row.__dict__ for row in chunk]
+                    )
+                    # Drop internal SQLAlchemy state if it exists
+                    df = df.loc[:, ~df.columns.str.contains('_sa_instance_state')]
+
+                    # Reorder columns to match the model's order
+                    df = df[ordered_columns]
+
+                    # Fill NaN values
+                    if fillna_value is not None:
+                        df = df.fillna(fillna_value)
+
+                    # Convert timezone-aware columns to naive
+                    for col in df.columns:
+                        if isinstance(df[col].dtype, pd.DatetimeTZDtype):
+                            df[col] = df[col].dt.tz_localize(None)
+
+                    # Apply inferred dtypes
+                    df = df.astype(dtypes)
+                    # Create a Dask partition
+                    partitions.append(dd.from_pandas(df, npartitions=1))
+
+                # Concatenate all partitions
+                if partitions:
+                    dask_df = dd.concat(partitions, axis=0, ignore_index=True)
+                else:
+                    dask_df = dd.from_pandas(pd.DataFrame(columns=ordered_columns), npartitions=1)
+
+                self.logger.debug(f"Loaded {len(dask_df)} rows into Dask DataFrame.")
+
+                return dask_df
+
+            except Exception as e:
+                self.logger.error(f"Error executing query: {str(e)}")
+                self.logger.error(self.query)
+                return dd.from_pandas(pd.DataFrame(columns=ordered_columns), npartitions=1)

sibi_dst/v2/df_helper/backends/sqlalchemy/_load_from_db.py

@@ -0,0 +1,142 @@
+import dask.dataframe as dd
+import pandas as pd
+
+from sibi_dst.v2.df_helper.core import ParamsConfig, QueryConfig
+from sibi_dst.v2.utils import Logger
+from ._io_dask import SQLAlchemyDask
+from ._db_connection import SqlAlchemyConnectionConfig
+
+
+class SqlAlchemyLoadFromDb:
+    """
+    The SqlAlchemyLoadFromDb class provides functionality to load data from a
+    database using SQLAlchemy into a Dask DataFrame. It is capable of handling
+    large datasets efficiently by utilizing the Dask framework for parallel
+    computations.
+
+    This class is initialized with a database connection configuration, query
+    configuration, optional parameters, and a logger. It can execute a query
+    using the specified configurations and read the results into a Dask
+    DataFrame. This is useful for processing and analyzing large-scale data.
+
+    :ivar df: Dask DataFrame to store the loaded data.
+    :type df: dd.DataFrame
+    :ivar db_connection: Database connection configuration object, containing details
+        such as the table, model, and engine to be used for the query.
+    :type db_connection: SqlAlchemyConnectionConfig
+    :ivar table_name: Name of the database table being queried.
+    :type table_name: str
+    :ivar model: SQLAlchemy model associated with the database connection.
+    :type model: sqlalchemy.ext.declarative.api.DeclarativeMeta
+    :ivar engine: SQLAlchemy engine used for executing queries.
+    :type engine: sqlalchemy.engine.base.Engine
+    :ivar logger: Logger instance for logging debug and error information.
+    :type logger: Logger
+    :ivar query_config: Query configuration, including query-related details such
+        as the SQL query or query settings.
+    :type query_config: QueryConfig
+    :ivar params_config: Parameters configuration, including filter parameters for
+        the query.
+    :type params_config: ParamsConfig
+    :ivar debug: Debug flag indicating whether debug mode is enabled.
+    :type debug: bool
+    :ivar chunk_size: Size of data chunks to process at a time.
+    :type chunk_size: int
+    """
+    df: dd.DataFrame = None
+
+    def __init__(
+            self,
+            plugin_sqlalchemy: SqlAlchemyConnectionConfig,  # Expected to be an instance of SqlAlchemyConnection
+            plugin_query: QueryConfig = None,
+            plugin_params: ParamsConfig = None,
+            debug: bool = False,
+            logger: Logger = None,
+            **kwargs,
+    ):
+        """
+        Initializes an instance of the class, setting up a database connection,
+        query configuration, parameter configuration, and other optional settings
+        like debugging and logging. The class aims to manage the integration and
+        interaction with SQLAlchemy-based database operations.
+
+        :param plugin_sqlalchemy:
+            The SQLAlchemy connection configuration object, which provides
+            the connection details like engine, table name, and model
+            associated with the database operations.
+        :param plugin_query:
+            The query configuration object, used to define specific query
+            options or rules. Defaults to None.
+        :param plugin_params:
+            The parameters configuration object, used for any additional
+            parameterized settings or configurations. Defaults to None.
+        :param logger:
+            Optional logger instance for logging purposes. If not provided,
+            a default logger is instantiated using the standard logging system.
+        :param kwargs:
+            Optional additional keyword arguments for customization. Can
+            include optional settings like `debug` mode or `chunk_size`
+            for batch operations.
+        """
+        self.db_connection = plugin_sqlalchemy
+        self.table_name = self.db_connection.table
+        self.model = self.db_connection.model
+        self.engine = self.db_connection.engine
+        self.debug = debug
+        self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__, debug=self.debug)
+        self.query_config = plugin_query
+        self.params_config = plugin_params
+        self.chunk_size = kwargs.pop("chunk_size", 1000)
+
+    def build_and_load(self) -> dd.DataFrame:
+        """
+        Builds and returns the resulting dataframe after calling the internal
+        build and load function. This method triggers the `_build_and_load`
+        function to process and prepare the data before returning it as
+        a dask dataframe.
+
+        :raises RuntimeError: If any error occurs during the build or load process.
+
+        :return: The processed data in a dask dataframe.
+        :rtype: dd.DataFrame
+        """
+        self._build_and_load()
+        return self.df
+
+    def _build_and_load(self) -> dd.DataFrame:
+        """
+        Builds and loads a Dask DataFrame from a SQLAlchemy-compatible source.
+
+        This method initializes a SQLAlchemyDask object with the provided model,
+        filters, engine URL, logger, chunk size, and debug configuration.
+        It attempts to load the data using the ``read_frame`` method of
+        SQLAlchemyDask. If the data cannot be loaded or the query returns
+        no rows, it creates and returns an empty Dask DataFrame.
+
+        :raises Exception: On failure to load data or to create a DataFrame.
+
+        :return: A Dask DataFrame object containing the queried data or an
+                 empty DataFrame if the query returns no results or fails.
+        :rtype: dask.dataframe.DataFrame
+        """
+        try:
+            self.df = SQLAlchemyDask(
+                model=self.model,
+                filters=self.params_config.filters,
+                engine_url=self.engine.url,
+                logger=self.logger,
+                chunk_size=self.chunk_size,
+                debug=self.debug
+            ).read_frame()
+
+            if self.df is None or len(self.df.head().index) == 0:
+                self.logger.debug("Query returned no results.")
+                dask_df = dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+                return dask_df
+            return self.df
+        except Exception as e:
+            self.logger.debug(f"Failed to load data into Dask DataFrame.{e}")
+            dask_df = dd.from_pandas(pd.DataFrame(), npartitions=1)
+
+            return dask_df