meerschaum.connectors.sql

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

218    @property
219    def Session(self):
220        if '_Session' not in self.__dict__:
221            if self.engine is None:
222                return None
223
224            from meerschaum.utils.packages import attempt_import
225            sqlalchemy_orm = attempt_import('sqlalchemy.orm')
226            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
227            self._Session = sqlalchemy_orm.scoped_session(session_factory)
228
229        return self._Session

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

259    @property
260    def DATABASE_URL(self) -> str:
261        """
262        Return the URI connection string (alias for `SQLConnector.URI`.
263        """
264        _ = self.engine
265        return str(self._engine_str)

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

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

287    @property
288    def metadata(self):
289        """
290        Return the metadata bound to this configured schema.
291        """
292        from meerschaum.utils.packages import attempt_import
293        sqlalchemy = attempt_import('sqlalchemy')
294        if '_metadata' not in self.__dict__:
295            self._metadata = sqlalchemy.MetaData(schema=self.schema)
296        return self._metadata

299    @property
300    def instance_schema(self):
301        """
302        Return the schema name for Meerschaum tables. 
303        """
304        return self.schema

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

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

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

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

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