meerschaum.connectors.sql

106    def __init__(
107        self,
108        label: Optional[str] = None,
109        flavor: Optional[str] = None,
110        wait: bool = False,
111        connect: bool = False,
112        debug: bool = False,
113        **kw: Any
114    ):
115        """
116        Parameters
117        ----------
118        label: str, default 'main'
119            The identifying label for the connector.
120            E.g. for `sql:main`, 'main' is the label.
121            Defaults to 'main'.
122
123        flavor: Optional[str], default None
124            The database flavor, e.g.
125            `'sqlite'`, `'postgresql'`, `'cockroachdb'`, etc.
126            To see supported flavors, run the `bootstrap connectors` command.
127
128        wait: bool, default False
129            If `True`, block until a database connection has been made.
130            Defaults to `False`.
131
132        connect: bool, default False
133            If `True`, immediately attempt to connect the database and raise
134            a warning if the connection fails.
135            Defaults to `False`.
136
137        debug: bool, default False
138            Verbosity toggle.
139            Defaults to `False`.
140
141        kw: Any
142            All other arguments will be passed to the connector's attributes.
143            Therefore, a connector may be made without being registered,
144            as long enough parameters are supplied to the constructor.
145        """
146        if 'uri' in kw:
147            uri = kw['uri']
148            if uri.startswith('postgres') and not uri.startswith('postgresql'):
149                uri = uri.replace('postgres', 'postgresql', 1)
150            if uri.startswith('postgresql') and not uri.startswith('postgresql+'):
151                uri = uri.replace('postgresql://', 'postgresql+psycopg://', 1)
152            if uri.startswith('timescaledb://'):
153                uri = uri.replace('timescaledb://', 'postgresql+psycopg://', 1)
154                flavor = 'timescaledb'
155            if uri.startswith('postgis://'):
156                uri = uri.replace('postgis://', 'postgresql+psycopg://', 1)
157                flavor = 'postgis'
158            kw['uri'] = uri
159            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
160            label = label or from_uri_params.get('label', None)
161            _ = from_uri_params.pop('label', None)
162
163            ### Sometimes the flavor may be provided with a URI.
164            kw.update(from_uri_params)
165            if flavor:
166                kw['flavor'] = flavor
167
168        ### set __dict__ in base class
169        super().__init__(
170            'sql',
171            label = label or self.__dict__.get('label', None),
172            **kw
173        )
174
175        if self.__dict__.get('flavor', None) == 'sqlite':
176            self._reset_attributes()
177            self._set_attributes(
178                'sql',
179                label = label,
180                inherit_default = False,
181                **kw
182            )
183            ### For backwards compatability reasons, set the path for sql:local if its missing.
184            if self.label == 'local' and not self.__dict__.get('database', None):
185                from meerschaum.config._paths import SQLITE_DB_PATH
186                self.database = str(SQLITE_DB_PATH)
187
188        ### ensure flavor and label are set accordingly
189        if 'flavor' not in self.__dict__:
190            if flavor is None and 'uri' not in self.__dict__:
191                raise ValueError(
192                    f"    Missing flavor. Provide flavor as a key for '{self}'."
193                )
194            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
195
196        if self.flavor == 'postgres':
197            self.flavor = 'postgresql'
198
199        self._debug = debug
200        ### Store the PID and thread at initialization
201        ### so we can dispose of the Pool in child processes or threads.
202        import os
203        import threading
204        self._pid = os.getpid()
205        self._thread_ident = threading.current_thread().ident
206        self._sessions = {}
207        self._locks = {'_sessions': threading.RLock(), }
208
209        ### verify the flavor's requirements are met
210        if self.flavor not in self.flavor_configs:
211            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
212        if not self.__dict__.get('uri'):
213            self.verify_attributes(
214                self.flavor_configs[self.flavor].get('requirements', set()),
215                debug=debug,
216            )
217
218        if wait:
219            from meerschaum.connectors.poll import retry_connect
220            retry_connect(connector=self, debug=debug)
221
222        if connect:
223            if not self.test_connection(debug=debug):
224                warn(f"Failed to connect with connector '{self}'!", stack=False)

226    @property
227    def Session(self):
228        if '_Session' not in self.__dict__:
229            if self.engine is None:
230                return None
231
232            from meerschaum.utils.packages import attempt_import
233            sqlalchemy_orm = attempt_import('sqlalchemy.orm', lazy=False)
234            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
235            self._Session = sqlalchemy_orm.scoped_session(session_factory)
236
237        return self._Session

239    @property
240    def engine(self):
241        """
242        Return the SQLAlchemy engine connected to the configured database.
243        """
244        import os
245        import threading
246        if '_engine' not in self.__dict__:
247            self._engine, self._engine_str = self.create_engine(include_uri=True)
248
249        same_process = os.getpid() == self._pid
250        same_thread = threading.current_thread().ident == self._thread_ident
251
252        ### handle child processes
253        if not same_process:
254            self._pid = os.getpid()
255            self._thread = threading.current_thread()
256            warn("Different PID detected. Disposing of connections...")
257            self._engine.dispose()
258
259        ### handle different threads
260        if not same_thread:
261            if self.flavor == 'duckdb':
262                warn("Different thread detected.")
263                self._engine.dispose()
264
265        return self._engine

267    @property
268    def DATABASE_URL(self) -> str:
269        """
270        Return the URI connection string (alias for `SQLConnector.URI`.
271        """
272        _ = self.engine
273        return str(self._engine_str)

275    @property
276    def URI(self) -> str:
277        """
278        Return the URI connection string.
279        """
280        _ = self.engine
281        return str(self._engine_str)

283    @property
284    def IS_THREAD_SAFE(self) -> str:
285        """
286        Return whether this connector may be multithreaded.
287        """
288        if self.flavor in ('duckdb', 'oracle'):
289            return False
290        if self.flavor == 'sqlite':
291            return ':memory:' not in self.URI
292        return True

294    @property
295    def metadata(self):
296        """
297        Return the metadata bound to this configured schema.
298        """
299        from meerschaum.utils.packages import attempt_import
300        sqlalchemy = attempt_import('sqlalchemy', lazy=False)
301        if '_metadata' not in self.__dict__:
302            self._metadata = sqlalchemy.MetaData(schema=self.schema)
303        return self._metadata

305    @property
306    def instance_schema(self):
307        """
308        Return the schema name for Meerschaum tables. 
309        """
310        return self.schema

312    @property
313    def internal_schema(self):
314        """
315        Return the schema name for internal tables. 
316        """
317        from meerschaum.config.static import STATIC_CONFIG
318        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
319        schema_name = self.__dict__.get('internal_schema', None) or (
320            STATIC_CONFIG['sql']['internal_schema']
321            if self.flavor not in NO_SCHEMA_FLAVORS
322            else self.schema
323        )
324
325        if '_internal_schema' not in self.__dict__:
326            self._internal_schema = schema_name
327        return self._internal_schema

329    @property
330    def db(self) -> Optional[databases.Database]:
331        from meerschaum.utils.packages import attempt_import
332        databases = attempt_import('databases', lazy=False, install=True)
333        url = self.DATABASE_URL
334        if 'mysql' in url:
335            url = url.replace('+pymysql', '')
336        if '_db' not in self.__dict__:
337            try:
338                self._db = databases.Database(url)
339            except KeyError:
340                ### Likely encountered an unsupported flavor.
341                from meerschaum.utils.warnings import warn
342                self._db = None
343        return self._db

345    @property
346    def db_version(self) -> Union[str, None]:
347        """
348        Return the database version.
349        """
350        _db_version = self.__dict__.get('_db_version', None)
351        if _db_version is not None:
352            return _db_version
353
354        from meerschaum.utils.sql import get_db_version
355        self._db_version = get_db_version(self)
356        return self._db_version

358    @property
359    def schema(self) -> Union[str, None]:
360        """
361        Return the default schema to use.
362        A value of `None` will not prepend a schema.
363        """
364        if 'schema' in self.__dict__:
365            return self.__dict__['schema']
366
367        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
368        if self.flavor in NO_SCHEMA_FLAVORS:
369            self.__dict__['schema'] = None
370            return None
371
372        sqlalchemy = mrsm.attempt_import('sqlalchemy', lazy=False)
373        _schema = sqlalchemy.inspect(self.engine).default_schema_name
374        self.__dict__['schema'] = _schema
375        return _schema

193def create_engine(
194    self,
195    include_uri: bool = False,
196    debug: bool = False,
197    **kw
198) -> 'sqlalchemy.engine.Engine':
199    """Create a sqlalchemy engine by building the engine string."""
200    from meerschaum.utils.packages import attempt_import
201    from meerschaum.utils.warnings import error, warn
202    sqlalchemy = attempt_import('sqlalchemy', lazy=False)
203    import urllib
204    import copy
205    ### Install and patch required drivers.
206    if self.flavor in install_flavor_drivers:
207        _ = attempt_import(
208            *install_flavor_drivers[self.flavor],
209            debug=debug,
210            lazy=False,
211            warn=False,
212        )
213        if self.flavor == 'mssql':
214            _init_mssql_sqlalchemy()
215    if self.flavor in require_patching_flavors:
216        from meerschaum.utils.packages import determine_version, _monkey_patch_get_distribution
217        import pathlib
218        for install_name, import_name in require_patching_flavors[self.flavor]:
219            pkg = attempt_import(
220                import_name,
221                debug=debug,
222                lazy=False,
223                warn=False
224            )
225            _monkey_patch_get_distribution(
226                install_name, determine_version(pathlib.Path(pkg.__file__), venv='mrsm')
227            )
228
229    ### supplement missing values with defaults (e.g. port number)
230    for a, value in flavor_configs[self.flavor]['defaults'].items():
231        if a not in self.__dict__:
232            self.__dict__[a] = value
233
234    ### Verify that everything is in order.
235    if self.flavor not in flavor_configs:
236        error(f"Cannot create a connector with the flavor '{self.flavor}'.")
237
238    _engine = flavor_configs[self.flavor].get('engine', None)
239    _username = self.__dict__.get('username', None)
240    _password = self.__dict__.get('password', None)
241    _host = self.__dict__.get('host', None)
242    _port = self.__dict__.get('port', None)
243    _database = self.__dict__.get('database', None)
244    _options = self.__dict__.get('options', {})
245    if isinstance(_options, str):
246        _options = dict(urllib.parse.parse_qsl(_options))
247    _uri = self.__dict__.get('uri', None)
248
249    ### Handle registering specific dialects (due to installing in virtual environments).
250    if self.flavor in flavor_dialects:
251        sqlalchemy.dialects.registry.register(*flavor_dialects[self.flavor])
252
253    ### self._sys_config was deepcopied and can be updated safely
254    if self.flavor in ("sqlite", "duckdb"):
255        engine_str = f"{_engine}:///{_database}" if not _uri else _uri
256        if 'create_engine' not in self._sys_config:
257            self._sys_config['create_engine'] = {}
258        if 'connect_args' not in self._sys_config['create_engine']:
259            self._sys_config['create_engine']['connect_args'] = {}
260        self._sys_config['create_engine']['connect_args'].update({"check_same_thread": False})
261    else:
262        engine_str = (
263            _engine + "://" + (_username if _username is not None else '') +
264            ((":" + urllib.parse.quote_plus(_password)) if _password is not None else '') +
265            "@" + _host + ((":" + str(_port)) if _port is not None else '') +
266            (("/" + _database) if _database is not None else '')
267            + (("?" + urllib.parse.urlencode(_options)) if _options else '')
268        ) if not _uri else _uri
269
270        ### Sometimes the timescaledb:// flavor can slip in.
271        if _uri and self.flavor in _uri:
272            if self.flavor in ('timescaledb', 'postgis'):
273                engine_str = engine_str.replace(self.flavor, 'postgresql', 1)
274            elif _uri.startswith('postgresql://'):
275                engine_str = engine_str.replace('postgresql://', 'postgresql+psycopg2://')
276
277    if debug:
278        dprint(
279            (
280                (engine_str.replace(':' + _password, ':' + ('*' * len(_password))))
281                    if _password is not None else engine_str
282            ) + '\n' + f"{self._sys_config.get('create_engine', {}).get('connect_args', {})}"
283        )
284
285    _kw_copy = copy.deepcopy(kw)
286
287    ### NOTE: Order of inheritance:
288    ###       1. Defaults
289    ###       2. System configuration
290    ###       3. Connector configuration
291    ###       4. Keyword arguments
292    _create_engine_args = flavor_configs.get(self.flavor, {}).get('create_engine', {})
293    def _apply_create_engine_args(update):
294        if 'ALL' not in flavor_configs[self.flavor].get('omit_create_engine', {}):
295            _create_engine_args.update(
296                { k: v for k, v in update.items()
297                    if 'omit_create_engine' not in flavor_configs[self.flavor]
298                        or k not in flavor_configs[self.flavor].get('omit_create_engine')
299                }
300            )
301    _apply_create_engine_args(self._sys_config.get('create_engine', {}))
302    _apply_create_engine_args(self.__dict__.get('create_engine', {}))
303    _apply_create_engine_args(_kw_copy)
304
305    try:
306        engine = sqlalchemy.create_engine(
307            engine_str,
308            ### I know this looks confusing, and maybe it's bad code,
309            ### but it's simple. It dynamically parses the config string
310            ### and splits it to separate the class name (QueuePool)
311            ### from the module name (sqlalchemy.pool).
312            poolclass    = getattr(
313                attempt_import(
314                    ".".join(self._sys_config['poolclass'].split('.')[:-1])
315                ),
316                self._sys_config['poolclass'].split('.')[-1]
317            ),
318            echo         = debug,
319            **_create_engine_args
320        )
321    except Exception:
322        warn(f"Failed to create connector '{self}':\n{traceback.format_exc()}", stack=False)
323        engine = None
324
325    if include_uri:
326        return engine, engine_str
327    return engine

 28def read(
 29    self,
 30    query_or_table: Union[str, sqlalchemy.Query],
 31    params: Union[Dict[str, Any], List[str], None] = None,
 32    dtype: Optional[Dict[str, Any]] = None,
 33    coerce_float: bool = True,
 34    chunksize: Optional[int] = -1,
 35    workers: Optional[int] = None,
 36    chunk_hook: Optional[Callable[[pandas.DataFrame], Any]] = None,
 37    as_hook_results: bool = False,
 38    chunks: Optional[int] = None,
 39    schema: Optional[str] = None,
 40    as_chunks: bool = False,
 41    as_iterator: bool = False,
 42    as_dask: bool = False,
 43    index_col: Optional[str] = None,
 44    silent: bool = False,
 45    debug: bool = False,
 46    **kw: Any
 47) -> Union[
 48    pandas.DataFrame,
 49    dask.DataFrame,
 50    List[pandas.DataFrame],
 51    List[Any],
 52    None,
 53]:
 54    """
 55    Read a SQL query or table into a pandas dataframe.
 56
 57    Parameters
 58    ----------
 59    query_or_table: Union[str, sqlalchemy.Query]
 60        The SQL query (sqlalchemy Query or string) or name of the table from which to select.
 61
 62    params: Optional[Dict[str, Any]], default None
 63        `List` or `Dict` of parameters to pass to `pandas.read_sql()`.
 64        See the pandas documentation for more information:
 65        https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_sql.html
 66
 67    dtype: Optional[Dict[str, Any]], default None
 68        A dictionary of data types to pass to `pandas.read_sql()`.
 69        See the pandas documentation for more information:
 70        https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.read_sql_query.html
 71
 72    chunksize: Optional[int], default -1
 73        How many chunks to read at a time. `None` will read everything in one large chunk.
 74        Defaults to system configuration.
 75
 76        **NOTE:** DuckDB does not allow for chunking.
 77
 78    workers: Optional[int], default None
 79        How many threads to use when consuming the generator.
 80        Only applies if `chunk_hook` is provided.
 81
 82    chunk_hook: Optional[Callable[[pandas.DataFrame], Any]], default None
 83        Hook function to execute once per chunk, e.g. writing and reading chunks intermittently.
 84        See `--sync-chunks` for an example.
 85        **NOTE:** `as_iterator` MUST be False (default).
 86
 87    as_hook_results: bool, default False
 88        If `True`, return a `List` of the outputs of the hook function.
 89        Only applicable if `chunk_hook` is not None.
 90
 91        **NOTE:** `as_iterator` MUST be `False` (default).
 92
 93    chunks: Optional[int], default None
 94        Limit the number of chunks to read into memory, i.e. how many chunks to retrieve and
 95        return into a single dataframe.
 96        For example, to limit the returned dataframe to 100,000 rows,
 97        you could specify a `chunksize` of `1000` and `chunks` of `100`.
 98
 99    schema: Optional[str], default None
100        If just a table name is provided, optionally specify the table schema.
101        Defaults to `SQLConnector.schema`.
102
103    as_chunks: bool, default False
104        If `True`, return a list of DataFrames.
105        Otherwise return a single DataFrame.
106
107    as_iterator: bool, default False
108        If `True`, return the pandas DataFrame iterator.
109        `chunksize` must not be `None` (falls back to 1000 if so),
110        and hooks are not called in this case.
111
112    index_col: Optional[str], default None
113        If using Dask, use this column as the index column.
114        If omitted, a Pandas DataFrame will be fetched and converted to a Dask DataFrame.
115
116    silent: bool, default False
117        If `True`, don't raise warnings in case of errors.
118        Defaults to `False`.
119
120    Returns
121    -------
122    A `pd.DataFrame` (default case), or an iterator, or a list of dataframes / iterators,
123    or `None` if something breaks.
124
125    """
126    if chunks is not None and chunks <= 0:
127        return []
128    from meerschaum.utils.sql import sql_item_name, truncate_item_name
129    from meerschaum.utils.dtypes import are_dtypes_equal, coerce_timezone
130    from meerschaum.utils.dtypes.sql import TIMEZONE_NAIVE_FLAVORS
131    from meerschaum.utils.packages import attempt_import, import_pandas
132    from meerschaum.utils.pool import get_pool
133    from meerschaum.utils.dataframe import chunksize_to_npartitions, get_numeric_cols
134    import warnings
135    import traceback
136    from decimal import Decimal
137    pd = import_pandas()
138    dd = None
139    is_dask = 'dask' in pd.__name__
140    pandas = attempt_import('pandas')
141    is_dask = dd is not None
142    npartitions = chunksize_to_npartitions(chunksize)
143    if is_dask:
144        chunksize = None
145    schema = schema or self.schema
146    utc_dt_cols = [
147        col
148        for col, typ in dtype.items()
149        if are_dtypes_equal(typ, 'datetime') and 'utc' in typ.lower()
150    ] if dtype else []
151
152    if dtype and utc_dt_cols and self.flavor in TIMEZONE_NAIVE_FLAVORS:
153        dtype = dtype.copy()
154        for col in utc_dt_cols:
155            dtype[col] = 'datetime64[ns]'
156
157    pool = get_pool(workers=workers)
158    sqlalchemy = attempt_import("sqlalchemy", lazy=False)
159    default_chunksize = self._sys_config.get('chunksize', None)
160    chunksize = chunksize if chunksize != -1 else default_chunksize
161    if chunksize is None and as_iterator:
162        if not silent and self.flavor not in _disallow_chunks_flavors:
163            warn(
164                "An iterator may only be generated if chunksize is not None.\n"
165                + "Falling back to a chunksize of 1000.", stacklevel=3,
166            )
167        chunksize = 1000
168    if chunksize is not None and self.flavor in _max_chunks_flavors:
169        if chunksize > _max_chunks_flavors[self.flavor]:
170            if chunksize != default_chunksize:
171                warn(
172                    f"The specified chunksize of {chunksize} exceeds the maximum of "
173                    + f"{_max_chunks_flavors[self.flavor]} for flavor '{self.flavor}'.\n"
174                    + f"    Falling back to a chunksize of {_max_chunks_flavors[self.flavor]}.",
175                    stacklevel=3,
176                )
177            chunksize = _max_chunks_flavors[self.flavor]
178
179    if chunksize is not None and self.flavor in _disallow_chunks_flavors:
180        chunksize = None
181
182    if debug:
183        import time
184        start = time.perf_counter()
185        dprint(f"[{self}]\n{query_or_table}")
186        dprint(f"[{self}] Fetching with chunksize: {chunksize}")
187
188    ### This might be sqlalchemy object or the string of a table name.
189    ### We check for spaces and quotes to see if it might be a weird table.
190    if (
191        ' ' not in str(query_or_table)
192        or (
193            ' ' in str(query_or_table)
194            and str(query_or_table).startswith('"')
195            and str(query_or_table).endswith('"')
196        )
197    ):
198        truncated_table_name = truncate_item_name(str(query_or_table), self.flavor)
199        if truncated_table_name != str(query_or_table) and not silent:
200            warn(
201                f"Table '{query_or_table}' is too long for '{self.flavor}',"
202                + f" will instead read the table '{truncated_table_name}'."
203            )
204
205        query_or_table = sql_item_name(str(query_or_table), self.flavor, schema)
206        if debug:
207            dprint(f"[{self}] Reading from table {query_or_table}")
208        formatted_query = sqlalchemy.text("SELECT * FROM " + str(query_or_table))
209        str_query = f"SELECT * FROM {query_or_table}"
210    else:
211        str_query = query_or_table
212
213    formatted_query = (
214        sqlalchemy.text(str_query)
215        if not is_dask and isinstance(str_query, str)
216        else format_sql_query_for_dask(str_query)
217    )
218
219    chunk_list = []
220    chunk_hook_results = []
221    def _process_chunk(_chunk, _retry_on_failure: bool = True):
222        if self.flavor in TIMEZONE_NAIVE_FLAVORS:
223            for col in utc_dt_cols:
224                _chunk[col] = coerce_timezone(_chunk[col], strip_timezone=False)
225        if not as_hook_results:
226            chunk_list.append(_chunk)
227        if chunk_hook is None:
228            return None
229
230        result = None
231        try:
232            result = chunk_hook(
233                _chunk,
234                workers=workers,
235                chunksize=chunksize,
236                debug=debug,
237                **kw
238            )
239        except Exception:
240            result = False, traceback.format_exc()
241            from meerschaum.utils.formatting import get_console
242            if not silent:
243                get_console().print_exception()
244
245        ### If the chunk fails to process, try it again one more time.
246        if isinstance(result, tuple) and result[0] is False:
247            if _retry_on_failure:
248                return _process_chunk(_chunk, _retry_on_failure=False)
249
250        return result
251
252    try:
253        stream_results = not as_iterator and chunk_hook is not None and chunksize is not None
254        with warnings.catch_warnings():
255            warnings.filterwarnings('ignore', 'case sensitivity issues')
256
257            read_sql_query_kwargs = {
258                'params': params,
259                'dtype': dtype,
260                'coerce_float': coerce_float,
261                'index_col': index_col,
262            }
263            if is_dask:
264                if index_col is None:
265                    dd = None
266                    pd = attempt_import('pandas')
267                    read_sql_query_kwargs.update({
268                        'chunksize': chunksize,
269                    })
270            else:
271                read_sql_query_kwargs.update({
272                    'chunksize': chunksize,
273                })
274
275            if is_dask and dd is not None:
276                ddf = dd.read_sql_query(
277                    formatted_query,
278                    self.URI,
279                    **read_sql_query_kwargs
280                )
281            else:
282
283                def get_chunk_generator(connectable):
284                    chunk_generator = pd.read_sql_query(
285                        formatted_query,
286                        self.engine,
287                        **read_sql_query_kwargs
288                    )
289                    to_return = (
290                        chunk_generator
291                        if as_iterator or chunksize is None
292                        else (
293                            list(pool.imap(_process_chunk, chunk_generator))
294                            if as_hook_results
295                            else None
296                        )
297                    )
298                    return chunk_generator, to_return
299
300                if self.flavor in SKIP_READ_TRANSACTION_FLAVORS:
301                    chunk_generator, to_return = get_chunk_generator(self.engine)
302                else:
303                    with self.engine.begin() as transaction:
304                        with transaction.execution_options(stream_results=stream_results) as connection:
305                            chunk_generator, to_return = get_chunk_generator(connection)
306
307                if to_return is not None:
308                    return to_return
309
310    except Exception as e:
311        if debug:
312            dprint(f"[{self}] Failed to execute query:\n\n{query_or_table}\n\n")
313        if not silent:
314            warn(str(e), stacklevel=3)
315        from meerschaum.utils.formatting import get_console
316        if not silent:
317            get_console().print_exception()
318
319        return None
320
321    if is_dask and dd is not None:
322        ddf = ddf.reset_index()
323        return ddf
324
325    chunk_list = []
326    read_chunks = 0
327    chunk_hook_results = []
328    if chunksize is None:
329        chunk_list.append(chunk_generator)
330    elif as_iterator:
331        return chunk_generator
332    else:
333        try:
334            for chunk in chunk_generator:
335                if chunk_hook is not None:
336                    chunk_hook_results.append(
337                        chunk_hook(chunk, chunksize=chunksize, debug=debug, **kw)
338                    )
339                chunk_list.append(chunk)
340                read_chunks += 1
341                if chunks is not None and read_chunks >= chunks:
342                    break
343        except Exception as e:
344            warn(f"[{self}] Failed to retrieve query results:\n" + str(e), stacklevel=3)
345            from meerschaum.utils.formatting import get_console
346            if not silent:
347                get_console().print_exception()
348
349    read_chunks = 0
350    try:
351        for chunk in chunk_generator:
352            if chunk_hook is not None:
353                chunk_hook_results.append(
354                    chunk_hook(chunk, chunksize=chunksize, debug=debug, **kw)
355                )
356            chunk_list.append(chunk)
357            read_chunks += 1
358            if chunks is not None and read_chunks >= chunks:
359                break
360    except Exception as e:
361        warn(f"[{self}] Failed to retrieve query results:\n" + str(e), stacklevel=3)
362        from meerschaum.utils.formatting import get_console
363        if not silent:
364            get_console().print_exception()
365
366        return None
367
368    ### If no chunks returned, read without chunks
369    ### to get columns
370    if len(chunk_list) == 0:
371        with warnings.catch_warnings():
372            warnings.filterwarnings('ignore', 'case sensitivity issues')
373            _ = read_sql_query_kwargs.pop('chunksize', None)
374            with self.engine.begin() as connection:
375                chunk_list.append(
376                    pd.read_sql_query(
377                        formatted_query,
378                        connection,
379                        **read_sql_query_kwargs
380                    )
381                )
382
383    ### call the hook on any missed chunks.
384    if chunk_hook is not None and len(chunk_list) > len(chunk_hook_results):
385        for c in chunk_list[len(chunk_hook_results):]:
386            chunk_hook_results.append(
387                chunk_hook(c, chunksize=chunksize, debug=debug, **kw)
388            )
389
390    ### chunksize is not None so must iterate
391    if debug:
392        end = time.perf_counter()
393        dprint(f"Fetched {len(chunk_list)} chunks in {round(end - start, 2)} seconds.")
394
395    if as_hook_results:
396        return chunk_hook_results
397    
398    ### Skip `pd.concat()` if `as_chunks` is specified.
399    if as_chunks:
400        for c in chunk_list:
401            c.reset_index(drop=True, inplace=True)
402            for col in get_numeric_cols(c):
403                c[col] = c[col].apply(lambda x: x.canonical() if isinstance(x, Decimal) else x)
404        return chunk_list
405
406    df = pd.concat(chunk_list).reset_index(drop=True)
407    ### NOTE: The calls to `canonical()` are to drop leading and trailing zeroes.
408    for col in get_numeric_cols(df):
409        df[col] = df[col].apply(lambda x: x.canonical() if isinstance(x, Decimal) else x)
410
411    return df