meerschaum.connectors.sql

153    def __init__(
154        self,
155        label: Optional[str] = None,
156        flavor: Optional[str] = None,
157        wait: bool = False,
158        connect: bool = False,
159        debug: bool = False,
160        **kw: Any
161    ):
162        """
163        Parameters
164        ----------
165        label: str, default 'main'
166            The identifying label for the connector.
167            E.g. for `sql:main`, 'main' is the label.
168            Defaults to 'main'.
169
170        flavor: Optional[str], default None
171            The database flavor, e.g.
172            `'sqlite'`, `'postgresql'`, `'cockroachdb'`, etc.
173            To see supported flavors, run the `bootstrap connectors` command.
174
175        wait: bool, default False
176            If `True`, block until a database connection has been made.
177            Defaults to `False`.
178
179        connect: bool, default False
180            If `True`, immediately attempt to connect the database and raise
181            a warning if the connection fails.
182            Defaults to `False`.
183
184        debug: bool, default False
185            Verbosity toggle.
186            Defaults to `False`.
187
188        kw: Any
189            All other arguments will be passed to the connector's attributes.
190            Therefore, a connector may be made without being registered,
191            as long enough parameters are supplied to the constructor.
192        """
193        if 'uri' in kw:
194            uri = kw['uri']
195            if uri.startswith('postgres') and not uri.startswith('postgresql'):
196                uri = uri.replace('postgres', 'postgresql', 1)
197            if uri.startswith('postgresql') and not uri.startswith('postgresql+'):
198                uri = uri.replace('postgresql://', 'postgresql+psycopg://', 1)
199            if uri.startswith('timescaledb://'):
200                uri = uri.replace('timescaledb://', 'postgresql+psycopg://', 1)
201                flavor = 'timescaledb'
202            if uri.startswith('timescaledb-ha://'):
203                uri = uri.replace('timescaledb-ha://', 'postgresql+psycopg://', 1)
204                flavor = 'timescaledb-ha'
205            if uri.startswith('postgis://'):
206                uri = uri.replace('postgis://', 'postgresql+psycopg://', 1)
207                flavor = 'postgis'
208            kw['uri'] = uri
209            from_uri_params = self.from_uri(kw['uri'], as_dict=True)
210            label = label or from_uri_params.get('label', None)
211            _ = from_uri_params.pop('label', None)
212
213            ### Sometimes the flavor may be provided with a URI.
214            kw.update(from_uri_params)
215            if flavor:
216                kw['flavor'] = flavor
217
218        ### set __dict__ in base class
219        super().__init__(
220            'sql',
221            label = label or self.__dict__.get('label', None),
222            **kw
223        )
224
225        if self.__dict__.get('flavor', None) in ('sqlite', 'geopackage'):
226            self._reset_attributes()
227            self._set_attributes(
228                'sql',
229                label = label,
230                inherit_default = False,
231                **kw
232            )
233            ### For backwards compatability reasons, set the path for sql:local if its missing.
234            if (
235                self.label == 'local'
236                and self.__dict__.get('database', None) in (None, '{SQLITE_DB_PATH}')
237            ):
238                import meerschaum.config.paths as paths
239                self.database = paths.SQLITE_DB_PATH.as_posix()
240
241        ### ensure flavor and label are set accordingly
242        if 'flavor' not in self.__dict__:
243            if flavor is None and 'uri' not in self.__dict__:
244                raise ValueError(
245                    f"    Missing flavor. Provide flavor as a key for '{self}'."
246                )
247            self.flavor = flavor or self.parse_uri(self.__dict__['uri']).get('flavor', None)
248
249        if self.flavor == 'postgres':
250            self.flavor = 'postgresql'
251
252        self._debug = debug
253        ### Store the PID and thread at initialization
254        ### so we can dispose of the Pool in child processes or threads.
255        import os
256        import threading
257        self._pid = os.getpid()
258        self._thread_ident = threading.current_thread().ident
259        self._sessions = {}
260        self._locks = {'_sessions': threading.RLock(), }
261
262        ### verify the flavor's requirements are met
263        if self.flavor not in self.flavor_configs:
264            error(f"Flavor '{self.flavor}' is not supported by Meerschaum SQLConnector")
265        if not self.__dict__.get('uri'):
266            self.verify_attributes(
267                self.flavor_configs[self.flavor].get('requirements', set()),
268                debug=debug,
269            )
270
271        if wait:
272            from meerschaum.connectors.poll import retry_connect
273            retry_connect(connector=self, debug=debug)
274
275        if connect:
276            if not self.test_connection(debug=debug):
277                warn(f"Failed to connect with connector '{self}'!", stack=False)

279    @property
280    def Session(self):
281        if '_Session' not in self.__dict__:
282            if self.engine is None:
283                return None
284
285            from meerschaum.utils.packages import attempt_import
286            sqlalchemy_orm = attempt_import('sqlalchemy.orm', lazy=False)
287            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
288            self._Session = sqlalchemy_orm.scoped_session(session_factory)
289
290        return self._Session

292    @property
293    def engine(self):
294        """
295        Return the SQLAlchemy engine connected to the configured database.
296        """
297        import os
298        import threading
299        if '_engine' not in self.__dict__:
300            self._engine, self._engine_str = self.create_engine(include_uri=True)
301
302        same_process = os.getpid() == self._pid
303        same_thread = threading.current_thread().ident == self._thread_ident
304
305        ### handle child processes
306        if not same_process:
307            self._pid = os.getpid()
308            self._thread = threading.current_thread()
309            warn("Different PID detected. Disposing of connections...")
310            self._engine.dispose()
311
312        ### handle different threads
313        if not same_thread:
314            if self.flavor == 'duckdb':
315                warn("Different thread detected.")
316                self._engine.dispose()
317
318        return self._engine

320    @property
321    def DATABASE_URL(self) -> str:
322        """
323        Return the URI connection string (alias for `SQLConnector.URI`.
324        """
325        _ = self.engine
326        return str(self._engine_str)

328    @property
329    def URI(self) -> str:
330        """
331        Return the URI connection string.
332        """
333        _ = self.engine
334        return str(self._engine_str)

336    @property
337    def IS_THREAD_SAFE(self) -> str:
338        """
339        Return whether this connector may be multithreaded.
340        """
341        if self.flavor in ('duckdb', 'oracle'):
342            return False
343        if self.flavor in ('sqlite', 'geopackage'):
344            return ':memory:' not in self.URI
345        return True

347    @property
348    def metadata(self):
349        """
350        Return the metadata bound to this configured schema.
351        """
352        from meerschaum.utils.packages import attempt_import
353        sqlalchemy = attempt_import('sqlalchemy', lazy=False)
354        if '_metadata' not in self.__dict__:
355            self._metadata = sqlalchemy.MetaData(schema=self.schema)
356        return self._metadata

358    @property
359    def instance_schema(self):
360        """
361        Return the schema name for Meerschaum tables. 
362        """
363        return self.schema

365    @property
366    def internal_schema(self):
367        """
368        Return the schema name for internal tables. 
369        """
370        from meerschaum._internal.static import STATIC_CONFIG
371        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
372        schema_name = self.__dict__.get('internal_schema', None) or (
373            STATIC_CONFIG['sql']['internal_schema']
374            if self.flavor not in NO_SCHEMA_FLAVORS
375            else self.schema
376        )
377
378        if '_internal_schema' not in self.__dict__:
379            self._internal_schema = schema_name
380        return self._internal_schema

382    @property
383    def db(self) -> Optional[databases.Database]:
384        from meerschaum.utils.packages import attempt_import
385        databases = attempt_import('databases', lazy=False, install=True)
386        url = self.DATABASE_URL
387        if 'mysql' in url:
388            url = url.replace('+pymysql', '')
389        if '_db' not in self.__dict__:
390            try:
391                self._db = databases.Database(url)
392            except KeyError:
393                ### Likely encountered an unsupported flavor.
394                from meerschaum.utils.warnings import warn
395                self._db = None
396        return self._db

398    @property
399    def db_version(self) -> Union[str, None]:
400        """
401        Return the database version.
402        """
403        _db_version = self.__dict__.get('_db_version', None)
404        if _db_version is not None:
405            return _db_version
406
407        from meerschaum.utils.sql import get_db_version
408        self._db_version = get_db_version(self)
409        return self._db_version

411    @property
412    def schema(self) -> Union[str, None]:
413        """
414        Return the default schema to use.
415        A value of `None` will not prepend a schema.
416        """
417        if 'schema' in self.__dict__:
418            return self.__dict__['schema']
419
420        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
421        if self.flavor in NO_SCHEMA_FLAVORS:
422            self.__dict__['schema'] = None
423            return None
424
425        sqlalchemy = mrsm.attempt_import('sqlalchemy', lazy=False)
426        _schema = sqlalchemy.inspect(self.engine).default_schema_name
427        self.__dict__['schema'] = _schema
428        return _schema

430    def get_metadata_cache_path(self, kind: str = 'json') -> pathlib.Path:
431        """
432        Return the path to the file to which to write metadata cache.
433        """
434        import meerschaum.config.paths as paths
435        filename = (
436            f'{self.label}-metadata.pkl'
437            if kind == 'pkl'
438            else f'{self.label}.json'
439        )
440        return paths.SQL_CONN_CACHE_RESOURCES_PATH / filename

 45def create_engine(
 46    self,
 47    include_uri: bool = False,
 48    debug: bool = False,
 49    **kw
 50) -> 'sqlalchemy.engine.Engine':
 51    """Create a sqlalchemy engine by building the engine string."""
 52    from meerschaum.utils.packages import attempt_import
 53    from meerschaum.utils.warnings import error, warn
 54    sqlalchemy = attempt_import('sqlalchemy', lazy=False)
 55    import urllib
 56    import copy
 57    ### Install and patch required drivers.
 58    if self.flavor in install_flavor_drivers:
 59        _ = attempt_import(
 60            *install_flavor_drivers[self.flavor],
 61            debug=debug,
 62            lazy=False,
 63            warn=False,
 64        )
 65        if self.flavor == 'mssql':
 66            _init_mssql_sqlalchemy()
 67
 68    ### supplement missing values with defaults (e.g. port number)
 69    for a, value in flavor_configs[self.flavor]['defaults'].items():
 70        if a not in self.__dict__:
 71            self.__dict__[a] = value
 72
 73    ### Verify that everything is in order.
 74    if self.flavor not in flavor_configs:
 75        error(f"Cannot create a connector with the flavor '{self.flavor}'.")
 76
 77    _engine = flavor_configs[self.flavor].get('engine', None)
 78    _username = self.__dict__.get('username', None)
 79    _password = self.__dict__.get('password', None)
 80    _host = self.__dict__.get('host', None)
 81    _port = self.__dict__.get('port', None)
 82    _database = self.__dict__.get('database', None)
 83    if _database == '{SQLITE_DB_PATH}':
 84        import meerschaum.config.paths as paths
 85        _database = paths.SQLITE_DB_PATH.as_posix()
 86    _options = self.__dict__.get('options', {})
 87    if isinstance(_options, str):
 88        _options = dict(urllib.parse.parse_qsl(_options))
 89    _uri = self.__dict__.get('uri', None)
 90
 91    ### Handle registering specific dialects (due to installing in virtual environments).
 92    if self.flavor in flavor_dialects:
 93        sqlalchemy.dialects.registry.register(*flavor_dialects[self.flavor])
 94
 95    ### self._sys_config was deepcopied and can be updated safely
 96    if self.flavor in ("sqlite", "duckdb", "geopackage"):
 97        engine_str = f"{_engine}:///{_database}" if not _uri else _uri
 98        if 'create_engine' not in self._sys_config:
 99            self._sys_config['create_engine'] = {}
100        if 'connect_args' not in self._sys_config['create_engine']:
101            self._sys_config['create_engine']['connect_args'] = {}
102        self._sys_config['create_engine']['connect_args'].update({"check_same_thread": False})
103    else:
104        engine_str = (
105            _engine + "://" + (_username if _username is not None else '') +
106            ((":" + urllib.parse.quote_plus(_password)) if _password is not None else '') +
107            "@" + _host + ((":" + str(_port)) if _port is not None else '') +
108            (("/" + _database) if _database is not None else '')
109            + (("?" + urllib.parse.urlencode(_options)) if _options else '')
110        ) if not _uri else _uri
111
112        ### Sometimes the timescaledb:// flavor can slip in.
113        if _uri and self.flavor in _uri:
114            if self.flavor in ('timescaledb', 'timescaledb-ha', 'postgis'):
115                engine_str = engine_str.replace(self.flavor, 'postgresql', 1)
116            elif _uri.startswith('postgresql://'):
117                engine_str = engine_str.replace('postgresql://', 'postgresql+psycopg2://')
118
119    if debug:
120        dprint(
121            (
122                (engine_str.replace(':' + _password, ':' + ('*' * len(_password))))
123                    if _password is not None else engine_str
124            ) + '\n' + f"{self._sys_config.get('create_engine', {}).get('connect_args', {})}"
125        )
126
127    _kw_copy = copy.deepcopy(kw)
128
129    ### NOTE: Order of inheritance:
130    ###       1. Defaults
131    ###       2. System configuration
132    ###       3. Connector configuration
133    ###       4. Keyword arguments
134    _create_engine_args = flavor_configs.get(self.flavor, {}).get('create_engine', {})
135    def _apply_create_engine_args(update):
136        if 'ALL' not in flavor_configs[self.flavor].get('omit_create_engine', {}):
137            _create_engine_args.update(
138                { k: v for k, v in update.items()
139                    if 'omit_create_engine' not in flavor_configs[self.flavor]
140                        or k not in flavor_configs[self.flavor].get('omit_create_engine')
141                }
142            )
143    _apply_create_engine_args(self._sys_config.get('create_engine', {}))
144    _apply_create_engine_args(self.__dict__.get('create_engine', {}))
145    _apply_create_engine_args(_kw_copy)
146
147    try:
148        engine = sqlalchemy.create_engine(
149            engine_str,
150            ### I know this looks confusing, and maybe it's bad code,
151            ### but it's simple. It dynamically parses the config string
152            ### and splits it to separate the class name (QueuePool)
153            ### from the module name (sqlalchemy.pool).
154            poolclass    = getattr(
155                attempt_import(
156                    ".".join(self._sys_config['poolclass'].split('.')[:-1])
157                ),
158                self._sys_config['poolclass'].split('.')[-1]
159            ),
160            echo         = debug,
161            **_create_engine_args
162        )
163    except Exception:
164        warn(f"Failed to create connector '{self}':\n{traceback.format_exc()}", stack=False)
165        engine = None
166
167    if include_uri:
168        return engine, engine_str
169    return engine

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