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

sibi_dst/df_helper/backends/sqlalchemy/__init__.py

@@ -1,4 +1,3 @@
-from ._filter_handler import SqlAlchemyFilterHandler
 from ._db_connection import SqlAlchemyConnectionConfig
 from ._load_from_db import SqlAlchemyLoadFromDb
 from ._sql_model_builder import SqlAlchemyModelBuilder
@@ -7,5 +6,4 @@ __all__ = [
     'SqlAlchemyConnectionConfig',
     'SqlAlchemyModelBuilder',
     'SqlAlchemyLoadFromDb',
-    'SqlAlchemyFilterHandler'
 ]

sibi_dst/df_helper/backends/sqlalchemy/_db_connection.py

@@ -1,51 +1,85 @@
 from __future__ import annotations
-from typing import Any, Optional, ClassVar, Generator, Type
+
 import threading
 from contextlib import contextmanager
-from pydantic import BaseModel, field_validator, ValidationError, model_validator
+from typing import Any, Optional, ClassVar, Generator, Type, Dict
+
+from pydantic import (
+    BaseModel,
+    field_validator,
+    model_validator,
+    ConfigDict,
+)
 from sqlalchemy import create_engine, event, text
 from sqlalchemy.engine import url as sqlalchemy_url
 from sqlalchemy.engine import Engine
-from sqlalchemy.exc import OperationalError
-from sqlalchemy.pool import QueuePool, NullPool, StaticPool
+from sqlalchemy.exc import OperationalError, SQLAlchemyError
 from sqlalchemy.orm import sessionmaker, Session
+from sqlalchemy.pool import QueuePool, NullPool, StaticPool
+
+# Assuming these are your project's internal modules
 from sibi_dst.utils import Logger
 from ._sql_model_builder import SqlAlchemyModelBuilder
 
 
 class SqlAlchemyConnectionConfig(BaseModel):
     """
-    Thread-safe, registry-backed SQLAlchemy connection manager with:
-      - Shared engine reuse
-      - Active connection tracking
-      - Idle-pool and database-level cleanup
-      - Dynamic ORM model building via SqlAlchemyModelBuilder
-      - Optional session factory
+    A thread-safe, registry-backed SQLAlchemy connection manager.
+
+    This class encapsulates database connection configuration and provides robust,
+    shared resource management. It is designed to be used as a context manager
+    to ensure resources are always released correctly.
+
+    Recommended Usage is via the `with` statement.
+    with SqlAlchemyConnectionConfig(...) as config:
+        session = config.get_session()
+        # ... do work ...
+    # config.close() is called automatically upon exiting the block.
+
+    Key Features:
+      - Context Manager Support: Guarantees resource cleanup.
+      - Shared Engine & Pool: Reuses a single SQLAlchemy Engine for identical
+        database URLs and pool settings, improving application performance.
+      - Reference Counting: Safely manages the lifecycle of the shared engine,
+        disposing of it only when the last user has closed its connection config.
     """
+    # --- Public Configuration ---
     connection_url: str
     table: Optional[str] = None
-    model: Optional[Any] = None
-    engine: Optional[Engine] = None
-    logger: Logger = None
     debug: bool = False
 
+    # --- Pool Configuration ---
     pool_size: int = 5
     max_overflow: int = 10
     pool_timeout: int = 30
-    pool_recycle: int = 300
+    pool_recycle: int = 1800
     pool_pre_ping: bool = True
-    poolclass: Type = QueuePool
+    poolclass: Type[QueuePool] = QueuePool
 
+    # --- Internal & Runtime State ---
+    model: Optional[Type[Any]] = None
+    engine: Optional[Engine] = None
+    logger: Optional[Logger] = None
     session_factory: Optional[sessionmaker] = None
-    _owns_engine: bool = False
 
-    _engine_registry: ClassVar[dict[tuple, Engine]] = {}
+    # --- Private State ---
+    _engine_key_instance: tuple = ()
+    _closed: bool = False  # Flag to prevent double-closing.
+
+    # --- Class-level Shared Resources ---
+    _engine_registry: ClassVar[Dict[tuple, Dict[str, Any]]] = {}
     _registry_lock: ClassVar[threading.Lock] = threading.Lock()
-    _active_connections: ClassVar[int] = 0
 
-    class Config:
-        arbitrary_types_allowed = True
-        underscore_attrs_are_private = True
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    # Add __enter__ and __exit__ for context manager protocol
+    def __enter__(self) -> SqlAlchemyConnectionConfig:
+        """Enter the runtime context, returning self."""
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """Exit the runtime context, ensuring that close() is called."""
+        self.close()
 
     @field_validator("pool_size", "max_overflow", "pool_timeout", "pool_recycle")
     @classmethod
@@ -56,78 +90,139 @@ class SqlAlchemyConnectionConfig(BaseModel):
 
     @model_validator(mode="after")
     def _init_all(self) -> SqlAlchemyConnectionConfig:
+        """Orchestrates the initialization process after Pydantic validation."""
         self._init_logger()
+        self._engine_key_instance = self._get_engine_key()
         self._init_engine()
         self._validate_conn()
         self._build_model()
-        self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False)
+        if self.engine:
+            self.session_factory = sessionmaker(bind=self.engine, expire_on_commit=False)
         return self
 
     def _init_logger(self) -> None:
-        self.logger = self.logger or Logger.default_logger(logger_name=self.__class__.__name__)
-        self.logger.set_level(Logger.DEBUG if self.debug else Logger.INFO)
-
-    def _engine_key(self) -> tuple:
+        """Initializes the logger for this instance."""
+        if self.logger is None:
+            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)
+
+    def _get_engine_key(self) -> tuple:
+        """Generates a unique, normalized key for an engine configuration."""
         parsed = sqlalchemy_url.make_url(self.connection_url)
         query = {k: v for k, v in parsed.query.items() if not k.startswith("pool_")}
-        normalized = parsed.set(query=query)
-        key = [str(normalized)]
+        normalized_url = parsed.set(query=query)
+        key_parts = [str(normalized_url)]
         if self.poolclass not in (NullPool, StaticPool):
-            key += [self.pool_size, self.max_overflow, self.pool_timeout, self.pool_recycle, self.pool_pre_ping, self.table]
-        return tuple(key)
+            key_parts += [
+                self.pool_size, self.max_overflow, self.pool_timeout,
+                self.pool_recycle, self.pool_pre_ping
+            ]
+        return tuple(key_parts)
 
     def _init_engine(self) -> None:
-        key = self._engine_key()
+        """Initializes or reuses a shared SQLAlchemy Engine."""
         with self._registry_lock:
-            existing = self._engine_registry.get(key)
-            if existing:
-                self.engine = existing
-                self._owns_engine = False
-                self.logger.debug(f"Reusing engine {key}")
+            engine_wrapper = self._engine_registry.get(self._engine_key_instance)
+            if engine_wrapper:
+                self.engine = engine_wrapper['engine']
+                engine_wrapper['ref_count'] += 1
+                self.logger.debug(f"Reusing engine. Ref count: {engine_wrapper['ref_count']}.")
             else:
-                self.logger.debug(f"Creating engine {key}")
-                self.engine = create_engine(
-                    self.connection_url,
-                    pool_size=self.pool_size,
-                    max_overflow=self.max_overflow,
-                    pool_timeout=self.pool_timeout,
-                    pool_recycle=self.pool_recycle,
-                    pool_pre_ping=self.pool_pre_ping,
-                    poolclass=self.poolclass,
-                )
-                self._attach_events()
-                self._engine_registry[key] = self.engine
-                self._owns_engine = True
+                self.logger.debug(f"Creating new engine for key: {self._engine_key_instance}")
+                try:
+                    new_engine = create_engine(
+                        self.connection_url, pool_size=self.pool_size,
+                        max_overflow=self.max_overflow, pool_timeout=self.pool_timeout,
+                        pool_recycle=self.pool_recycle, pool_pre_ping=self.pool_pre_ping,
+                        poolclass=self.poolclass,
+                    )
+                    self.engine = new_engine
+                    self._attach_events()
+                    self._engine_registry[self._engine_key_instance] = {
+                        'engine': new_engine, 'ref_count': 1, 'active_connections': 0
+                    }
+                except Exception as e:
+                    self.logger.error(f"Failed to create engine: {e}")
+                    raise SQLAlchemyError(f"Engine creation failed: {e}") from e
+
+    def close(self) -> None:
+        """
+        Decrements the engine's reference count and disposes of the engine
+        if the count reaches zero. This is now typically called automatically
+        when exiting a `with` block.
+        """
+        # Prevent the method from running more than once per instance.
+        if self._closed:
+            self.logger.debug("Attempted to close an already-closed config instance.")
+            return
+
+        with self._registry_lock:
+            key = self._engine_key_instance
+            engine_wrapper = self._engine_registry.get(key)
+
+            if not engine_wrapper:
+                self.logger.warning("Attempted to close a config whose engine is not in the registry.")
+                return
+
+            engine_wrapper['ref_count'] -= 1
+            self.logger.debug(f"Closing config. Ref count is now {engine_wrapper['ref_count']}.")
+
+            if engine_wrapper['ref_count'] <= 0:
+                self.logger.debug(f"Disposing engine as reference count is zero. Key: {key}")
+                engine_wrapper['engine'].dispose()
+                del self._engine_registry[key]
+
+        # 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:
-        event.listen(self.engine, "checkout", self._on_checkout)
-        event.listen(self.engine, "checkin", self._on_checkin)
+        """Attaches checkout/checkin events to the engine for connection tracking."""
+        if self.engine:
+            event.listen(self.engine, "checkout", self._on_checkout)
+            event.listen(self.engine, "checkin", self._on_checkin)
 
     def _on_checkout(self, *args) -> None:
+        """Event listener for when a connection is checked out from the pool."""
         with self._registry_lock:
-            type(self)._active_connections += 1
-        self.logger.debug(f"Checked out, active: {self.active_connections}")
+            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}")
 
     def _on_checkin(self, *args) -> None:
+        """Event listener for when a connection is returned to the pool."""
         with self._registry_lock:
-            type(self)._active_connections = max(type(self)._active_connections - 1, 0)
-        self.logger.debug(f"Checked in, active: {self.active_connections}")
+            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}")
 
     @property
     def active_connections(self) -> int:
-        return type(self)._active_connections
+        """Returns the number of active connections for this instance's engine."""
+        with self._registry_lock:
+            wrapper = self._engine_registry.get(self._engine_key_instance)
+            return wrapper['active_connections'] if wrapper else 0
 
     def _validate_conn(self) -> None:
+        """Tests the database connection by executing a simple query."""
         try:
             with self.managed_connection() as conn:
                 conn.execute(text("SELECT 1"))
-            self.logger.debug("Connection OK")
+            self.logger.debug("Database connection validated successfully.")
         except OperationalError as e:
-            self.logger.error(f"Connection failed: {e}")
-            raise ValidationError(f"DB connection failed: {e}")
+            self.logger.error(f"Database connection failed: {e}")
+            raise ValueError(f"DB connection failed: {e}") from e
 
     @contextmanager
-    def managed_connection(self) -> Generator[Any, None, Any]:
+    def managed_connection(self) -> Generator[Any, None, None]:
+        """Provides a single database connection from the engine pool."""
+        if not self.engine:
+            raise RuntimeError("Engine not initialized. Cannot get a connection.")
         conn = self.engine.connect()
         try:
             yield conn
@@ -135,68 +230,19 @@ class SqlAlchemyConnectionConfig(BaseModel):
             conn.close()
 
     def get_session(self) -> Session:
+        """Returns a new SQLAlchemy Session from the session factory."""
         if not self.session_factory:
-            raise RuntimeError("Session factory not initialized")
+            raise RuntimeError("Session factory not initialized. Cannot get a session.")
         return self.session_factory()
 
     def _build_model(self) -> None:
-        """Dynamically build and assign the ORM model if table is set"""
+        """Dynamically builds an ORM model if `self.table` is set."""
         if not self.table or not self.engine:
             return
         try:
             builder = SqlAlchemyModelBuilder(self.engine, self.table)
             self.model = builder.build_model()
-            self.logger.debug(f"Model built for table: {self.table}")
+            self.logger.debug(f"Successfully built ORM model for table: {self.table}")
         except Exception as e:
-            self.logger.error(f"Model build failed: {e}")
-            raise ValidationError(f"Model construction error: {e}") from e
-
-    def dispose_idle_connections(self) -> int:
-        key = self._engine_key()
-        with self._registry_lock:
-            if self._engine_registry.get(key) is not self.engine:
-                self.logger.debug("Engine changed")
-                return 0
-            pool = self.engine.pool
-            if isinstance(pool, QueuePool):
-                count = pool.checkedin()
-                pool.dispose()
-                self.logger.debug(f"Disposed {count}")
-                return count
-            self.logger.warning(f"No idle dispose for {type(pool).__name__}")
-            return 0
-
-    def terminate_idle_connections(self, idle_seconds: int = 300) -> int:
-        terminated = 0
-        dialect = self.engine.dialect.name
-        with self.managed_connection() as conn:
-            if dialect == 'postgresql':
-                res = conn.execute(text(
-                    f"SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
-                    f"WHERE state='idle' AND (now() - query_start) > interval '{idle_seconds} seconds' "
-                    f"AND pid<>pg_backend_pid()"
-                ))
-                terminated = res.rowcount
-            elif dialect == 'mysql':
-                for row in conn.execute(text("SHOW PROCESSLIST")):
-                    if row.Command == 'Sleep' and row.Time > idle_seconds:
-                        conn.execute(text(f"KILL {row.Id}"))
-                        terminated += 1
-            else:
-                self.logger.warning(f"Idle termination not supported: {dialect}")
-        self.logger.debug(f"Terminated {terminated}")
-        return terminated
-
-    def close(self) -> None:
-        with self._registry_lock:
-            key = self._engine_key()
-            if not self._owns_engine:
-                self.logger.debug("Not owner, skipping close")
-                return
-            if self._engine_registry.get(key) != self.engine:
-                self.logger.debug("Engine not in registry")
-                return
-            self.engine.dispose()
-            del self._engine_registry[key]
-            type(self)._active_connections = 0
-            self.logger.debug(f"Engine closed {key}")
+            self.logger.error(f"Failed to build ORM model for table '{self.table}': {e}")
+            raise ValueError(f"Model construction failed for table '{self.table}': {e}") from e

sibi_dst/df_helper/backends/sqlalchemy/_io_dask.py

@@ -1,135 +1,179 @@
-import itertools
+from typing import Type
 
+import dask
 import dask.dataframe as dd
 import pandas as pd
-from sqlalchemy import create_engine, inspect, select
-from sqlalchemy.orm import sessionmaker
-
+from sqlalchemy import (
+    inspect,
+    select,
+    func,
+)
+from sqlalchemy.engine import Engine
+from sqlalchemy.orm import declarative_base
+import time
+from sqlalchemy.exc import TimeoutError
+import sqlalchemy as sa
 from sibi_dst.df_helper.core import FilterHandler
 from sibi_dst.utils import Logger
 
 
 class SQLAlchemyDask:
-    def __init__(self, model, filters, engine_url, chunk_size=1000, logger=None, debug=False):
+    """
+    Loads data from a database into a Dask DataFrame using a memory-safe,
+    non-parallel, paginated approach.
+
+    This class avoids using a numeric `index_col for parallel loading.
+    """
+
+    _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",
+    }
+
+    def __init__(
+            self,
+            model: Type[declarative_base()],
+            filters: dict,
+            engine: Engine,
+            chunk_size: int = 1000,
+            logger=None,
+            debug: bool = 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.
+        Initializes the data loader.
+
+        Args:
+            model: The SQLAlchemy ORM model for the table.
+            filters: A dictionary of filters to apply to the query.
+            engine: An SQLAlchemy Engine instance.
+            chunk_size: The number of records to fetch in each database query.
+            logger: A logger instance.
+            debug: Whether to enable detailed logging.
         """
-        self.query = None
         self.model = model
         self.filters = filters
+        self.engine = engine
         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)
+        self.logger.set_level(Logger.DEBUG if debug else Logger.INFO)
+        self.filter_handler_cls = FilterHandler
 
-    @staticmethod
-    def infer_dtypes_from_model(model):
+    @classmethod
+    def infer_meta_from_model(cls, model: Type[declarative_base()]) -> dict:
         """
-        Infer data types for Dask DataFrame based on SQLAlchemy ORM model columns.
+        Infers a metadata dictionary for Dask based on the SQLAlchemy model.
+        This helps Dask understand the DataFrame structure without reading data.
         """
         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')
+            dtype_str = str(column.type).upper().split("(")[0]
+            dtype = cls._SQLALCHEMY_TO_DASK_DTYPE.get(dtype_str, "object")
             dtypes[column.name] = dtype
-
         return dtypes
 
-    def read_frame(self, fillna_value=None):
+    def read_frame(self, fillna_value=None) -> dd.DataFrame:
         """
-        Load data from an SQLAlchemy query into a Dask DataFrame.
+        Builds and executes a query to load data into a Dask DataFrame.
 
-        :param fillna_value: Value to replace NaN or NULL values with, if any.
-        :return: Dask DataFrame.
+        This method works by first running a COUNT query to get the total
+        size, then creating a series of delayed tasks that each fetch a
+        chunk of data using LIMIT/OFFSET.
+
+        Args:
+            fillna_value: Value to replace NaN or NULL values with, if any.
+
+        Returns:
+            A lazy Dask DataFrame.
         """
-        with self.Session() as session:
+        # 1. Build the base query and apply filters
+        query = select(self.model)
+        if self.filters:
+            query = self.filter_handler_cls(
+                backend="sqlalchemy", logger=self.logger, debug=self.debug
+            ).apply_filters(query, model=self.model, filters=self.filters)
+
+        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]
+        meta_dtypes = self.infer_meta_from_model(self.model)
+        meta_df = pd.DataFrame(columns=ordered_columns).astype(meta_dtypes)
+
+        # 3. Get the total record count to calculate the number of chunks
+        # try:
+        #     with self.engine.connect() as connection:
+        #         count_query = select(func.count()).select_from(query.alias())
+        #         total_records = connection.execute(count_query).scalar_one()
+        # except Exception as e:
+        #     self.logger.error(f"Failed to count records for pagination: {e}", exc_info=True)
+        #     return dd.from_pandas(meta_df, npartitions=1)
+        retry_attempts = 3
+        backoff_factor = 0.5  # start with a 0.5-second delay
+
+        for attempt in range(retry_attempts):
             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)
+                with self.engine.connect() as connection:
+                    count_query = sa.select(sa.func.count()).select_from(query.alias())
+                    total_records = connection.execute(count_query).scalar_one()
+
+                # If successful, break the loop
+                break
+
+            except TimeoutError:
+                if attempt < retry_attempts - 1:
+                    self.logger.warning(
+                        f"Connection pool limit reached. Retrying in {backoff_factor} seconds..."
+                    )
+                    time.sleep(backoff_factor)
+                    backoff_factor *= 2  # Double the backoff time for the next attempt
                 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]
+                    self.logger.error(
+                        "Failed to get a connection from the pool after several retries.",
+                        exc_info=True
                     )
-                    # Drop internal SQLAlchemy state if it exists
-                    df = df.loc[:, ~df.columns.str.contains('_sa_instance_state')]
+                    return dd.from_pandas(meta_df, npartitions=1)
+
+            except Exception as e:
+                self.logger.error(f"An unexpected error occurred: {e}", exc_info=True)
+                return dd.from_pandas(meta_df, npartitions=1)
 
-                    # Reorder columns to match the model's order
-                    df = df[ordered_columns]
+        if total_records == 0:
+            self.logger.warning("Query returned 0 records.")
+            return dd.from_pandas(meta_df, npartitions=1)
 
-                    # Fill NaN values
-                    if fillna_value is not None:
-                        df = df.fillna(fillna_value)
+        self.logger.debug(f"Total records to fetch: {total_records}. Chunk size: {self.chunk_size}.")
 
-                    # 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)
+        # 4. Create a list of Dask Delayed objects, one for each chunk
+        @dask.delayed
+        def get_chunk(sql_query, chunk_offset):
+            """A Dask-delayed function to fetch one chunk of data."""
+            # LIMIT/OFFSET must be applied in the delayed function
+            paginated_query = sql_query.limit(self.chunk_size).offset(chunk_offset)
+            df = pd.read_sql(paginated_query, self.engine)
 
-                    # Apply inferred dtypes
-                    df = df.astype(dtypes)
-                    # Create a Dask partition
-                    partitions.append(dd.from_pandas(df, npartitions=1))
+            if fillna_value is not None:
+                df = df.fillna(fillna_value)
 
-                # 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)
+            # Ensure column order and types match the meta
+            return df[ordered_columns].astype(meta_dtypes)
 
-                self.logger.debug(f"Loaded {len(dask_df)} rows into Dask DataFrame.")
+        offsets = range(0, total_records, self.chunk_size)
+        delayed_chunks = [get_chunk(query, offset) for offset in offsets]
 
-                return dask_df
+        # 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.")
 
-            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)
+        return ddf