sibi-dst 0.3.63__py3-none-any.whl → 2025.1.1__py3-none-any.whl

sibi_dst/df_helper/backends/sqlalchemy/_sql_model_builder.py

@@ -1,193 +1,206 @@
 import re
+import keyword
+import threading
+from sqlalchemy import MetaData, Engine
+from sqlalchemy.orm import DeclarativeBase
 
-from sqlalchemy import MetaData, Table
-from sqlalchemy.orm import declarative_base, relationship
 
-# Base class for dynamically created models
-Base = declarative_base()
+class Base(DeclarativeBase):
+    """Shared declarative base for all ORM models."""
+    pass
 
-apps_label = "datacubes"
+
+apps_label = "datacubes.models"
 
 
 class SqlAlchemyModelBuilder:
     """
-    Provides functionality for building SQLAlchemy ORM models dynamically from
-    reflected database tables. This class is intended for use with a SQLAlchemy
-    engine and metadata to automatically generate ORM models for specified
-    database tables.
-
-    The primary purpose of this class is to simplify the process of creating
-    SQLAlchemy ORM models by reflecting tables from a connected database,
-    dynamically generating model classes, and handling relationships between
-    tables.
-
-    :ivar engine: SQLAlchemy engine connected to the database.
-    :type engine: Engine
-    :ivar table_name: Name of the table for which the model is generated.
-    :type table_name: str
-    :ivar metadata: SQLAlchemy MetaData instance for reflecting tables.
-    :type metadata: MetaData
-    :ivar table: Reflected SQLAlchemy Table object for the specified table name.
-    :type table: Optional[Table]
-    :ivar class_name: Dynamically normalized class name derived from table_name.
-    :type class_name: str
+    Builds a single SQLAlchemy ORM model from a specific database table.
+    This class is thread-safe and caches reflected table metadata to
+    improve performance across multiple instantiations.
     """
-    _model_cache = {}  # Local cache for model classes
+    _lock = threading.Lock()
+    _metadata_cache: dict[str, MetaData] = {}
 
-    def __init__(self, engine, table_name):
+    def __init__(self, engine: Engine, table_name: str):
         """
-        Initialize the model builder with a database engine and specific table.
+        Initializes the model builder for a specific table.
 
         Args:
-            engine: SQLAlchemy engine connected to the database.
-            table_name (str): Name of the table to generate the model for.
+            engine: The SQLAlchemy engine connected to the database.
+            table_name: The name of the table to generate the model for.
         """
         self.engine = engine
         self.table_name = table_name
-        self.metadata = MetaData()
-        self.table = None  # Placeholder for the specific table
-        self.class_name = self.normalize_class_name(self.table_name)
+        self.class_name = self._normalize_class_name(self.table_name)
 
-    def build_model(self) -> type:
-        """
-        Builds and returns a database model class corresponding to the specified table name.
-        The method checks if the model is already registered in the ORM's registry. If not,
-        it reflects the database schema of the specified table and dynamically creates the
-        model class.
-
-        :raises ValueError: If the specified table does not exist in the database.
-        :return: A database model class corresponding to the specified table name.
-        :rtype: type
-        """
-        # Check if the model is already registered
-        model = Base.registry._class_registry.get(self.class_name)
-        if model:
-            return model
+        engine_key = str(engine.url)
 
-        self.metadata.reflect(only=[self.table_name], bind=self.engine)
-        self.table = self.metadata.tables.get(self.table_name)
-        if self.table is None:
-            raise ValueError(f"Table '{self.table_name}' does not exist in the database.")
+        # ✅ REFACTOR: Acquire lock to make cache access and creation atomic,
+        # preventing a race condition between multiple threads.
+        with self._lock:
+            if engine_key not in self._metadata_cache:
+                self._metadata_cache[engine_key] = MetaData()
+            self.metadata = self._metadata_cache[engine_key]
 
-        model = self.create_model()
-        return model
-
-    def create_model(self) -> type:
+    def build_model(self) -> type:
         """
-        Generates a SQLAlchemy model class dynamically based on the specified table and
-        its columns. The method extracts column information, defines the necessary
-        attributes, and creates the model class if it doesn't already exist in the
-        SQLAlchemy base registry.
+        Builds and returns a database model class for the specified table.
+        This process is atomic and thread-safe.
 
-        :raises KeyError: If the table or table name does not exist in the provided
-            schema.
-        :raises Exception: If the model creation fails for any reason.
-
-        :return: The dynamically created or fetched model class.
-        :rtype: type
+        Raises:
+            ValueError: If the specified table does not exist in the database.
+        Returns:
+            The dynamically created ORM model class.
         """
-        # Normalize the class name from the table name
-        columns = self.get_columns(self.table)
-
-        # Define attributes for the model class
-        attrs = {
-            "__tablename__": self.table_name,
-            "__table__": self.table,
-            "__module__": f"{apps_label}.models",
-            "__mapper_args__": {"eager_defaults": True},
-        }
-
-        # Add columns and relationships to the model
-        attrs.update(columns)
-        #self.add_relationships(attrs, self.table)
-        model = Base.registry._class_registry.get(self.class_name)
-        if not model:
+        with self._lock:
+            # ✅ REFACTOR: Add a comment acknowledging the risk of using an
+            # internal API. This is a maintenance warning for future developers.
+            # NOTE: Using a private SQLAlchemy API. This is a performance
+            # optimization but may break in future versions of the library.
+            registered_model = Base.registry._class_registry.get(self.class_name)
+            if registered_model:
+                return registered_model
+
+            # Check if the table's schema is in our metadata cache
+            table = self.metadata.tables.get(self.table_name)
+
+            # If not cached, reflect it from the database
+            if table is None:
+                self.metadata.reflect(bind=self.engine, only=[self.table_name])
+                table = self.metadata.tables.get(self.table_name)
+
+            if table is None:
+                raise ValueError(
+                    f"Table '{self.table_name}' does not exist in the database."
+                )
+
+            # Create the model class dynamically.
+            attrs = {
+                "__tablename__": table.name,
+                "__table__": table,
+                "__module__": apps_label,
+            }
             model = type(self.class_name, (Base,), attrs)
-            # Add the class to Base.registry so it is registered
-            Base.registry._class_registry[self.class_name] = model
-        return model
-
-    def get_columns(self, table: Table):
-        """
-        Extracts and returns a dictionary of column names and their corresponding column
-        objects from a given table, excluding reserved names. Reserved names are used
-        internally and should not overlap with column names in the provided table. The
-        method ensures sanitized column names through normalization and filters out any
-        column matching reserved keywords.
-
-        :param table: The table object from which columns are to be extracted.
-        :type table: Table
-        :return: A dictionary containing the sanitized column names as keys and their
-            corresponding column objects as values, excluding reserved names.
-        :rtype: dict
-        """
-        columns = {}
-        reserved_names = ["metadata", "class_", "table"]
-
-        for column in table.columns:
-            column_name = self.normalize_column_name(column.name)
-            if column_name not in reserved_names:
-                columns[column_name] = column
-        return columns
-
-    def add_relationships(self, attrs, table: Table):
-        """
-        Adds relationships to the provided attributes dictionary for a given database table.
-
-        This method iterates through the foreign keys of the provided table, constructs
-        relationship attributes, and updates the attributes dictionary with relationships
-        that connect the current table to related tables.
-
-        :param attrs: Dictionary of attributes to which relationships will be added.
-                      The dictionary will be updated with new relationship mappings.
-        :type attrs: dict
-        :param table: A database table object containing foreign key relationships.
-                      The method will use this table to establish relationships.
-        :return: None
-        """
-        for fk in table.foreign_keys:
-            related_table_name = fk.column.table.name
-            related_class_name = self.normalize_class_name(related_table_name)
-            relationship_name = self.normalize_column_name(related_table_name)
-            attrs[relationship_name] = relationship(related_class_name, back_populates=None)
 
+            return model
 
     @staticmethod
-    def normalize_class_name(table_name: str) -> str:
-        """
-        Generate a normalized class name from a given table name by capitalizing
-        each word separated by underscores and concatenating them.
-
-        This static method takes a string representation of a table name, where
-        words are separated by underscores, and converts it into a camel case
-        class name. It processes the string by capitalizing the first letter of
-        each word and removing the underscores. The normalized class name
-        returned can be used programmatically for various purposes, such as
-        class generation or naming conventions.
-
-        :param table_name: The table name to normalize, with words separated by
-            underscores. E.g., 'sample_table' becomes 'SampleTable'.
-        :type table_name: str
-        :return: A normalized class name in camel case format.
-        :rtype: str
-        """
+    def _normalize_class_name(table_name: str) -> str:
+        """Converts a snake_case table_name to a CamelCase class name."""
         return "".join(word.capitalize() for word in table_name.split("_"))
 
     @staticmethod
-    def normalize_column_name(column_name: str) -> str:
+    def _normalize_column_name(column_name: str) -> str:
         """
-        Normalize a column name by replacing any non-word characters or leading numbers
-        with underscores, while ensuring it does not conflict with reserved keywords
-        such as 'class', 'def', 'return', etc. If the normalized name conflicts with
-        a Python reserved keyword, "_field" is appended to it.
-
-        :param column_name: The original name of the column to be normalized.
-        :type column_name: str
-        :return: A normalized column name that is safe and compatible for usage
-            in various contexts such as database columns or Python code.
-        :rtype: str
+        Sanitizes a column name to be a valid Python identifier.
+        (Kept for utility, though not used in the final model creation).
         """
-        column_name = re.sub(r"\W|^(?=\d)", "_", column_name)
-        if column_name in {"class", "def", "return", "yield", "global"}:
-            column_name += "_field"
-        return column_name
+        sane_name = re.sub(r"\W", "_", column_name)
+        sane_name = re.sub(r"^\d", r"_\g<0>", sane_name)
+
+        if keyword.iskeyword(sane_name):
+            return f"{sane_name}_field"
+        return sane_name
+
+# import re
+# import keyword
+# import threading
+# from sqlalchemy import MetaData, Engine
+# from sqlalchemy.orm import DeclarativeBase
+#
+#
+#
+# class Base(DeclarativeBase):
+#     """shared declarative base for all ORM models."""
+#     pass
+#
+#
+# apps_label = "datacubes.models"
+#
+#
+# class SqlAlchemyModelBuilder:
+#     """
+#     Builds a single SQLAlchemy ORM model from a specific database table.
+#     This class is thread-safe and caches reflected table metadata to
+#     improve performance across multiple instantiations.
+#     """
+#     _lock = threading.Lock()
+#     _metadata_cache: dict[str, MetaData] = {}
+#
+#     def __init__(self, engine: Engine, table_name: str):
+#         """
+#         Initializes the model builder for a specific table.
+#
+#         Args:
+#             engine: The SQLAlchemy engine connected to the database.
+#             table_name: The name of the table to generate the model for.
+#         """
+#         self.engine = engine
+#         self.table_name = table_name
+#         self.class_name = self._normalize_class_name(self.table_name)
+#
+#         # Use or create a cached MetaData object for this engine to avoid
+#         # re-reading the schema for tables that are already known.
+#         engine_key = str(engine.url)
+#         if engine_key not in self._metadata_cache:
+#             self._metadata_cache[engine_key] = MetaData()
+#         self.metadata = self._metadata_cache[engine_key]
+#
+#     def build_model(self) -> type:
+#         """
+#         Builds and returns a database model class for the specified table.
+#         This process is atomic and thread-safe.
+#
+#         Raises:
+#             ValueError: If the specified table does not exist in the database.
+#         Returns:
+#             The dynamically created ORM model class.
+#         """
+#         with self._lock:
+#             # First, check if the model class is already registered in SQLAlchemy
+#             registered_model = Base.registry._class_registry.get(self.class_name)
+#             if registered_model:
+#                 return registered_model
+#
+#             # Next, check if the table's schema is in our metadata cache
+#             table = self.metadata.tables.get(self.table_name)
+#
+#             # If not cached, reflect it from the database
+#             if table is None:
+#                 self.metadata.reflect(bind=self.engine, only=[self.table_name])
+#                 table = self.metadata.tables.get(self.table_name)
+#
+#             if table is None:
+#                 raise ValueError(
+#                     f"Table '{self.table_name}' does not exist in the database."
+#                 )
+#
+#             # Create the model class dynamically.
+#             # No need to add columns manually; __table__ handles it.
+#             attrs = {
+#                 "__tablename__": table.name,
+#                 "__table__": table,
+#                 "__module__": apps_label,
+#             }
+#             model = type(self.class_name, (Base,), attrs)
+#
+#             return model
+#
+#     @staticmethod
+#     def _normalize_class_name(table_name: str) -> str:
+#         """Converts a snake_case table_name to a CamelCase class name."""
+#         return "".join(word.capitalize() for word in table_name.split("_"))
+#
+#     @staticmethod
+#     def _normalize_column_name(column_name: str) -> str:
+#         """
+#         Sanitizes a column name to be a valid Python identifier.
+#         (Kept for utility, though not used in the final model creation).
+#         """
+#         sane_name = re.sub(r"\W", "_", column_name)
+#         sane_name = re.sub(r"^\d", r"_\g<0>", sane_name)
+#
+#         if keyword.iskeyword(sane_name):
+#             return f"{sane_name}_field"
+#         return sane_name

sibi_dst/df_helper/core/__init__.py

@@ -1,8 +1,6 @@
 from __future__ import annotations
 
 from ._defaults import (
-    django_field_conversion_map_pandas,
-    django_field_conversion_map_dask,
     sqlalchemy_field_conversion_map_dask,
     normalize_sqlalchemy_type)
 from ._filter_handler import FilterHandler
@@ -12,8 +10,6 @@ from ._query_config import QueryConfig
 __all__ = [
     "ParamsConfig",
     "QueryConfig",
-    "django_field_conversion_map_pandas",
-    "django_field_conversion_map_dask",
     "sqlalchemy_field_conversion_map_dask",
     "normalize_sqlalchemy_type",
     "FilterHandler",

sibi_dst/df_helper/core/_defaults.py

@@ -13,56 +13,7 @@ from sqlalchemy.dialects.mysql import TINYINT, MEDIUMTEXT
 # conversion_map is a dictionary that maps the field types to their corresponding data type conversion functions.
 # Each entry in the dictionary is a pair of a field type (as a string) and a callable function that performs the
 # conversion. This mapping is used to convert the values in a pandas DataFrame to the appropriate data types based on
-# the Django field type.
-
-django_field_conversion_map_pandas: Dict[str, callable] = {
-    "CharField": lambda x: x.astype(str),
-    "TextField": lambda x: x.astype(str),
-    "IntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "AutoField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "BigAutoField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "BigIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "SmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "PositiveIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "PositiveSmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "FloatField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "DecimalField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "BooleanField": lambda x: x.astype(bool),
-    "NullBooleanField": lambda x: x.astype(bool),
-    "DateTimeField": lambda x: pd.to_datetime(x, errors="coerce"),
-    "DateField": lambda x: pd.to_datetime(x, errors="coerce").dt.date,
-    "TimeField": lambda x: pd.to_datetime(x, errors="coerce").dt.time,
-    "DurationField": lambda x: pd.to_timedelta(x, errors="coerce"),
-    # for JSONField, assuming JSON objects are represented as string in df
-    "JSONField": lambda x: x.apply(json.loads),
-    "ArrayField": lambda x: x.apply(eval),
-    "UUIDField": lambda x: x.astype(str),
-}
-
-django_field_conversion_map_dask: Dict[str, callable] = {
-    "CharField": lambda x: x.astype(str),
-    "TextField": lambda x: x.astype(str),
-    "IntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "AutoField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "BigAutoField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "BigIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "SmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "PositiveIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "PositiveSmallIntegerField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "FloatField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "DecimalField": lambda x: pd.to_numeric(x, errors="coerce"),
-    "BooleanField": lambda x: x.astype(bool),
-    "NullBooleanField": lambda x: x.astype(bool),
-    "DateTimeField": lambda x: pd.to_datetime(x, errors="coerce"),
-    "DateField": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda x: x.dt.date,
-                                                                             meta=("date", "object")),
-    "TimeField": lambda x: pd.to_datetime(x, errors="coerce").map_partitions(lambda x: x.dt.time,
-                                                                             meta=("time", "object")),
-    "DurationField": lambda x: pd.to_timedelta(x, errors="coerce"),
-    "JSONField": lambda x: x.map_partitions(lambda s: s.apply(json.loads), meta=("json", "object")),
-    "ArrayField": lambda x: x.map_partitions(lambda s: s.apply(eval), meta=("array", "object")),
-    "UUIDField": lambda x: x.astype(str),
-}
+# the db field type.
 
 sqlalchemy_field_conversion_map_dask: Dict[str, callable] = {
     String.__name__: lambda x: x.astype(str).fillna(""),

sibi_dst/df_helper/core/_query_config.py

@@ -7,8 +7,8 @@ class QueryConfig(BaseModel):
     use_exclude: bool = False
     n_records: int = 100
     dt_field: Optional[str] = None
-    use_dask: bool = False
-    as_dask: bool = False
+    use_dask: bool = True
+    as_dask: bool = True
 
     @model_validator(mode='after')
     def check_n_records(self):

sibi_dst/utils/__init__.py

@@ -10,7 +10,6 @@ from .df_utils import DfUtils
 from .storage_manager import StorageManager
 from .parquet_saver import ParquetSaver
 from .clickhouse_writer import ClickHouseWriter
-from .airflow_manager import AirflowDAGManager
 from .credentials import *
 from .update_planner import UpdatePlanner
 from .data_wrapper import DataWrapper
@@ -35,7 +34,6 @@ __all__ = [
     "StorageManager",
     "DfUtils",
     "ClickHouseWriter",
-    "AirflowDAGManager",
     "StorageConfig",
     "FsRegistry",
     "DataFromHttpSource",

sibi_dst/utils/data_wrapper.py

@@ -38,7 +38,7 @@ class DataWrapper:
             logger: Logger = None,
             show_progress: bool = False,
             timeout: float = 30,
-            max_threads: int = 1,
+            max_threads: int = 3,
             **kwargs: Any,
     ):
         self.dataclass = dataclass
@@ -66,6 +66,7 @@ class DataWrapper:
         self.benchmarks: Dict[datetime.date, Dict[str, float]] = {}
         self.mmanifest = kwargs.get("mmanifest", None)
         self.update_planner=kwargs.get("update_planner", None)
+        self.datacls = self.dataclass(**self.class_params)
 
     def __enter__(self):
         """Context manager entry"""
@@ -164,28 +165,24 @@ class DataWrapper:
     def _process_single_date(self, date: datetime.date):
         """Core date processing logic with load/save timing and thread reporting"""
         path = f"{self.data_path}{date.year}/{date.month:02d}/{date.day:02d}/"
-        self.logger.info(f"Processing date {date.isoformat()} for {path}")
-        # self.logger.info(f"Path {path} in {self.skipped}: {path in self.skipped}")
+        self.logger.debug(f"Processing date {date.isoformat()} for {path}")
         if path in self.update_planner.skipped and self.update_planner.ignore_missing:
             self.logger.info(f"Skipping {date} as it exists in the skipped list")
             return
         full_path = f"{path}{self.parquet_filename}"
 
         thread_name = threading.current_thread().name
-        self.logger.info(f"[{thread_name}] Executing date: {date} -> saving to: {full_path}")
+        self.logger.debug(f"[{thread_name}] Executing date: {date} -> saving to: {full_path}")
 
         overall_start = time.perf_counter()
         try:
             load_start = time.perf_counter()
-            with self.dataclass(**self.class_params) as data:
-                df = data.load_period(
-                    dt_field=self.date_field,
-                    start=date,
-                    end=date,
-                    **self.load_params
-                )
+            date_filter = {f"{self.date_field}__date": {date.isoformat()}}
+            self.logger.debug(f"Loading data for {date} with filter: {date_filter}")
+            # Load data using the dataclass with the provided date filter
+            self.load_params.update(date_filter)
+            df = self.datacls.load(**self.load_params)
             load_time = time.perf_counter() - load_start
-
             if df.head(1, compute=True).empty:
                 if self.mmanifest:
                     schema = df._meta.dtypes.astype(str).to_dict()

sibi_dst/utils/log_utils.py

@@ -115,22 +115,26 @@ class Logger:
         """
         self.logger.setLevel(level)
 
-    def debug(self, msg: str):
+    def debug(self, msg: str, *args, **kwargs):
         """Log a debug message."""
-        self.logger.debug(msg)
+        self.logger.debug(msg, *args, **kwargs)
 
-    def info(self, msg: str):
+    def info(self, msg: str, *args, **kwargs):
         """Log an info message."""
-        self.logger.info(msg)
+        self.logger.info(msg, *args, **kwargs)
 
-    def warning(self, msg: str):
+    def warning(self, msg: str, *args, **kwargs):
         """Log a warning message."""
-        self.logger.warning(msg)
+        self.logger.warning(msg, *args, **kwargs)
 
-    def error(self, msg: str):
-        """Log an error message."""
-        self.logger.error(msg)
+    def error(self, msg: str, *args, **kwargs):
+        """
+        Log an error message.
+
+        To log exception information, use the `exc_info=True` keyword argument.
+        """
+        self.logger.error(msg, *args, **kwargs)
 
-    def critical(self, msg: str):
+    def critical(self, msg: str, *args, **kwargs):
         """Log a critical message."""
-        self.logger.critical(msg)
+        self.logger.critical(msg, *args, **kwargs)

sibi_dst/utils/update_planner.py

@@ -73,6 +73,8 @@ class UpdatePlanner:
         self.show_progress = show_progress
         self.logger = logger or Logger.default_logger(logger_name="update_planner")
         self.logger.set_level(Logger.DEBUG if debug else Logger.INFO)
+        self.debug = debug
+        self.verbose = verbose
 
         # Filesystem and age helper
         self.fs = fs or fsspec.filesystem(filesystem_type, **(filesystem_options or {}))