meerschaum.connectors.sql

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

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

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

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

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

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

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

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

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

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

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

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

432    def get_metadata_cache_path(self, kind: str = 'json') -> pathlib.Path:
433        """
434        Return the path to the file to which to write metadata cache.
435        """
436        import meerschaum.config.paths as paths
437        filename = (
438            f'{self.label}-metadata.pkl'
439            if kind == 'pkl'
440            else f'{self.label}.json'
441        )
442        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    is_dask = 'dask' in pd.__name__
149    dd = pd if is_dask else None
150    pandas = attempt_import('pandas')
151    npartitions = chunksize_to_npartitions(chunksize)
152    if is_dask:
153        chunksize = None
154
155    schema = schema or self.schema
156    utc_dt_cols = [
157        col
158        for col, typ in dtype.items()
159        if are_dtypes_equal(typ, 'datetime') and 'utc' in typ.lower()
160    ] if dtype else []
161
162    if dtype and utc_dt_cols and self.flavor in TIMEZONE_NAIVE_FLAVORS:
163        dtype = dtype.copy()
164        for col in utc_dt_cols:
165            dtype[col] = 'datetime64[us]'
166
167    pool = get_pool(workers=workers)
168    sqlalchemy = attempt_import("sqlalchemy", lazy=False)
169    default_chunksize = self._sys_config.get('chunksize', None)
170    chunksize = chunksize if chunksize != -1 else default_chunksize
171    if chunksize is None and as_iterator:
172        if not silent and self.flavor not in _disallow_chunks_flavors:
173            warn(
174                "An iterator may only be generated if chunksize is not None.\n"
175                + "Falling back to a chunksize of 1000.", stacklevel=3,
176            )
177        chunksize = 1000
178    if chunksize is not None and self.flavor in _max_chunks_flavors:
179        if chunksize > _max_chunks_flavors[self.flavor]:
180            if chunksize != default_chunksize:
181                warn(
182                    f"The specified chunksize of {chunksize} exceeds the maximum of "
183                    + f"{_max_chunks_flavors[self.flavor]} for flavor '{self.flavor}'.\n"
184                    + f"    Falling back to a chunksize of {_max_chunks_flavors[self.flavor]}.",
185                    stacklevel=3,
186                )
187            chunksize = _max_chunks_flavors[self.flavor]
188
189    if chunksize is not None and self.flavor in _disallow_chunks_flavors:
190        chunksize = None
191
192    if debug:
193        import time
194        start = time.perf_counter()
195        dprint(f"[{self}]\n{query_or_table}")
196        dprint(f"[{self}] Fetching with chunksize: {chunksize}")
197
198    ### This might be sqlalchemy object or the string of a table name.
199    ### We check for spaces and quotes to see if it might be a weird table.
200    if (
201        ' ' not in str(query_or_table)
202        or (
203            ' ' in str(query_or_table)
204            and str(query_or_table).startswith('"')
205            and str(query_or_table).endswith('"')
206        )
207    ):
208        truncated_table_name = truncate_item_name(str(query_or_table), self.flavor)
209        if truncated_table_name != str(query_or_table) and not silent:
210            if self.flavor not in ('oracle', 'mysql', 'mariadb'):
211                warn(
212                    f"Table '{query_or_table}' is too long for '{self.flavor}',"
213                    + f" will instead read the table '{truncated_table_name}'."
214                )
215
216        query_or_table = sql_item_name(str(query_or_table), self.flavor, schema)
217        if debug:
218            dprint(f"[{self}] Reading from table {query_or_table}")
219        formatted_query = sqlalchemy.text("SELECT * FROM " + str(query_or_table))
220        str_query = f"SELECT * FROM {query_or_table}"
221    else:
222        str_query = query_or_table
223
224    formatted_query = (
225        sqlalchemy.text(str_query)
226        if not is_dask and isinstance(str_query, str)
227        else format_sql_query_for_dask(str_query)
228    )
229
230    def _get_chunk_args_kwargs(_chunk):
231        return filter_arguments(
232            chunk_hook,
233            _chunk,
234            workers=workers,
235            chunksize=chunksize,
236            debug=debug,
237            **kw
238        )
239
240    chunk_list = []
241    chunk_hook_results = []
242    def _process_chunk(_chunk, _retry_on_failure: bool = True):
243        if self.flavor in TIMEZONE_NAIVE_FLAVORS:
244            for col in utc_dt_cols:
245                _chunk[col] = coerce_timezone(_chunk[col], strip_utc=False)
246        if not as_hook_results:
247            chunk_list.append(_chunk)
248
249        if chunk_hook is None:
250            return None
251
252        chunk_args, chunk_kwargs = _get_chunk_args_kwargs(_chunk)
253
254        result = None
255        try:
256            result = chunk_hook(*chunk_args, **chunk_kwargs)
257        except Exception:
258            result = False, traceback.format_exc()
259            from meerschaum.utils.formatting import get_console
260            if not silent:
261                get_console().print_exception()
262
263        ### If the chunk fails to process, try it again one more time.
264        if isinstance(result, tuple) and result[0] is False:
265            if _retry_on_failure:
266                return _process_chunk(_chunk, _retry_on_failure=False)
267
268        return result
269
270    try:
271        stream_results = not as_iterator and chunk_hook is not None and chunksize is not None
272        with warnings.catch_warnings():
273            warnings.filterwarnings('ignore', 'case sensitivity issues')
274
275            read_sql_query_kwargs = {
276                'params': params,
277                'dtype': dtype,
278                'coerce_float': coerce_float,
279                'index_col': index_col,
280            }
281            if is_dask:
282                if index_col is None:
283                    dd = None
284                    pd = attempt_import('pandas')
285                    read_sql_query_kwargs.update({
286                        'chunksize': chunksize,
287                    })
288            else:
289                read_sql_query_kwargs.update({
290                    'chunksize': chunksize,
291                })
292
293            if is_dask and dd is not None:
294                ddf = dd.read_sql_query(
295                    formatted_query,
296                    self.URI,
297                    **read_sql_query_kwargs
298                )
299            else:
300
301                def get_chunk_generator(connectable):
302                    chunk_generator = pd.read_sql_query(
303                        formatted_query,
304                        connectable, # NOTE: test this against `self.engine`.
305                        **read_sql_query_kwargs
306                    )
307
308                    to_return = (
309                        iter((chunk_generator,))
310                        if chunksize is None and as_iterator
311                        else (
312                            chunk_generator
313                            if as_iterator
314                            else (
315                                list(pool.imap(_process_chunk, chunk_generator))
316                                if as_hook_results and chunksize is not None
317                                else None
318                            )
319                        )
320                    )
321                    return chunk_generator, to_return
322
323                if self.flavor in SKIP_READ_TRANSACTION_FLAVORS:
324                    chunk_generator, to_return = get_chunk_generator(self.engine)
325                else:
326                    with self.engine.begin() as transaction:
327                        with transaction.execution_options(
328                            stream_results=stream_results,
329                        ) as connection:
330                            chunk_generator, to_return = get_chunk_generator(connection)
331
332                if to_return is not None:
333                    return to_return
334
335    except Exception as e:
336        if debug:
337            dprint(f"[{self}] Failed to execute query:\n\n{query_or_table}\n\n")
338        if not silent:
339            warn(str(e), stacklevel=3)
340        from meerschaum.utils.formatting import get_console
341        if not silent:
342            get_console().print_exception()
343
344        return None
345
346    if is_dask and dd is not None:
347        ddf = ddf.reset_index()
348        return ddf
349
350    chunk_list = []
351    read_chunks = 0
352    chunk_hook_results = []
353    if chunksize is None:
354        chunk_list.append(chunk_generator)
355    elif as_iterator:
356        return chunk_generator
357    else:
358        try:
359            for chunk in chunk_generator:
360                if chunk_hook is not None:
361                    chunk_args, chunk_kwargs = _get_chunk_args_kwargs(chunk)
362                    chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
363                chunk_list.append(chunk)
364                read_chunks += 1
365                if chunks is not None and read_chunks >= chunks:
366                    break
367        except Exception as e:
368            warn(f"[{self}] Failed to retrieve query results:\n" + str(e), stacklevel=3)
369            from meerschaum.utils.formatting import get_console
370            if not silent:
371                get_console().print_exception()
372
373    ### If no chunks returned, read without chunks
374    ### to get columns
375    if len(chunk_list) == 0:
376        with warnings.catch_warnings():
377            warnings.filterwarnings('ignore', 'case sensitivity issues')
378            _ = read_sql_query_kwargs.pop('chunksize', None)
379            with self.engine.begin() as connection:
380                chunk_list.append(
381                    pd.read_sql_query(
382                        formatted_query,
383                        connection,
384                        **read_sql_query_kwargs
385                    )
386                )
387
388    ### call the hook on any missed chunks.
389    if chunk_hook is not None and len(chunk_list) > len(chunk_hook_results):
390        for c in chunk_list[len(chunk_hook_results):]:
391            chunk_args, chunk_kwargs = _get_chunk_args_kwargs(c)
392            chunk_hook_results.append(chunk_hook(*chunk_args, **chunk_kwargs))
393
394    ### chunksize is not None so must iterate
395    if debug:
396        end = time.perf_counter()
397        dprint(f"Fetched {len(chunk_list)} chunks in {round(end - start, 2)} seconds.")
398
399    if as_hook_results:
400        return chunk_hook_results
401    
402    ### Skip `pd.concat()` if `as_chunks` is specified.
403    if as_chunks:
404        for c in chunk_list:
405            c.reset_index(drop=True, inplace=True)
406            for col in get_numeric_cols(c):
407                c[col] = c[col].apply(lambda x: x.canonical() if isinstance(x, Decimal) else x)
408        return chunk_list
409
410    df = pd.concat(chunk_list).reset_index(drop=True)
411    ### NOTE: The calls to `canonical()` are to drop leading and trailing zeroes.
412    for col in get_numeric_cols(df):
413        df[col] = df[col].apply(lambda x: x.canonical() if isinstance(x, Decimal) else x)
414
415    return df