sibi-dst 2025.1.3__py3-none-any.whl → 2025.1.5__py3-none-any.whl

sibi_dst/df_helper/_df_helper.py

@@ -4,12 +4,12 @@ import warnings
 from typing import Any, Dict, Optional, Union, TypeVar
 
 import dask.dataframe as dd
-import fsspec
 import pandas as pd
 from pydantic import BaseModel
 
 from sibi_dst.df_helper.core import QueryConfig, ParamsConfig, FilterHandler
-from sibi_dst.utils import Logger, ParquetSaver, ClickHouseWriter
+from sibi_dst.utils import ManagedResource
+from sibi_dst.utils import ParquetSaver, ClickHouseWriter
 from .backends.http import HttpConfig
 from .backends.parquet import ParquetConfig
 from .backends.sqlalchemy import SqlAlchemyConnectionConfig, SqlAlchemyLoadFromDb
@@ -27,18 +27,19 @@ class BaseBackend:
         self.helper = helper
         self.logger = helper.logger
         self.debug = helper.debug
+        self.total_records = helper.total_records  # no records loaded yet
 
-    def load(self, **options) -> dd.DataFrame | pd.DataFrame:
+    def load(self, **options) -> tuple[Any, Any] | Union[dd.DataFrame | pd.DataFrame]:
         """Synchronous data loading method. Must be implemented by sync backends."""
         raise NotImplementedError(f"Backend '{self.__class__.__name__}' does not support synchronous loading.")
 
-    async def aload(self, **options) -> dd.DataFrame | pd.DataFrame:
+    async def aload(self, **options) -> tuple[Any, Any] | Union[dd.DataFrame | pd.DataFrame]:
         """Asynchronous data loading method. By default, it calls the sync version."""
         return self.load(**options)
 
 
 class SqlAlchemyBackend(BaseBackend):
-    def load(self, **options) -> dd.DataFrame:
+    def load(self, **options) -> tuple[Any, Any] | Union[dd.DataFrame | pd.DataFrame]:
         try:
             # Process incoming filter options into the ParamsConfig object
             if options and hasattr(self.helper._backend_params, 'parse_params'):
@@ -49,45 +50,51 @@ class SqlAlchemyBackend(BaseBackend):
                 plugin_query=self.helper._backend_query,
                 plugin_params=self.helper._backend_params,
                 logger=self.logger,
-                debug= self.debug
+                debug=self.debug
             )
-            return db_loader.build_and_load()
+            self.total_records, result = db_loader.build_and_load()
+            return self.total_records, result
         except Exception as e:
             self.logger.error(f"Failed to load data from sqlalchemy: {e}", exc_info=self.debug)
-            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+            return -1, dd.from_pandas(pd.DataFrame(), npartitions=1)
 
 
 class ParquetBackend(BaseBackend):
     """This backend is also purely synchronous."""
 
-    def load(self, **options) -> dd.DataFrame | pd.DataFrame:
+    def load(self, **options) -> tuple[Any, Any] | Union[dd.DataFrame | pd.DataFrame]:
         try:
             df = self.helper.backend_parquet.load_files()
             if options and df is not None:
                 df = FilterHandler('dask', logger=self.logger, debug=False).apply_filters(df, filters=options)
-            return df
+            self.total_records = len(df)
+            return self.total_records, df
         except Exception as e:
+            self.total_records = -1  # Reset total_records on failure
             self.logger.error(f"Failed to load data from parquet: {e}", exc_info=True)
-            return dd.from_pandas(pd.DataFrame(), npartitions=1)
+            return -1, dd.from_pandas(pd.DataFrame(), npartitions=1)
 
 
 class HttpBackend(BaseBackend):
     """This backend is purely asynchronous."""
 
-    def load(self, **options) -> dd.DataFrame | pd.DataFrame:
+    def load(self, **options) -> tuple[Any, Any] | Union[dd.DataFrame | pd.DataFrame]:
         # This will correctly fail by raising NotImplementedError from the base class.
         return self.helper.backend_http.fetch_data(**options)
 
-    async def aload(self, **options) -> Union[pd.DataFrame, dd.DataFrame]:
+    async def aload(self, **options) -> tuple[Any, Any] | Union[pd.DataFrame, dd.DataFrame]:
         if not self.helper.backend_http:
             self.logger.warning("HTTP plugin not configured properly.")
-            return dd.from_pandas(pd.DataFrame(), npartitions=1)
-        return await self.helper.backend_http.fetch_data(**options)
+            self.total_records = -1
+            return self.total_records, dd.from_pandas(pd.DataFrame(), npartitions=1)
+        result = await self.helper.backend_http.fetch_data(**options)
+        self.total_records = len(result)
+        return self.total_records, result
 
 
 # --- Main DfHelper Facade Class ---
 
-class DfHelper:
+class DfHelper(ManagedResource):
     """
     A reusable utility for loading data. It provides both sync (`load`) and
     async (`aload`) methods to accommodate different backends.
@@ -103,13 +110,14 @@ class DfHelper:
     def __init__(self, backend='sqlalchemy', **kwargs):
         self.default_config = self.default_config or {}
         kwargs = {**self.default_config.copy(), **kwargs}
+        super().__init__(**kwargs)
         self.backend = backend
-        self.debug = kwargs.get("debug", False)
-        self.logger = kwargs.get("logger", Logger.default_logger(logger_name=self.__class__.__name__))
-        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
-        self.fs = kwargs.get("fs", fsspec.filesystem('file'))
+
+        # Need to set default values for backend-specific configurations
+        kwargs.setdefault("debug", self.debug)
         kwargs.setdefault("fs", self.fs)
         kwargs.setdefault("logger", self.logger)
+        self.total_records = -1  # Initialize total_records to -1 to indicate no records loaded yet
         self._backend_query = self._get_config(QueryConfig, kwargs)
         self._backend_params = self._get_config(ParamsConfig, kwargs)
         self.backend_db_connection: Optional[SqlAlchemyConnectionConfig] = None
@@ -127,11 +135,9 @@ class DfHelper:
         if not strategy_class: raise ValueError(f"Unsupported backend: {self.backend}")
         self.backend_strategy = strategy_class(self)
 
-    def __enter__(self):
-        return self
-
     def __exit__(self, exc_type, exc_value, traceback):
         self._cleanup()
+        super().__exit__(exc_type, exc_value, traceback)
 
     def _cleanup(self):
         active_config = getattr(self, f"backend_{self.backend}", None)
@@ -147,14 +153,14 @@ class DfHelper:
     def load(self, as_pandas=False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
         """Loads data synchronously. Fails if backend is async-only."""
         self.logger.debug(f"Loading data from {self.backend} backend with options: {options}")
-        df = self.backend_strategy.load(**options)
+        self.total_records, df = self.backend_strategy.load(**options)
         df = self._process_loaded_data(df)
         df = self._post_process_df(df)
         return df.compute() if as_pandas else df
 
     async def aload(self, as_pandas=False, **options) -> Union[pd.DataFrame, dd.DataFrame]:
         """Loads data asynchronously from any backend."""
-        df = await self.backend_strategy.aload(**options)
+        self.total_records, df = await self.backend_strategy.aload(**options)
         df = self._process_loaded_data(df)
         df = self._post_process_df(df)
         return df.compute() if as_pandas else df
@@ -195,15 +201,27 @@ class DfHelper:
             return
         fs = kwargs.pop('fs', self.fs)
         path = kwargs.pop('parquet_storage_path', self.backend_parquet.parquet_storage_path)
-        ParquetSaver(df, path, self.logger, fs).save_to_parquet(parquet_filename)
-        self.logger.debug(f"Parquet saved to {parquet_filename} in path: {path}.")
+        writer_config = {
+            'df_result': df,
+            'parquet_storage_path': path,
+            'fs': fs,
+            'debug': self.debug,
+            'logger': self.logger,
+            'verbose': self.verbose,
+        }
+        with ParquetSaver(**writer_config) as saver:
+            saver.save_to_parquet(parquet_filename)
+
+        self.logger.debug(f"Parquet saved to {parquet_filename} in {path}.")
 
     def save_to_clickhouse(self, df: dd.DataFrame, **credentials):
         if hasattr(df, 'npartitions') and df.npartitions == 1 and not len(df.head(1)):
             self.logger.warning("Cannot write to ClickHouse; DataFrame is empty.")
             return
-        ClickHouseWriter(self.logger, **credentials).save_to_clickhouse(df)
-        self.logger.debug("Save to ClickHouse completed.")
+
+        with ClickHouseWriter(debug=self.debug, logger=self.logger, verbose=self.verbose, **credentials) as writer:
+            writer.save_to_clickhouse(df)
+            self.logger.debug("Save to ClickHouse completed.")
 
     def load_period(self, dt_field: str, start: str, end: str, **kwargs) -> Union[pd.DataFrame, dd.DataFrame]:
         """Synchronous convenience method for loading a date range."""
@@ -229,4 +247,3 @@ class DfHelper:
             kwargs[f"{mapped_field}__date__range"] = [start_date, end_date]
         self.logger.debug(f"Period load generated filters: {kwargs}")
         return kwargs
-

sibi_dst/df_helper/_parquet_artifact.py

@@ -1,3 +1,4 @@
+from __future__ import annotations
 import datetime
 import logging
 import threading
@@ -7,7 +8,7 @@ import dask.dataframe as dd
 import fsspec
 
 from sibi_dst.df_helper import DfHelper
-from sibi_dst.utils import DataWrapper, DateUtils, Logger, ParquetSaver, UpdatePlanner
+from sibi_dst.utils import DataWrapper, DateUtils, UpdatePlanner
 from sibi_dst.utils import MissingManifestManager
 
 
@@ -58,14 +59,6 @@ class ParquetArtifact(DfHelper):
         'backend': 'parquet'
     }
 
-    # DEFAULT_UPDATE_PLANNER_CONFIG = {
-    #     'reverse_order': True,
-    #     'overwrite': False,
-    #     'ignore_missing': True,
-    #     'history_days_threshold': 30,
-    #     'max_age_minutes': 10,
-    #     'show_progress': False
-    # }
 
     def __init__(self, data_wrapper_class, **kwargs):
         """
@@ -97,7 +90,9 @@ class ParquetArtifact(DfHelper):
             **kwargs,
         }
         self.df: Optional[dd.DataFrame] = None
-        self._setup_logging()
+        super().__init__(**self.config)
+        #self._own_logger = False
+        #self._setup_logging()
         self.data_wrapper_class = data_wrapper_class
 
         self.date_field = self._validate_required('date_field')
@@ -107,35 +102,37 @@ class ParquetArtifact(DfHelper):
         self.parquet_end_date = self._validate_required('parquet_end_date')
 
         # Filesystem setup
-        self.filesystem_type = self.config.setdefault('filesystem_type', 'file')
-        self.filesystem_options = self.config.setdefault('filesystem_options', {})
-        self.fs = self.config.setdefault('fs', None)
-        self._own_fs = self.fs is None
-        if self.fs is None:
-            self.fs = fsspec.filesystem(self.filesystem_type, **self.filesystem_options)
-            self._own_fs = True
-        self.config.setdefault('fs', self.fs)
+        #self.filesystem_type = self.config.setdefault('filesystem_type', 'file')
+        #self.filesystem_options = self.config.setdefault('filesystem_options', {})
+        #self.fs = self.config.setdefault('fs', None)
+        #self._own_fs = self.fs is None
+        #if self.fs is None:
+        #    self.fs = fsspec.filesystem(self.filesystem_type, **self.filesystem_options)
+        #    self._own_fs = True
+        #self.config.setdefault('fs', self.fs)
         ## Populate to parameters to pass to data_wrapper_class
         self.class_params = self.config.pop('class_params', {
             'debug': self.debug,
             'logger': self.logger,
             'fs': self.fs,
+            'verbose': self.verbose,
         })
         # Populate parameters to pass to load method of DataWrapper class
         self.load_params = self.config.setdefault('load_params', {})
         # Ensure the directory exists
         self.ensure_directory_exists(self.parquet_storage_path)
-        super().__init__(**self.config)
+        #super().__init__(**self.config)
         self.update_planner_params = {}
         self.datawrapper_params = {}
 
-    def _setup_logging(self):
-        """Initialize logger and debug settings."""
-        self.debug = self.config.get('debug', False)
-        self.logger = self.config.get('logger',
-                                      Logger.default_logger(
-                                          logger_name=f'Parquet_Artifact_InstanceOf_{self.__class__.__name__}'))
-        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
+    #def _setup_logging(self):
+    #    """Initialize logger and debug settings."""
+    #    self.debug = self.config.get('debug', False)
+    #    logger = self.config.get('logger', None)
+    #    self._own_logger = logger is None
+    #    self.logger = logger or Logger.default_logger(
+    #                                      logger_name=f'Parquet_Artifact_InstanceOf_{self.__class__.__name__}')
+    #    self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
 
     def _validate_required(self, key: str) -> Any:
         """Validate required configuration fields."""
@@ -150,7 +147,9 @@ class ParquetArtifact(DfHelper):
         self.mmanifest = MissingManifestManager(
             fs=self.fs,
             manifest_path=self.missing_manifest_path,
-            clear_existing=overwrite
+            clear_existing=overwrite,
+            debug= self.debug,
+            logger=self.logger
         )
 
         # Initialize skipped files
@@ -158,7 +157,7 @@ class ParquetArtifact(DfHelper):
         if not manifest_exists:
             self.logger.info(f"Creating new manifest at {self.missing_manifest_path}")
             self.mmanifest.save()
-            self.mmanifest.cleanup_temp_manifests()
+            #self.mmanifest.cleanup_temp_manifests()
         else:
             self.logger.info(f"Manifest already exists at {self.missing_manifest_path}")
 
@@ -172,7 +171,7 @@ class ParquetArtifact(DfHelper):
     def _setup_update_planner(self, **kwargs) -> None:
         self._prepare_update_params(**kwargs)
         self.update_planner = UpdatePlanner(**self.update_planner_params)
-        self.update_planner.generate_plan(self.start_date, self.end_date)
+        self.update_planner.generate_plan(start=self.start_date,end= self.end_date)
 
     def load(self, **kwargs):
         with self._lock:
@@ -197,36 +196,45 @@ class ParquetArtifact(DfHelper):
             with DataWrapper(self.data_wrapper_class, **params) as dw:
                 dw.process()
 
-    def __enter__(self):
-        if getattr(self, "_entered", False):
-            return self
-        self._entered = True
-        return self
-
     def __exit__(self, exc_type, exc_value, traceback):
         try:
             if self.mmanifest and self.mmanifest._new_records:
                 self.mmanifest.save()
-                self.mmanifest.cleanup_temp_manifests()
-            if getattr(self, "_entered", False) and self.fs and self._own_fs:
-                self.fs.close()
         except Exception as e:
             self.logger.warning(f"Error closing filesystem: {e}")
         finally:
-            self._entered = False
+            super().__exit__(exc_type, exc_value, traceback)
         # return False so exceptions aren’t suppressed
         return False
 
-    @classmethod
-    def get_size_estimate(cls, parquet_path: str, **kwargs) -> int:
+    def get_size_estimate(self, **kwargs) -> int:
         """
-        Estimate complexity as total bytes of all .parquet files under parquet_path.
-        Returns size in megabytes (so you can cap or scale priority sensibly).
+        Synchronously estimates artifact size for use in multi-threaded environments.
+
+        This method uses the filesystem's own .sync() method to safely execute
+        asynchronous I/O operations from a synchronous context, preventing
+        event loop conflicts.
         """
-        fs, _, paths = cls.fs.get_fs_token_paths(parquet_path)
-        files = fs.glob(f"{parquet_path}/*.parquet")
-        total_bytes = sum(fs.size(f) for f in files)
-        # convert to “units” (e.g. MB) so priorities stay in a reasonable range
+
+        async def _get_total_bytes_async():
+            """A helper async coroutine to perform the I/O."""
+            import asyncio
+
+            # Use the async versions of fsspec methods (e.g., _glob, _size)
+            files = await self.fs._glob(f"{self.parquet_storage_path}/*.parquet")
+            if not files:
+                return 0
+
+            # Concurrently gather the size of all files for performance
+            size_tasks = [self.fs._size(f) for f in files]
+            sizes = await asyncio.gather(*size_tasks)
+            return sum(s for s in sizes if s is not None)
+
+        # Use the filesystem's own built-in sync method. This is the most
+        # reliable way to bridge the sync/async gap for fsspec.
+        total_bytes = self.fs.sync(_get_total_bytes_async())
+
+        # Convert to megabytes, ensuring a minimum of 1
         return max(1, int(total_bytes / (1024 ** 2)))
 
     def update_parquet(self, period: str = 'today', **kwargs) -> None:
@@ -316,6 +324,8 @@ class ParquetArtifact(DfHelper):
             'parquet_filename': self.parquet_filename,
             'data_path': self.parquet_storage_path,
             'fs': self.fs,
+            'debug': self.debug,
+            'logger': self.logger,
             'class_params': self.class_params,
             'date_field': self.date_field,
             'load_params': self.load_params,

sibi_dst/df_helper/_parquet_reader.py

@@ -54,9 +54,9 @@ class ParquetReader(DfHelper):
             **kwargs,
         }
         self.df: Optional[dd.DataFrame] = None
-        self.debug = self.config.setdefault('debug', False)
-        self.logger = self.config.setdefault('logger', Logger.default_logger(logger_name=self.__class__.__name__))
-        self.logger.set_level(logging.DEBUG if self.debug else logging.INFO)
+        #self.debug = self.config.setdefault('debug', False)
+        #self.logger = self.config.setdefault('logger', Logger.default_logger(logger_name=self.__class__.__name__))
+        #self.logger.set_level(logging.DEBUG if self.debug else logging.INFO)
         self.parquet_storage_path = self.config.setdefault('parquet_storage_path', None)
         if self.parquet_storage_path is None:
             raise ValueError('parquet_storage_path must be set')
@@ -69,12 +69,12 @@ class ParquetReader(DfHelper):
             raise ValueError('parquet_end_date must be set')
 
         # Filesystem setup
-        self.filesystem_type = filesystem_type
-        self.filesystem_options = filesystem_options or {}
-        self.fs = self.config.setdefault('fs', None)
-        if self.fs is None:
-            self.fs = fsspec.filesystem(self.filesystem_type, **self.filesystem_options)
-        self.config.setdefault('fs', self.fs)
+        #self.filesystem_type = filesystem_type
+        #self.filesystem_options = filesystem_options or {}
+        #self.fs = self.config.setdefault('fs', None)
+        #if self.fs is None:
+        #    self.fs = fsspec.filesystem(self.filesystem_type, **self.filesystem_options)
+        #self.config.setdefault('fs', self.fs)
 
         if not self.directory_exists():
             raise ValueError(f"{self.parquet_storage_path} does not exist")
@@ -92,7 +92,3 @@ class ParquetReader(DfHelper):
         except FileNotFoundError:
             return False
 
-    def __exit__(self, exc_type, exc_value, traceback):
-        # Ensure resources are cleaned up
-        if self.fs:
-            self.fs.close()

sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-
+import os
 import threading
 from contextlib import contextmanager
 from typing import Any, Optional, ClassVar, Generator, Type, Dict
@@ -49,10 +49,10 @@ class SqlAlchemyConnectionConfig(BaseModel):
     debug: bool = False
 
     # --- Pool Configuration ---
-    pool_size: int = 5
-    max_overflow: int = 10
-    pool_timeout: int = 30
-    pool_recycle: int = 1800
+    pool_size: int = int(os.environ.get("DB_POOL_SIZE", 5))
+    max_overflow: int = int(os.environ.get("DB_MAX_OVERFLOW",10))
+    pool_timeout: int = int(os.environ.get("DB_POOL_TIMEOUT", 30))
+    pool_recycle: int = int(os.environ.get("DB_POOL_RECYCLE", 1800))
     pool_pre_ping: bool = True
     poolclass: Type[QueuePool] = QueuePool
 
@@ -60,6 +60,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
     model: Optional[Type[Any]] = None
     engine: Optional[Engine] = None
     logger: Optional[Logger] = None
+    _own_logger: bool = False  # Indicates if this instance owns the logger.
     session_factory: Optional[sessionmaker] = None
 
     # --- Private State ---
@@ -102,10 +103,13 @@ class SqlAlchemyConnectionConfig(BaseModel):
 
     def _init_logger(self) -> None:
         """Initializes the logger for this instance."""
+        # This is not a ManagedResource subclass, so we handle logger initialization directly.
+        # unless a logger is provided, we create our own.
         if self.logger is None:
+            self._own_logger = True
             self.logger = Logger.default_logger(logger_name=self.__class__.__name__)
-        log_level = Logger.DEBUG if self.debug else Logger.INFO
-        self.logger.set_level(log_level)
+            log_level = Logger.DEBUG if self.debug else Logger.INFO
+            self.logger.set_level(log_level)
 
     def _get_engine_key(self) -> tuple:
         """Generates a unique, normalized key for an engine configuration."""
@@ -146,6 +150,8 @@ class SqlAlchemyConnectionConfig(BaseModel):
                     self.logger.error(f"Failed to create engine: {e}")
                     raise SQLAlchemyError(f"Engine creation failed: {e}") from e
 
+            #self.logger.debug(f"Connections Active: {self.active_connections}")
+
     def close(self) -> None:
         """
         Decrements the engine's reference count and disposes of the engine
@@ -176,8 +182,6 @@ class SqlAlchemyConnectionConfig(BaseModel):
         # Mark this instance as closed to prevent subsequent calls.
         self._closed = True
 
-    # ... (the rest of your methods like _attach_events, _on_checkout, get_session, etc. remain unchanged)
-    # They are omitted here for brevity but should be included in your final file.
 
     def _attach_events(self) -> None:
         """Attaches checkout/checkin events to the engine for connection tracking."""
@@ -191,7 +195,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
             wrapper = self._engine_registry.get(self._engine_key_instance)
             if wrapper:
                 wrapper['active_connections'] += 1
-        self.logger.debug(f"Connection checked out. Active: {self.active_connections}")
+        # self.logger.debug(f"Connection checked out. Active: {self.active_connections}")
 
     def _on_checkin(self, *args) -> None:
         """Event listener for when a connection is returned to the pool."""
@@ -199,7 +203,7 @@ class SqlAlchemyConnectionConfig(BaseModel):
             wrapper = self._engine_registry.get(self._engine_key_instance)
             if wrapper:
                 wrapper['active_connections'] = max(0, wrapper['active_connections'] - 1)
-        self.logger.debug(f"Connection checked in. Active: {self.active_connections}")
+        # self.logger.debug(f"Connection checked in. Active: {self.active_connections}")
 
     @property
     def active_connections(self) -> int:

sibi_dst/df_helper/backends/sqlalchemy/_io_dask.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
 
-from typing import Type
+from typing import Type, Any
 
 import dask
 import dask.dataframe as dd
@@ -13,11 +14,12 @@ from sqlalchemy.orm import declarative_base
 import time
 from sqlalchemy.exc import TimeoutError as SASQLTimeoutError, OperationalError
 import sqlalchemy as sa
+
+from sibi_dst.utils import ManagedResource
 from sibi_dst.df_helper.core import FilterHandler
-from sibi_dst.utils import Logger
 
 
-class SQLAlchemyDask:
+class SQLAlchemyDask(ManagedResource):
     """
     Loads data from a database into a Dask DataFrame using a memory-safe,
     non-parallel, paginated approach.
@@ -46,8 +48,7 @@ class SQLAlchemyDask:
             filters: dict,
             engine: Engine,
             chunk_size: int = 1000,
-            logger=None,
-            debug: bool = False,
+            **kwargs
     ):
         """
         Initializes the data loader.
@@ -60,14 +61,13 @@ class SQLAlchemyDask:
             logger: A logger instance.
             debug: Whether to enable detailed logging.
         """
+        super().__init__(**kwargs)
         self.model = model
         self.filters = filters
         self.engine = engine
         self.chunk_size = chunk_size
-        self.debug = debug
-        self.logger = logger or Logger.default_logger(logger_name=self.__class__.__name__)
-        self.logger.set_level(Logger.DEBUG if debug else Logger.INFO)
         self.filter_handler_cls = FilterHandler
+        self.total_records = -1 # Initialize to -1 to indicate uncounted
 
     @classmethod
     def infer_meta_from_model(cls, model: Type[declarative_base()]) -> dict:
@@ -83,7 +83,7 @@ class SQLAlchemyDask:
             dtypes[column.name] = dtype
         return dtypes
 
-    def read_frame(self, fillna_value=None) -> dd.DataFrame:
+    def read_frame(self, fillna_value=None) -> tuple[int | Any, Any] | Any:
         """
         Builds and executes a query to load data into a Dask DataFrame.
 
@@ -105,7 +105,8 @@ class SQLAlchemyDask:
             ).apply_filters(query, model=self.model, filters=self.filters)
         else:
             query = query.limit(self.chunk_size)
-        self.logger.debug(f"Base query for pagination: {query}")
+        if self.verbose:
+            self.logger.debug(f"Base query for pagination: {query}")
 
         # 2. Get metadata for the Dask DataFrame structure
         ordered_columns = [column.name for column in self.model.__table__.columns]
@@ -116,6 +117,7 @@ class SQLAlchemyDask:
 
         retry_attempts = 3
         backoff_factor = 0.5  # start with a 0.5-second delay
+        total_records = 0
 
         for attempt in range(retry_attempts):
             try:
@@ -134,11 +136,12 @@ class SQLAlchemyDask:
                     time.sleep(backoff_factor)
                     backoff_factor *= 2  # Double the backoff time for the next attempt
                 else:
+                    self.total_records = -1  # Indicate failure to count records
                     self.logger.error(
                         "Failed to get a connection from the pool after several retries.",
                         exc_info=True
                     )
-                    return dd.from_pandas(meta_df, npartitions=1)
+                    return self.total_records, dd.from_pandas(meta_df, npartitions=1)
             except OperationalError as oe:
                 # sometimes the DB driver wraps timeouts in OperationalError
                 if "timeout" in str(oe).lower():
@@ -147,15 +150,18 @@ class SQLAlchemyDask:
                     backoff_factor *= 2
                     continue
                 else:
+                    self.total_records = -1  # Indicate failure to count records
                     self.logger.error("OperationalError", exc_info=True)
-                    return dd.from_pandas(meta_df, npartitions=1)
+                    return self.total_records, dd.from_pandas(meta_df, npartitions=1)
             except Exception as e:
+                self.total_records = -1  # Indicate failure to count records
                 self.logger.error(f"An unexpected error occurred: {e}", exc_info=True)
-                return dd.from_pandas(meta_df, npartitions=1)
+                return self.total_records, dd.from_pandas(meta_df, npartitions=1)
 
+        self.total_records = total_records
         if total_records == 0:
             self.logger.warning("Query returned 0 records.")
-            return dd.from_pandas(meta_df, npartitions=1)
+            return self.total_records, dd.from_pandas(meta_df, npartitions=1)
 
         self.logger.debug(f"Total records to fetch: {total_records}. Chunk size: {self.chunk_size}.")
 
@@ -179,8 +185,9 @@ class SQLAlchemyDask:
         # 5. Construct the final lazy Dask DataFrame from the delayed chunks
         ddf = dd.from_delayed(delayed_chunks, meta=meta_df)
         self.logger.debug(f"Successfully created a lazy Dask DataFrame with {ddf.npartitions} partitions.")
-
-        return ddf
+        if not self._entered:
+            super().cleanup()
+        return self.total_records, ddf
 
 ## Dask-Only Solution to test in better hardware