meerschaum.connectors.sql

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

232    @property
233    def Session(self):
234        if '_Session' not in self.__dict__:
235            if self.engine is None:
236                return None
237
238            from meerschaum.utils.packages import attempt_import
239            sqlalchemy_orm = attempt_import('sqlalchemy.orm', lazy=False)
240            session_factory = sqlalchemy_orm.sessionmaker(self.engine)
241            self._Session = sqlalchemy_orm.scoped_session(session_factory)
242
243        return self._Session

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

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

281    @property
282    def URI(self) -> str:
283        """
284        Return the URI connection string.
285        """
286        _ = self.engine
287        return str(self._engine_str)

289    @property
290    def IS_THREAD_SAFE(self) -> str:
291        """
292        Return whether this connector may be multithreaded.
293        """
294        if self.flavor in ('duckdb', 'oracle'):
295            return False
296        if self.flavor in ('sqlite', 'geopackage'):
297            return ':memory:' not in self.URI
298        return True

300    @property
301    def metadata(self):
302        """
303        Return the metadata bound to this configured schema.
304        """
305        from meerschaum.utils.packages import attempt_import
306        sqlalchemy = attempt_import('sqlalchemy', lazy=False)
307        if '_metadata' not in self.__dict__:
308            self._metadata = sqlalchemy.MetaData(schema=self.schema)
309        return self._metadata

311    @property
312    def instance_schema(self):
313        """
314        Return the schema name for Meerschaum tables. 
315        """
316        return self.schema

318    @property
319    def internal_schema(self):
320        """
321        Return the schema name for internal tables. 
322        """
323        from meerschaum._internal.static import STATIC_CONFIG
324        from meerschaum.utils.sql import NO_SCHEMA_FLAVORS
325        schema_name = self.__dict__.get('internal_schema', None) or (
326            STATIC_CONFIG['sql']['internal_schema']
327            if self.flavor not in NO_SCHEMA_FLAVORS
328            else self.schema
329        )
330
331        if '_internal_schema' not in self.__dict__:
332            self._internal_schema = schema_name
333        return self._internal_schema

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

351    @property
352    def db_version(self) -> Union[str, None]:
353        """
354        Return the database version.
355        """
356        _db_version = self.__dict__.get('_db_version', None)
357        if _db_version is not None:
358            return _db_version
359
360        from meerschaum.utils.sql import get_db_version
361        self._db_version = get_db_version(self)
362        return self._db_version

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

383    def get_metadata_cache_path(self, kind: str = 'json') -> pathlib.Path:
384        """
385        Return the path to the file to which to write metadata cache.
386        """
387        from meerschaum.config.paths import SQL_CONN_CACHE_RESOURCES_PATH
388        filename = (
389            f'{self.label}-metadata.pkl'
390            if kind == 'pkl'
391            else f'{self.label}.json'
392        )
393        return 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        from meerschaum.config.paths import SQLITE_DB_PATH
 85        _database = 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            warn(
213                f"Table '{query_or_table}' is too long for '{self.flavor}',"
214                + f" will instead read the table '{truncated_table_name}'."
215            )
216
217        query_or_table = sql_item_name(str(query_or_table), self.flavor, schema)
218        if debug:
219            dprint(f"[{self}] Reading from table {query_or_table}")
220        formatted_query = sqlalchemy.text("SELECT * FROM " + str(query_or_table))
221        str_query = f"SELECT * FROM {query_or_table}"
222    else:
223        str_query = query_or_table
224
225    formatted_query = (
226        sqlalchemy.text(str_query)
227        if not is_dask and isinstance(str_query, str)
228        else format_sql_query_for_dask(str_query)
229    )
230
231    def _get_chunk_args_kwargs(_chunk):
232        return filter_arguments(
233            chunk_hook,
234            _chunk,
235            workers=workers,
236            chunksize=chunksize,
237            debug=debug,
238            **kw
239        )
240
241    chunk_list = []
242    chunk_hook_results = []
243    def _process_chunk(_chunk, _retry_on_failure: bool = True):
244        if self.flavor in TIMEZONE_NAIVE_FLAVORS:
245            for col in utc_dt_cols:
246                _chunk[col] = coerce_timezone(_chunk[col], strip_utc=False)
247        if not as_hook_results:
248            chunk_list.append(_chunk)
249
250        if chunk_hook is None:
251            return None
252
253        chunk_args, chunk_kwargs = _get_chunk_args_kwargs(_chunk)
254
255        result = None
256        try:
257            result = chunk_hook(*chunk_args, **chunk_kwargs)
258        except Exception:
259            result = False, traceback.format_exc()
260            from meerschaum.utils.formatting import get_console
261            if not silent:
262                get_console().print_exception()
263
264        ### If the chunk fails to process, try it again one more time.
265        if isinstance(result, tuple) and result[0] is False:
266            if _retry_on_failure:
267                return _process_chunk(_chunk, _retry_on_failure=False)
268
269        return result
270
271    try:
272        stream_results = not as_iterator and chunk_hook is not None and chunksize is not None
273        with warnings.catch_warnings():
274            warnings.filterwarnings('ignore', 'case sensitivity issues')
275
276            read_sql_query_kwargs = {
277                'params': params,
278                'dtype': dtype,
279                'coerce_float': coerce_float,
280                'index_col': index_col,
281            }
282            if is_dask:
283                if index_col is None:
284                    dd = None
285                    pd = attempt_import('pandas')
286                    read_sql_query_kwargs.update({
287                        'chunksize': chunksize,
288                    })
289            else:
290                read_sql_query_kwargs.update({
291                    'chunksize': chunksize,
292                })
293
294            if is_dask and dd is not None:
295                ddf = dd.read_sql_query(
296                    formatted_query,
297                    self.URI,
298                    **read_sql_query_kwargs
299                )
300            else:
301
302                def get_chunk_generator(connectable):
303                    chunk_generator = pd.read_sql_query(
304                        formatted_query,
305                        connectable, # NOTE: test this against `self.engine`.
306                        **read_sql_query_kwargs
307                    )
308
309                    to_return = (
310                        (
311                            chunk_generator
312                            if not (as_hook_results or chunksize is None)
313                            else (
314                                _process_chunk(_chunk)
315                                for _chunk in chunk_generator
316                            )
317                        )
318                        if as_iterator or chunksize is None
319                        else (
320                            list(pool.imap(_process_chunk, chunk_generator))
321                            if as_hook_results
322                            else None
323                        )
324                    )
325                    return chunk_generator, to_return
326
327                if self.flavor in SKIP_READ_TRANSACTION_FLAVORS:
328                    chunk_generator, to_return = get_chunk_generator(self.engine)
329                else:
330                    with self.engine.begin() as transaction:
331                        with transaction.execution_options(
332                            stream_results=stream_results,
333                        ) as connection:
334                            chunk_generator, to_return = get_chunk_generator(connection)
335
336                if to_return is not None:
337                    return to_return
338
339    except Exception as e:
340        if debug:
341            dprint(f"[{self}] Failed to execute query:\n\n{query_or_table}\n\n")
342        if not silent:
343            warn(str(e), stacklevel=3)
344        from meerschaum.utils.formatting import get_console
345        if not silent:
346            get_console().print_exception()
347
348        return None
349
350    if is_dask and dd is not None:
351        ddf = ddf.reset_index()
352        return ddf
353
354    chunk_list = []
355    read_chunks = 0
356    chunk_hook_results = []
357    if chunksize is None:
358        chunk_list.append(chunk_generator)
359    elif as_iterator:
360        return chunk_generator
361    else:
362        try:
363            for chunk in chunk_generator:
364                if chunk_hook is not None:
365                    chunk_args, chunk_kwargs = _get_chunk_args_kwargs(chunk)
366                    chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
367                chunk_list.append(chunk)
368                read_chunks += 1
369                if chunks is not None and read_chunks >= chunks:
370                    break
371        except Exception as e:
372            warn(f"[{self}] Failed to retrieve query results:\n" + str(e), stacklevel=3)
373            from meerschaum.utils.formatting import get_console
374            if not silent:
375                get_console().print_exception()
376
377    read_chunks = 0
378    try:
379        for chunk in chunk_generator:
380            if chunk_hook is not None:
381                chunk_args, chunk_kwargs = _get_chunk_args_kwargs(chunk)
382                chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
383            chunk_list.append(chunk)
384            read_chunks += 1
385            if chunks is not None and read_chunks >= chunks:
386                break
387    except Exception as e:
388        warn(f"[{self}] Failed to retrieve query results:\n" + str(e), stacklevel=3)
389        from meerschaum.utils.formatting import get_console
390        if not silent:
391            get_console().print_exception()
392
393        return None
394
395    ### If no chunks returned, read without chunks
396    ### to get columns
397    if len(chunk_list) == 0:
398        with warnings.catch_warnings():
399            warnings.filterwarnings('ignore', 'case sensitivity issues')
400            _ = read_sql_query_kwargs.pop('chunksize', None)
401            with self.engine.begin() as connection:
402                chunk_list.append(
403                    pd.read_sql_query(
404                        formatted_query,
405                        connection,
406                        **read_sql_query_kwargs
407                    )
408                )
409
410    ### call the hook on any missed chunks.
411    if chunk_hook is not None and len(chunk_list) > len(chunk_hook_results):
412        for c in chunk_list[len(chunk_hook_results):]:
413            chunk_args, chunk_kwargs = _get_chunk_args_kwargs(c)
414            chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
415
416    ### chunksize is not None so must iterate
417    if debug:
418        end = time.perf_counter()
419        dprint(f"Fetched {len(chunk_list)} chunks in {round(end - start, 2)} seconds.")
420
421    if as_hook_results:
422        return chunk_hook_results
423    
424    ### Skip `pd.concat()` if `as_chunks` is specified.
425    if as_chunks:
426        for c in chunk_list:
427            c.reset_index(drop=True, inplace=True)
428            for col in get_numeric_cols(c):
429                c[col] = c[col].apply(lambda x: x.canonical() if isinstance(x, Decimal) else x)
430        return chunk_list
431
432    df = pd.concat(chunk_list).reset_index(drop=True)
433    ### NOTE: The calls to `canonical()` are to drop leading and trailing zeroes.
434    for col in get_numeric_cols(df):
435        df[col] = df[col].apply(lambda x: x.canonical() if isinstance(x, Decimal) else x)
436
437    return df