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            kw['uri'] = uri
156            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
157            label = label or from_uri_params.get('label', None)
158            _ = from_uri_params.pop('label', None)
159
160            ### Sometimes the flavor may be provided with a URI.
161            kw.update(from_uri_params)
162            if flavor:
163                kw['flavor'] = flavor
164
165        ### set __dict__ in base class
166        super().__init__(
167            'sql',
168            label = label or self.__dict__.get('label', None),
169            **kw
170        )
171
172        if self.__dict__.get('flavor', None) == 'sqlite':
173            self._reset_attributes()
174            self._set_attributes(
175                'sql',
176                label = label,
177                inherit_default = False,
178                **kw
179            )
180            ### For backwards compatability reasons, set the path for sql:local if its missing.
181            if self.label == 'local' and not self.__dict__.get('database', None):
182                from meerschaum.config._paths import SQLITE_DB_PATH
183                self.database = str(SQLITE_DB_PATH)
184
185        ### ensure flavor and label are set accordingly
186        if 'flavor' not in self.__dict__:
187            if flavor is None and 'uri' not in self.__dict__:
188                raise Exception(
189                    f"    Missing flavor. Provide flavor as a key for '{self}'."
190                )
191            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
192
193        if self.flavor == 'postgres':
194            self.flavor = 'postgresql'
195
196        self._debug = debug
197        ### Store the PID and thread at initialization
198        ### so we can dispose of the Pool in child processes or threads.
199        import os
200        import threading
201        self._pid = os.getpid()
202        self._thread_ident = threading.current_thread().ident
203        self._sessions = {}
204        self._locks = {'_sessions': threading.RLock(), }
205
206        ### verify the flavor's requirements are met
207        if self.flavor not in self.flavor_configs:
208            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
209        if not self.__dict__.get('uri'):
210            self.verify_attributes(
211                self.flavor_configs[self.flavor].get('requirements', set()),
212                debug=debug,
213            )
214
215        if wait:
216            from meerschaum.connectors.poll import retry_connect
217            retry_connect(connector=self, debug=debug)
218
219        if connect:
220            if not self.test_connection(debug=debug):
221                warn(f"Failed to connect with connector '{self}'!", stack=False)

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

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

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

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

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

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

302    @property
303    def instance_schema(self):
304        """
305        Return the schema name for Meerschaum tables. 
306        """
307        return self.schema

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

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

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

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

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